1use crossbeam::queue::SegQueue;
  2use std::ops::{Deref, DerefMut};
  3use std::path::PathBuf;
  4use std::sync::Arc;
  5
  6/// Trait for types that can be pooled and reused.
  7///
  8/// Types implementing this trait can be:
  9/// 1. Created empty via `Default`
 10/// 2. Reset to empty state via `clear()` (preserving capacity)
 11/// 3. Sent across threads (required by `SegQueue`)
 12pub trait Poolable: Default + Send {
 13    /// Reset the value to its empty state, preserving any allocated capacity.
 14    fn clear(&mut self);
 15}
 16
 17impl<T: Send> Poolable for Vec<T> {
 18    fn clear(&mut self) {
 19        Vec::clear(self);
 20    }
 21}
 22
 23impl Poolable for String {
 24    fn clear(&mut self) {
 25        String::clear(self);
 26    }
 27}
 28
 29impl Poolable for PathBuf {
 30    fn clear(&mut self) {
 31        PathBuf::clear(self);
 32    }
 33}
 34
 35/// A thread-safe, lock-free pool of reusable buffers.
 36#[derive(Clone)]
 37pub struct BufPool<T: Poolable> {
 38    queue: Arc<SegQueue<T>>,
 39}
 40
 41impl<T: Poolable> BufPool<T> {
 42    /// Creates a pool pre-populated with `capacity` empty buffers.
 43    pub fn with_capacity(capacity: usize) -> Self {
 44        let queue = Arc::new(SegQueue::new());
 45        for _ in 0..capacity {
 46            queue.push(T::default());
 47        }
 48        Self { queue }
 49    }
 50
 51    /// Checks out a buffer from the pool. If the pool is empty, allocates a new buffer.
 52    pub fn checkout(&self) -> BufGuard<T> {
 53        let buf = self.queue.pop().unwrap_or_else(T::default);
 54        BufGuard {
 55            buf,
 56            pool: self.queue.clone(),
 57        }
 58    }
 59}
 60
 61/// RAII guard that automatically returns the buffer to the pool when dropped.
 62#[derive(Debug)]
 63pub struct BufGuard<T: Poolable> {
 64    buf: T,
 65    pool: Arc<SegQueue<T>>,
 66}
 67
 68impl<T: Poolable> BufGuard<T> {
 69    /// Extracts the buffer, preventing automatic return to the pool.
 70    /// Useful if you need to move the buffer elsewhere.
 71    pub fn into_inner(mut self) -> T {
 72        let buf = std::mem::take(&mut self.buf);
 73        // SAFETY:
 74        // 1) We forget `self`, preventing it from
 75        // double-free'ing `pool`
 76        // 2) The 'pointer' is really a &mut, so
 77        // it is aligned, type-valid, and valid for reads.
 78        let pool = unsafe { std::ptr::read(&self.pool) };
 79        std::mem::forget(self);
 80        drop(pool);
 81        buf
 82    }
 83}
 84
 85impl<T: Poolable> Drop for BufGuard<T> {
 86    fn drop(&mut self) {
 87        let mut buf = std::mem::take(&mut self.buf);
 88        buf.clear();
 89        self.pool.push(buf);
 90    }
 91}
 92
 93impl<T: Poolable> Deref for BufGuard<T> {
 94    type Target = T;
 95
 96    fn deref(&self) -> &Self::Target {
 97        &self.buf
 98    }
 99}
100
101impl<T: Poolable> DerefMut for BufGuard<T> {
102    fn deref_mut(&mut self) -> &mut Self::Target {
103        &mut self.buf
104    }
105}