poolshark/

lib.rs

//! A high-performance object pool that reuses allocations instead of freeing them.
//!
//! # Quick Start
//!
//! ```
//! use poolshark::local::LPooled;
//! use std::collections::HashMap;
//!
//! // Take a HashMap from the thread-local pool (or create new if empty)
//! let mut map: LPooled<HashMap<String, i32>> = LPooled::take();
//! map.insert("answer".to_string(), 42);
//! // When dropped, the HashMap is cleared and returned to the pool
//! ```
//!
//! # Which Pool Should I Use?
//!
//! - **Use [`local::LPooled`]** (default choice): Faster, for objects created and dropped on the same thread(s)
//! - **Use [`global::GPooled`]**: When one thread creates objects and other threads drop them (producer-consumer)
//!
//! # Why Poolshark?
//!
//! - **Reduce allocations**: Reuse containers instead of repeatedly allocating and freeing
//! - **Predictable performance**: Consistent behavior across platforms, independent of allocator
//! - **Fast**: Local pools avoid atomic operations and are more ergonomic than `thread_local!`
//! - **Flexible**: Choose between fast thread-local pools or lock-free cross-thread pools
//!
//! # Pool Types
//!
//! ## Global Pooling
//!
//! Global pools share objects between threads (see [`global::GPooled`]).
//! An object taken from a global pool always returns to the pool it was
//! taken from, regardless of which thread drops it. Use this for producer-consumer
//! patterns where one thread creates objects and other threads consume them.
//!
//! There are several different ways to use global pools. You can use
//! [take](global::take) or [take_any](global::take_any) to just take objects
//! from thread local global pools. If you need better performance you can use
//! [pool](global::pool) or [pool_any](global::pool_any) and then store the pool
//! somewhere. If you don't have anywhere to store the pool you can use a static
//! [LazyLock](std::sync::LazyLock) for a truly global named pool. For example,
//!
//! ```no_run
//! use std::{sync::LazyLock, collections::HashMap};
//! use poolshark::global::{Pool, GPooled};
//!
//! type Widget = HashMap<usize, usize>;
//!
//! // create a global static widget pool that will accept up to 1024 widgets with
//! // up to 64 elements of capacity each
//! static WIDGETS: LazyLock<Pool<Widget>> = LazyLock::new(|| Pool::new(1024, 64));
//!
//! fn widget_maker() -> GPooled<Widget> {
//!     let mut w = WIDGETS.take();
//!     w.insert(42, 42);
//!     w
//! }
//!
//! fn widget_user(w: GPooled<Widget>) {
//!     drop(w) // puts the widget back in the WIDGETS pool
//! }
//! ```
//!
//! ## Local Pooling
//!
//! Local pools (see [`local::LPooled`]) always return dropped objects to a thread-local
//! pool on whichever thread drops them. They are significantly faster than global pools
//! because they avoid atomic operations. Use local pools by default unless you have
//! a cross-thread producer-consumer pattern.
//!
//! **Thread safety**: `LPooled<T>` is `Send + Sync` whenever `T` is `Send + Sync`, so you can
//! safely pass pooled objects between threads.
//!
//! Local pools require types to implement the unsafe trait [`IsoPoolable`], but all
//! standard containers (Vec, HashMap, String, etc.) already implement it.
//!
//! ```no_run
//! use poolshark::local::LPooled;
//! use std::collections::HashMap;
//!
//! type Widget = HashMap<usize, usize>;
//!
//! fn widget_maker() -> LPooled<Widget> {
//!     let mut w = LPooled::<Widget>::default(); // takes from the local pool
//!     w.insert(42, 42);
//!     w
//! }
//!
//! fn widget_user(w: LPooled<Widget>) {
//!     drop(w) // puts the widget back in the local pool
//! }
//! ```
use global::WeakPool;
pub use poolshark_derive::location_id;
use std::alloc::Layout;

pub mod global;
pub mod local;
pub mod pooled;

/// A globally unique id for a source code position
///
/// use poolshark_derive::location_id!() macro to generate one
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct LocationId(pub u16);

#[cfg(test)]
mod test;

