vecpool 0.1.0

//! Thread-local pool of reusable `Vec` buffers with cross-type reuse.
//!
//! `vecpool` eliminates repeated heap allocations by recycling `Vec` buffers.
//! When a [`PoolVec<T>`] is dropped, its underlying allocation is returned to a
//! thread-local pool. The next call to [`get`] or [`with_capacity`] can reuse that
//! buffer instead of asking the allocator for new memory.
//!
//! # Key Features
//!
//! - **Cross-type reuse**: Buffers are keyed by `(size_of::<T>(), align_of::<T>())`,
//!   so a `Vec<i32>` buffer can be reused as a `Vec<u32>` (same size and alignment).
//! - **Zero synchronization**: Each thread has its own pool. No mutexes, no atomics.
//! - **Transparent**: [`PoolVec<T>`] dereferences to `Vec<T>` — use all `Vec` methods directly.
//! - **Fast path**: Common power-of-two sizes (1–256 bytes) use a direct array
//!   lookup instead of hashing.
//!
//! # Quick Start
//!
//! ```
//! use vecpool::{get, with_capacity, PoolVec};
//!
//! // Get a vec from the pool (or a fresh one if the pool is empty)
//! let mut v = get::<u32>();
//! v.push(1);
//! v.push(2);
//! v.push(3);
//! assert_eq!(v.len(), 3);
//!
//! // On drop, the buffer is returned to the pool
//! drop(v);
//!
//! // The next request reuses the same buffer — no allocation!
//! let v2 = get::<u32>();
//! assert!(v2.capacity() >= 3);
//! ```
//!
//! # Cross-Type Reuse
//!
//! Types with the same memory layout share a pool lane:
//!
//! ```
//! use vecpool::get;
//!
//! let mut v = get::<i32>();
//! v.push(1);
//! v.push(2);
//! let cap = v.capacity();
//! drop(v);
//!
//! // u32 has the same size (4) and alignment (4) as i32 — reuses the buffer
//! let v2 = get::<u32>();
//! assert_eq!(v2.capacity(), cap);
//! ```
//!
//! # Pre-allocating Capacity
//!
//! ```
//! use vecpool::with_capacity;
//!
//! // If the pool has a buffer with >= 1000 capacity, reuse it.
//! // Otherwise, allocate fresh with Vec::with_capacity(1000).
//! let v = with_capacity::<u64>(1000);
//! assert!(v.capacity() >= 1000);
//! ```
//!
//! # Wrapping Existing Vecs
//!
//! ```
//! use vecpool::PoolVec;
//!
//! let data = vec![1u32, 2, 3, 4, 5];
//! let pooled = PoolVec::from(data);
//! // When `pooled` is dropped, its buffer enters the pool for future reuse.
//! ```
//!
//! # Escaping the Pool
//!
//! ```
//! use vecpool::get;
//!
//! let mut v = get::<u32>();
//! v.push(42);
//!
//! // Extract the inner Vec — buffer will NOT be returned to the pool
//! let owned: Vec<u32> = v.into_vec();
//! assert_eq!(owned, [42]);
//! ```
//!
//! Calling `.into_iter()` also escapes the pool (it calls `into_vec()` internally).
//!
//! # Collecting Iterators
//!
//! `PoolVec<T>` implements `FromIterator`, so `.collect()` works directly:
//!
//! ```
//! use vecpool::PoolVec;
//!
//! let v: PoolVec<i32> = (0..100).map(|x| x * 2).collect();
//! assert_eq!(v.len(), 100);
//! ```
//!
//! # Memory Management
//!
//! Pooled buffers persist for the lifetime of the thread. To reclaim memory:
//!
//! ```
//! use vecpool::{get, clear_pool};
//!
//! let mut v = get::<u32>();
//! v.push(1);
//! drop(v);
//!
//! // Free all pooled buffers on this thread
//! clear_pool();
//! ```
//!
//! # Thread Safety
//!
//! `PoolVec<T>` is `Send` when `T: Send`. If moved to another thread and dropped
//! there, the buffer is returned to *that* thread's pool. There is no cross-thread
//! synchronization.
//!
//! # Performance
//!
//! The pool uses a two-tier lookup:
//! - **Fast path**: Types with power-of-two size ≤ 256 and natural alignment
//!   (e.g., `u8`, `u32`, `u64`, `usize`) use a direct array index via
//!   `trailing_zeros()` — a single CPU instruction.
//! - **Slow path**: All other layouts fall back to an [`FxHashMap`].

