iris-core 1.1.5

//! Memory pool: reusable buffer allocation for GPU transfers and I/O.
//!
//! Provides `BufferPool` — a simple slab-style pool of `Vec<u8>` buffers
//! that avoids repeated allocation churn. Used internally by iris-core and
//! exposed for upper layers (iris-gpu, iris-js, etc.).

use std::sync::Mutex;

/// A pool of reusable byte buffers.
///
/// Each buffer starts at `chunk_size` bytes and grows as needed.
/// When returned via `give_back`, the buffer is reset (length = 0)
/// and made available for the next `take()` call.
pub struct BufferPool {
    /// Minimum size for each allocated buffer.
    chunk_size: usize,
    /// Maximum number of idle buffers to retain.
    max_idle: usize,
    /// Idle buffers ready for reuse.
    idle: Mutex<Vec<Vec<u8>>>,
}

impl BufferPool {
    /// Create a new buffer pool.
    ///
    /// * `chunk_size` — minimum capacity per buffer (default 4096).
    /// * `max_idle`   — max number of buffers to keep in the pool (default 64).
    pub fn new(chunk_size: usize, max_idle: usize) -> Self {
        Self {
            chunk_size,
            max_idle,
            idle: Mutex::new(Vec::with_capacity(max_idle)),
        }
    }

    /// Take a buffer from the pool, or allocate a fresh one.
    ///
    /// The returned buffer has length 0 and capacity >= `chunk_size`.
    pub fn take(&self) -> Vec<u8> {
        let mut idle = self.idle.lock().unwrap();
        idle.pop().unwrap_or_else(|| Vec::with_capacity(self.chunk_size))
    }

    /// Take a buffer pre-sized to at least `min_capacity`.
    pub fn take_sized(&self, min_capacity: usize) -> Vec<u8> {
        let mut buf = self.take();
        if buf.capacity() < min_capacity {
            let additional = min_capacity - buf.len();
            buf.reserve(additional);
        }
        buf
    }

    /// Return a buffer to the pool for reuse.
    ///
    /// The buffer is cleared (length set to 0) but its allocation is kept.
    pub fn give_back(&self, mut buf: Vec<u8>) {
        buf.clear();
        let cap = buf.capacity();
        // Only keep buffers that aren't excessively large (> 4× chunk_size).
        if cap <= self.chunk_size * 4 {
            let mut idle = self.idle.lock().unwrap();
            if idle.len() < self.max_idle {
                idle.push(buf);
            }
        }
    }

    /// Pre-allocate `count` buffers into the pool.
    pub fn warm(&self, count: usize) {
        let mut idle = self.idle.lock().unwrap();
        let remaining = self.max_idle.saturating_sub(idle.len());
        let to_add = count.min(remaining);
        for _ in 0..to_add {
            idle.push(Vec::with_capacity(self.chunk_size));
        }
    }

    /// Number of buffers currently idle in the pool.
    pub fn idle_count(&self) -> usize {
        self.idle.lock().unwrap().len()
    }
}

impl Default for BufferPool {
    fn default() -> Self {
        Self::new(4096, 64)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_take_and_give_back() {
        let pool = BufferPool::default();
        let buf = pool.take();
        assert!(buf.capacity() >= 4096);
        assert!(buf.is_empty());
        pool.give_back(buf);
        assert_eq!(pool.idle_count(), 1);
    }

    #[test]
    fn test_take_sized() {
        let pool = BufferPool::default();
        let buf = pool.take_sized(8192);
        assert!(buf.capacity() >= 8192);
    }

    #[test]
    fn test_warm() {
        let pool = BufferPool::new(1024, 16);
        pool.warm(8);
        assert_eq!(pool.idle_count(), 8);
    }

    #[test]
    fn test_max_idle_respected() {
        let pool = BufferPool::new(64, 2);
        let b1 = pool.take();
        let b2 = pool.take();
        let b3 = pool.take();
        pool.give_back(b1);
        pool.give_back(b2);
        pool.give_back(b3);
        assert_eq!(pool.idle_count(), 2);
    }
}