// msb 0 -> it's a layout
// msb 1 -> it's a size
//
// layout: 1 bit flag, 12 bit size, 3 bit align
//
// aligns
// 0x0 -> 1
// 0x1 -> 2
// 0x2 -> 4
// 0x3 -> 8
// 0x4 -> 16
//
// size: 1 bit flag, 15 bit size
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct ULayout(u16);

impl Default for ULayout {
    fn default() -> Self {
        Self(0)
    }
}

impl ULayout {
    const fn empty() -> Self {
        Self(0)
    }

    const fn is_empty(&self) -> bool {
        self.0 == 0
    }

    const fn new_layout<T>() -> Option<Self> {
        let l = Layout::new::<T>();
        let size = l.size();
        let align = l.align();
        if size > 0x0FFF {
            return None;
        }
        let align = match align {
            1 => 0x0,
            2 => 0x1,
            4 => 0x2,
            8 => 0x3,
            16 => 0x4,
            _ => return None,
        };
        Some(Self(((size << 3) | align) as u16))
    }

    const fn new_size(sz: usize) -> Option<Self> {
        if sz > 0x7FFF {
            None
        } else {
            Some(Self((0x8000 | sz) as u16))
        }
    }
}

macro_rules! add_param {
    ($d:expr, $p:ty) => {
        match $d.add_param::<$p>() {
            Some(d) => d,
            None => return None,
        }
    };
}

/// Type describing the layout, alignment, and type of a container
///
/// `Discriminant` is central to the safety and performance of local pooling. It
/// describes 2 things in just 8 bytes.
///
/// - The unique location in the source code of the implementation of
/// [IsoPoolable]. This is accomplished by a proc macro that generates a global
/// table of unique location ids for cross crate source code locations. This
/// unique id ensures that different container types can't be mixed in the same
/// pool.
///
/// - The layout and alignment of all the type parameters of the
/// container. Discriminant has 3 slots that can be filled with either
/// type parameters or const SIZE parameters. If your container has
/// more parameters than that then you can't locally pool it, and you
/// can't implement [IsoPoolable]. If you try you will likely cause
/// undefined behavior.
///
/// In order to squeeze all this information into just 8 bytes there are some
/// limitations.
///
/// - You can't have more than 0xFFFF implementations of [IsoPoolable] in the
/// same project. This includes all the crates depended on by the project.
///
/// - Your type parameters must have size <= 0x0FFF bytes and
///   alignment of 1, 2, 4, 8, or 16. Alignments > 16 will be rejected.
///
/// - const SIZE parameters must be <= 0x7FFF.
///
/// If any of these constraints are violated the `Discriminant` constructors
/// will return `None`. If you desire you may panic at that point to cause a
/// compile error. If you do not panic and instead leave `DISCRIMINANT` as
/// `None` then local pool operations on that type will work just fine, but
/// nothing will be pooled. Objects will be freed when they are dropped and
/// [take](local::take) will allocate new objects each time it is called.
///
/// # Discriminant Collisions and Why They're Safe
///
/// Two different types can have the same discriminant if they have the same size and
/// alignment. For example:
///
/// ```ignore
/// #[repr(C)]
/// struct Padded1 { a: u8, _pad: [u8; 7], b: u64 }  // size 16, align 8
///
/// #[repr(C)]
/// struct Padded2 { x: u64, y: u64 }                 // size 16, align 8
/// ```
///
/// If you pool `Vec<Padded1>` and `Vec<Padded2>`, they would get the same discriminant
/// because `Padded1` and `Padded2` have identical size and alignment. This means a
/// `Vec<Padded1>` allocation could be reused as a `Vec<Padded2>` allocation.
///
/// **This is safe** because:
///
/// 1. Containers are **always empty** when returned to pools (`reset()` ensures this)
/// 2. An empty `Vec<T>` only cares about `T`'s size and alignment for its allocation
/// 3. The actual bit patterns inside `T` don't matter when the Vec is empty
/// 4. When you take from the pool and populate it with your type, it's initialized correctly
///
/// The discriminant system is designed to ensure that different container **types** never
/// share pools (via the `LocationId`), and that the **memory layout** of type parameters
/// is compatible. As long as containers are properly emptied before pooling (which `reset()`
/// guarantees), the system is memory safe even with discriminant collisions.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Discriminant {
    container: LocationId,
    elements: [ULayout; 3],
}