use rustc_hash::FxHashMap;
use std::cell::UnsafeCell;
use std::fmt::Debug;
use std::mem::{self, ManuallyDrop};
use std::ops::{Deref, DerefMut};
use std::ptr::NonNull;

struct RawBuf {
    ptr: NonNull<u8>,
    capacity: usize, // in bytes
}

const FAST_LANES: usize = 9;

struct Pool {
    fast: [Vec<RawBuf>; FAST_LANES],
    slow: FxHashMap<(usize, usize), Vec<RawBuf>>,
}

impl Pool {
    fn new() -> Self {
        Self {
            fast: std::array::from_fn(|_| Vec::new()),
            slow: FxHashMap::default(),
        }
    }

    #[inline(always)]
    fn pop(&mut self, size: usize, align: usize) -> Option<RawBuf> {
        if let Some(idx) = fast_index(size, align) {
            self.fast[idx].pop()
        } else {
            self.slow.get_mut(&(size, align))?.pop()
        }
    }

    #[inline(always)]
    fn push(&mut self, size: usize, align: usize, buf: RawBuf) {
        if let Some(idx) = fast_index(size, align) {
            self.fast[idx].push(buf);
        } else {
            self.slow.entry((size, align)).or_default().push(buf);
        }
    }

    fn clear(&mut self) {
        for (idx, bufs) in self.fast.iter_mut().enumerate() {
            let align = 1 << idx;
            for buf in bufs.drain(..) {
                unsafe {
                    if buf.capacity > 0 {
                        let layout =
                            std::alloc::Layout::from_size_align_unchecked(buf.capacity, align);
                        std::alloc::dealloc(buf.ptr.as_ptr(), layout);
                    }
                }
            }
        }
        for (&(_, align), bufs) in self.slow.iter_mut() {
            for buf in bufs.drain(..) {
                unsafe {
                    if buf.capacity > 0 {
                        let layout =
                            std::alloc::Layout::from_size_align_unchecked(buf.capacity, align);
                        std::alloc::dealloc(buf.ptr.as_ptr(), layout);
                    }
                }
            }
        }
        self.slow.clear();
    }
}

#[inline(always)]
fn fast_index(size: usize, align: usize) -> Option<usize> {
    if size == align && size.is_power_of_two() && size <= 256 {
        Some(size.trailing_zeros() as usize)
    } else {
        None
    }
}

impl Drop for Pool {
    fn drop(&mut self) {
        self.clear();
    }
}

thread_local! {
    static POOL: UnsafeCell<Pool> = UnsafeCell::new(Pool::new());
}

/// A wrapper around [`Vec<T>`] that returns its buffer to the thread-local pool on drop.
///
/// `PoolVec<T>` dereferences to `Vec<T>`, so all `Vec` methods are available directly.
///
/// # Example
///
/// ```
/// use vecpool::PoolVec;
///
/// let mut v = PoolVec::<String>::new();
/// v.push("hello".to_string());
/// v.push("world".to_string());
/// assert_eq!(v.len(), 2);
/// // On drop, elements are dropped and the buffer is returned to the pool.
/// ```
pub struct PoolVec<T> {
    vec: ManuallyDrop<Vec<T>>,
}

