oxiphysics_gpu/compute/
buffer_pool.rs

1// Copyright 2026 COOLJAPAN OU (Team KitaSan)
2// SPDX-License-Identifier: Apache-2.0
3
4//! A power-of-two bucket pool for reusing CPU-side `Vec<f64>` buffers.
5//!
6//! Allocations are rounded up to the next power of two so that freed buffers
7//! can be reused for similarly-sized requests without reallocating.
8
9use std::collections::HashMap;
10
11/// Pool of reusable `Vec<f64>` allocations, bucketed by capacity (power of two).
12///
13/// This is used by the CPU-shadow (stub) path to avoid repeated heap allocation
14/// when the same buffer sizes are recycled across simulation steps.
15pub struct BufferPool {
16    /// Map from power-of-two bucket size → list of free buffers.
17    pub pools: HashMap<usize, Vec<Vec<f64>>>,
18    /// Maximum number of free buffers to retain per bucket before dropping.
19    pub cap_per_bucket: usize,
20}
21
22impl BufferPool {
23    /// Create a new pool with a default per-bucket cap of 64.
24    pub fn new() -> Self {
25        Self {
26            pools: HashMap::new(),
27            cap_per_bucket: 64,
28        }
29    }
30
31    /// Allocate a buffer of at least `len` elements.
32    ///
33    /// If a free buffer of the correct power-of-two bucket size exists it is
34    /// returned directly (contents may be stale); otherwise a fresh
35    /// zero-filled allocation is made.
36    pub fn alloc(&mut self, len: usize) -> Vec<f64> {
37        let bucket = len.next_power_of_two();
38        self.pools
39            .entry(bucket)
40            .or_default()
41            .pop()
42            .unwrap_or_else(|| vec![0.0; bucket])
43    }
44
45    /// Return a buffer to the pool for future reuse.
46    ///
47    /// If the pool for this bucket is already at capacity, the buffer is
48    /// silently dropped (deallocated).
49    pub fn free(&mut self, len: usize, buf: Vec<f64>) {
50        let bucket = len.next_power_of_two();
51        let pool = self.pools.entry(bucket).or_default();
52        if pool.len() < self.cap_per_bucket {
53            pool.push(buf);
54        }
55        // Otherwise drop `buf` — no-op, it is deallocated here.
56    }
57
58    /// Return the number of free buffers currently held by the pool.
59    pub fn total_free(&self) -> usize {
60        self.pools.values().map(|v| v.len()).sum()
61    }
62
63    /// Clear all pooled buffers, releasing memory back to the allocator.
64    pub fn clear(&mut self) {
65        self.pools.clear();
66    }
67}
68
69impl Default for BufferPool {
70    fn default() -> Self {
71        Self::new()
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[test]
80    fn test_pool_alloc_returns_correctly_sized_vec() {
81        let mut pool = BufferPool::new();
82        let buf = pool.alloc(10);
83        // bucket = 16 (next power of two ≥ 10)
84        assert!(buf.len() >= 10);
85        assert_eq!(buf.capacity(), 16);
86    }
87
88    #[test]
89    fn test_pool_alloc_free_reuses_buffer() {
90        let mut pool = BufferPool::new();
91        let buf = pool.alloc(8);
92        let ptr = buf.as_ptr();
93        pool.free(8, buf);
94        let buf2 = pool.alloc(8);
95        // The reused buffer should be the same allocation.
96        assert_eq!(buf2.as_ptr(), ptr);
97    }
98
99    #[test]
100    fn test_pool_bucket_cap_does_not_grow_unbounded() {
101        let mut pool = BufferPool::new();
102        pool.cap_per_bucket = 4;
103        // Free 8 buffers into the same bucket — only 4 should be retained.
104        for _ in 0..8 {
105            let buf = vec![0.0f64; 8];
106            pool.free(8, buf);
107        }
108        assert_eq!(pool.total_free(), 4);
109    }
110
111    #[test]
112    fn test_pool_clear_releases_all() {
113        let mut pool = BufferPool::new();
114        pool.free(4, vec![0.0; 4]);
115        pool.free(8, vec![0.0; 8]);
116        assert_eq!(pool.total_free(), 2);
117        pool.clear();
118        assert_eq!(pool.total_free(), 0);
119    }
120
121    #[test]
122    fn test_pool_alloc_zero_len_returns_power_of_two() {
123        let mut pool = BufferPool::new();
124        // 0.next_power_of_two() == 1 in Rust
125        let buf = pool.alloc(0);
126        assert!(!buf.is_empty());
127    }
128}
oxiphysics_gpu/compute/buffer_pool.rs

oxiphysics_gpu/compute/
buffer_pool.rs