impl Discriminant {
    /// return a new empty discriminant
    pub const fn empty(id: LocationId) -> Discriminant {
        Discriminant { container: id, elements: [ULayout::empty(); 3] }
    }

    /// build a discriminant for a type with no type variables (just a location
    /// id). Always returns Some
    pub const fn new(id: LocationId) -> Option<Discriminant> {
        Some(Self::empty(id))
    }

    /// Add a type parameter.
    ///
    /// Discriminant has 3 slots. Each slot can hold either a type
    /// parameter or a const SIZE. This will return None if the
    /// discriminant is full, or the type parameter's size or
    /// alignment are too big.
    pub const fn add_param<T>(mut self) -> Option<Self> {
        let l = match ULayout::new_layout::<T>() {
            None => return None,
            Some(l) => l,
        };
        let mut i = 0;
        while i < 3 {
            if self.elements[i].is_empty() {
                self.elements[i] = l;
                return Some(self);
            }
            i += 1
        }
        None
    }

    /// Add a const SIZE
    ///
    /// Discriminant has 3 slots. Each slot can hold either a type
    /// parameter or a const SIZE. This will return None if the
    /// discriminant is full, or if the size is too large.
    pub const fn add_size<const SIZE: usize>(mut self) -> Option<Self> {
        let l = match ULayout::new_size(SIZE) {
            None => return None,
            Some(l) => l,
        };
        let mut i = 0;
        while i < 3 {
            if self.elements[i].is_empty() {
                self.elements[i] = l;
                return Some(self);
            }
            i += 1
        }
        None
    }

    /// build a discriminant with one type param
    pub const fn new_p1<T>(id: LocationId) -> Option<Discriminant> {
        let d = Discriminant::empty(id);
        d.add_param::<T>()
    }

    /// build a discriminant with one type param and a size
    pub const fn new_p1_size<T, const SIZE: usize>(
        id: LocationId,
    ) -> Option<Discriminant> {
        let d = Discriminant::empty(id);
        let d = add_param!(d, T);
        d.add_size::<SIZE>()
    }

    /// build a discriminant with two type params
    pub const fn new_p2<T, U>(id: LocationId) -> Option<Discriminant> {
        let d = Discriminant::empty(id);
        let d = add_param!(d, T);
        d.add_param::<U>()
    }

    /// build a discriminant with two type params and a size
    pub const fn new_p2_size<T, U, const SIZE: usize>(
        id: LocationId,
    ) -> Option<Discriminant> {
        let d = Discriminant::empty(id);
        let d = add_param!(d, T);
        let d = add_param!(d, U);
        d.add_size::<SIZE>()
    }

    /// build a discriminant with three type params
    pub const fn new_p3<T, U, V>(id: LocationId) -> Option<Discriminant> {
        let d = Discriminant::empty(id);
        let d = add_param!(d, T);
        let d = add_param!(d, U);
        d.add_param::<V>()
    }
}

struct Opaque {
    t: *mut (),
    drop: Option<Box<dyn FnOnce(*mut ())>>,
}

impl Drop for Opaque {
    fn drop(&mut self) {
        if let Some(f) = self.drop.take() {
            f(self.t)
        }
    }
}

/// Trait for poolable objects
pub trait Poolable {
    /// allocate a new empty collection
    fn empty() -> Self;

    /// empty the collection and reset it to its default state so it
    /// can be put back in the pool.
    fn reset(&mut self);

    /// return the capacity of the collection
    fn capacity(&self) -> usize;

    /// return true if the object has really been dropped, e.g. if
    /// you're pooling an Arc then Arc::get_mut().is_some() == true.
    fn really_dropped(&mut self) -> bool {
        true
    }
}