/// Get an empty vec from the thread-local pool.
///
/// If a buffer with matching `(size_of::<T>(), align_of::<T>())` is available
/// in the pool, it is reused. Otherwise a fresh `Vec::new()` is returned.
///
/// # Example
///
/// ```
/// use vecpool::get;
///
/// let mut v = get::<u32>();
/// v.push(42);
/// assert_eq!(v[0], 42);
/// ```
#[must_use]
pub fn get<T>() -> PoolVec<T> {
    if mem::size_of::<T>() == 0 {
        return PoolVec {
            vec: ManuallyDrop::new(Vec::new()),
        };
    }

    let vec = POOL
        .try_with(|pool| {
            let pool = unsafe { &mut *pool.get() };
            pool.pop(mem::size_of::<T>(), mem::align_of::<T>())
                .map(|buf| unsafe {
                    Vec::from_raw_parts(
                        buf.ptr.as_ptr() as *mut T,
                        0,
                        buf.capacity / mem::size_of::<T>(),
                    )
                })
        })
        .ok()
        .flatten()
        .unwrap_or_default();

    PoolVec {
        vec: ManuallyDrop::new(vec),
    }
}

/// Get a vec from the pool with at least `capacity` elements of space.
///
/// If the pool has a buffer with sufficient capacity, it is reused.
/// Otherwise a fresh `Vec::with_capacity(capacity)` is allocated.
///
/// If the pool has a buffer but its capacity is less than requested,
/// the buffer remains in the pool and a fresh allocation is made.
///
/// # Example
///
/// ```
/// use vecpool::with_capacity;
///
/// let v = with_capacity::<u64>(1000);
/// assert!(v.capacity() >= 1000);
/// ```
#[must_use]
pub fn with_capacity<T>(capacity: usize) -> PoolVec<T> {
    if mem::size_of::<T>() == 0 {
        return PoolVec {
            vec: ManuallyDrop::new(Vec::new()),
        };
    }

    let vec = POOL
        .try_with(|pool| {
            let pool = unsafe { &mut *pool.get() };
            pool.pop(mem::size_of::<T>(), mem::align_of::<T>())
                .and_then(|buf| {
                    let cap = buf.capacity / mem::size_of::<T>();
                    if cap >= capacity {
                        Some(unsafe { Vec::from_raw_parts(buf.ptr.as_ptr() as *mut T, 0, cap) })
                    } else {
                        pool.push(mem::size_of::<T>(), mem::align_of::<T>(), buf);
                        None
                    }
                })
        })
        .ok()
        .flatten()
        .unwrap_or_else(|| Vec::with_capacity(capacity));

    PoolVec {
        vec: ManuallyDrop::new(vec),
    }
}

/// Drain all buffers from the thread-local pool, freeing their memory.
///
/// Call this when you know a workload is complete and want to reclaim memory.
///
/// # Example
///
/// ```
/// use vecpool::{get, clear_pool};
///
/// let mut v = get::<u32>();
/// v.push(1);
/// drop(v); // returned to pool
///
/// clear_pool(); // pool is now empty
///
/// let v2 = get::<u32>();
/// assert_eq!(v2.capacity(), 0); // pool is empty, returns a fresh Vec::new()
/// ```
pub fn clear_pool() {
    let _ = POOL.try_with(|pool| {
        let pool = unsafe { &mut *pool.get() };
        pool.clear();
    });
}

impl<T> Default for PoolVec<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> PoolVec<T> {
    /// Create a new empty `PoolVec`, reusing a buffer from the pool if available.
    ///
    /// This is equivalent to [`get::<T>()`](get).
    ///
    /// # Example
    ///
    /// ```
    /// use vecpool::PoolVec;
    ///
    /// let mut v = PoolVec::<i32>::new();
    /// v.push(10);
    /// ```
    pub fn new() -> Self {
        get()
    }