/// Low level global pool trait for maximum control
///
/// Implementing this trait allows full low level control over where the pool
/// pointer is stored. For example if you are pooling an allocated data
/// structure, you could store the pool pointer in the allocation to keep the
/// size of the handle struct to a minimum. E.G. you're pooling a
/// [triomphe::ThinArc]. Or, if you have a static global pool, then you would
/// not need to keep a pool pointer at all.
///
/// The object's drop implementation should return the object to the
/// pool instead of deallocating it
///
/// Implementing this trait correctly is extremely tricky, and requires unsafe
/// code, therefore it is marked as unsafe.
///
/// Most of the time you should use the [GPooled](global::GPooled) wrapper.
pub unsafe trait RawPoolable: Sized {
    /// allocate a new empty object and set it's pool pointer to `pool`
    fn empty(pool: WeakPool<Self>) -> Self;

    /// empty the collection and reset it to its default state so it
    /// can be put back in the pool
    fn reset(&mut self);

    /// return the capacity of the collection
    fn capacity(&self) -> usize;

    /// Actually drop the inner object, don't put it back in the pool,
    /// make sure you do not call both this method and the drop
    /// implementation that puts the object back in the pool!
    fn really_drop(self);
}

/// Trait for isomorphicly poolable objects.
///
/// That is objects that can safely be pooled by memory layout and container
/// type. For example two `HashMap`s, `HashMap<usize, usize>` and
/// `HashMap<ArcStr, ArcStr>` are isomorphic, their memory allocations can be
/// used interchangably so long as they are empty.
pub unsafe trait IsoPoolable: Poolable {
    /// # Getting the Layout Right
    ///
    /// You must pass every type variable that can effect the layout
    /// of the container's inner allocation to Discriminant. Take
    /// HashMap as an example. If you build the discriminant such as
    /// `Discriminant::new_p1::<HashMap<K, V>>()` it would always be
    /// the same for any `K`, `V`, because the `HashMap` struct
    /// doesn't actually contain any `K`s or `V`s, just a pointer to
    /// some `K`s and `V`s. If you implemented discriminant this way
    /// it would cause undefined behavior when you tried to pool two
    /// HashMap's with `K`, `V` types that aren't isomorphic. Instead
    /// you must pass `K` and `V` to `Discriminant::new_p2::<K, V>()`
    /// to get the real layout of the inner collection of
    /// `HashMap`. This is why this trait is unsafe to implement, if
    /// you aren't careful when you build the discriminant very bad
    /// things will happen.
    ///
    /// # Why not TypeId
    ///
    /// The reason why Discriminant is used instead of
    /// [`TypeId`](std::any::TypeId) (which would accomplish the same
    /// goal) is twofold. First Discriminant is 1 word on a 64 bit
    /// machine, and thus very fast to index, and second `TypeId` only
    /// supports types without references. However we often want to
    /// pool empty containers where the inner type is a reference,
    /// thus we cannot use `TypeId`.
    ///
    /// # Why Discriminant is an Option
    ///
    /// Discriminant is a compressed version of layout that squeezes 2
    /// layouts a size and a container type into 8 bytes. As such
    /// there are some layouts that are too big to fit in it, and the
    /// constructor will return None in those cases. For the purpose
    /// of pooling containers of small objects these tradeoffs seemed
    /// worth it. If you must pool containers of huge objects like
    /// this, you can use the global pools.
    ///
    /// # Arc
    ///
    /// It is not safe to implement this trait for
    /// [`Arc`](std::sync::Arc) or in general for any container that
    /// can't be totally empty. This is because having the same
    /// Discriminant only guarantees that two types are isomorphic, it
    /// does not guarantee that they have the same bit patterns.
    /// Normal container types are safe in spite of this because reset
    /// makes sure they are empty, and thus no errent bit patterns
    /// exist in the container and all we care about is that the
    /// container's allocation is isomorphic with respect to the types
    /// we want to put in it. However `Arc` can never be empty, and
    /// since notch optimization may change the bit pattern of `None`
    /// depending on the type of `T`, it is not even safe to pool
    /// `Arc<Option<T>>`. Because if `T` and `U` were isomorphic, but
    /// notch optimization used a different bit pattern for `None`,
    /// then pooling these objects could cause undefined behavior.
    const DISCRIMINANT: Option<Discriminant>;
}