    /// Consume the `PoolVec` and return the inner `Vec<T>`.
    ///
    /// The buffer is **not** returned to the pool. The caller takes full
    /// ownership of the allocation.
    ///
    /// # Example
    ///
    /// ```
    /// use vecpool::get;
    ///
    /// let mut v = get::<u32>();
    /// v.push(1);
    /// v.push(2);
    ///
    /// let vec: Vec<u32> = v.into_vec();
    /// assert_eq!(vec, [1, 2]);
    /// ```
    #[must_use]
    pub fn into_vec(mut self) -> Vec<T> {
        let vec = unsafe { ManuallyDrop::take(&mut self.vec) };
        mem::forget(self);
        vec
    }
}

/// Wrap an existing `Vec<T>` so its buffer is returned to the pool on drop.
///
/// # Example
///
/// ```
/// use vecpool::PoolVec;
///
/// let vec = vec![1u32, 2, 3];
/// let pooled = PoolVec::from(vec);
/// assert_eq!(pooled[0], 1);
/// // On drop, the buffer goes to the pool for reuse.
/// ```
impl<T> From<Vec<T>> for PoolVec<T> {
    fn from(vec: Vec<T>) -> Self {
        PoolVec {
            vec: ManuallyDrop::new(vec),
        }
    }
}

impl<T> Deref for PoolVec<T> {
    type Target = Vec<T>;
    fn deref(&self) -> &Vec<T> {
        &self.vec
    }
}

impl<T> DerefMut for PoolVec<T> {
    fn deref_mut(&mut self) -> &mut Vec<T> {
        &mut self.vec
    }
}

impl<T> Drop for PoolVec<T> {
    fn drop(&mut self) {
        if mem::size_of::<T>() == 0 {
            unsafe {
                ManuallyDrop::drop(&mut self.vec);
            }
            return;
        }

        self.vec.clear();

        let ptr = self.vec.as_mut_ptr();
        let capacity = self.vec.capacity();

        if capacity == 0 {
            return;
        }

        let buf = RawBuf {
            ptr: unsafe { NonNull::new_unchecked(ptr as *mut u8) },
            capacity: capacity * mem::size_of::<T>(),
        };

        let returned = POOL
            .try_with(|pool| {
                let pool = unsafe { &mut *pool.get() };
                pool.push(mem::size_of::<T>(), mem::align_of::<T>(), buf);
            })
            .is_ok();

        if !returned {
            unsafe {
                ManuallyDrop::drop(&mut self.vec);
            }
        }
    }
}

/// Consumes the `PoolVec` and returns an iterator over its elements.
///
/// The underlying buffer is **not** returned to the pool (same as [`into_vec`](PoolVec::into_vec)).
impl<T> IntoIterator for PoolVec<T> {
    type Item = T;
    type IntoIter = std::vec::IntoIter<T>;
    fn into_iter(self) -> Self::IntoIter {
        self.into_vec().into_iter()
    }
}

/// Collects an iterator into a `PoolVec`, reusing a buffer from the pool.
///
/// # Example
///
/// ```
/// use vecpool::PoolVec;
///
/// let v: PoolVec<u32> = (0..10).collect();
/// assert_eq!(v.len(), 10);
/// ```
impl<T> FromIterator<T> for PoolVec<T> {
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        let mut v = get::<T>();
        v.extend(iter);
        v
    }
}

impl<T: Debug> Debug for PoolVec<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        (**self).fmt(f)
    }
}

impl<T> AsRef<[T]> for PoolVec<T> {
    fn as_ref(&self) -> &[T] {
        &self.vec
    }
}

impl<T> AsMut<[T]> for PoolVec<T> {
    fn as_mut(&mut self) -> &mut [T] {
        &mut self.vec
    }
}

impl<T> AsRef<Vec<T>> for PoolVec<T> {
    fn as_ref(&self) -> &Vec<T> {
        &self.vec
    }
}

impl<T> AsMut<Vec<T>> for PoolVec<T> {
    fn as_mut(&mut self) -> &mut Vec<T> {
        &mut self.vec
    }
}