Skip to main content

sciforge_hub/tools/
arena.rs

1//! Bump-allocator arena for temporary buffers: slices, matrices, and scratch pools.
2//!
3//! Provides [`Arena`], [`ArenaSlice`], [`ArenaMatrix`], and [`ScratchPool`]
4//! to reuse memory without repeated dynamic allocation.
5
6use std::cell::RefCell;
7
8/// Contiguous `f64` bump allocator.
9///
10/// Pre-allocates a fixed-size vector and hands out slices
11/// (`ArenaSlice`) by advancing an internal cursor.
12pub struct Arena {
13    buf: RefCell<Vec<f64>>,
14    capacity: usize,
15    offset: RefCell<usize>,
16}
17
18impl Arena {
19    /// Creates an arena with room for `capacity` `f64` elements.
20    pub fn new(capacity: usize) -> Self {
21        Self {
22            buf: RefCell::new(vec![0.0; capacity]),
23            capacity,
24            offset: RefCell::new(0),
25        }
26    }
27
28    /// Allocates a slice of `n` elements, returns `None` if capacity is insufficient.
29    pub fn alloc_slice(&self, n: usize) -> Option<ArenaSlice> {
30        let mut off = self.offset.borrow_mut();
31        if *off + n > self.capacity {
32            return None;
33        }
34        let start = *off;
35        *off += n;
36        Some(ArenaSlice { start, len: n })
37    }
38
39    /// Writes `data` into the given slice (truncated if `data` exceeds the slice length).
40    pub fn write(&self, slice: &ArenaSlice, data: &[f64]) {
41        let mut buf = self.buf.borrow_mut();
42        let end = slice.start + data.len().min(slice.len);
43        buf[slice.start..end].copy_from_slice(&data[..end - slice.start]);
44    }
45
46    /// Reads the slice contents into a new `Vec<f64>`.
47    pub fn read(&self, slice: &ArenaSlice) -> Vec<f64> {
48        let buf = self.buf.borrow();
49        buf[slice.start..slice.start + slice.len].to_vec()
50    }
51
52    /// Returns the element at `index` within the slice, or `None` if out of bounds.
53    pub fn get(&self, slice: &ArenaSlice, index: usize) -> Option<f64> {
54        if index >= slice.len {
55            return None;
56        }
57        Some(self.buf.borrow()[slice.start + index])
58    }
59
60    /// Sets `value` at `index` within the slice.
61    pub fn set(&self, slice: &ArenaSlice, index: usize, value: f64) {
62        if index < slice.len {
63            self.buf.borrow_mut()[slice.start + index] = value;
64        }
65    }
66
67    /// Resets the cursor to zero, making all memory available for reuse.
68    pub fn reset(&self) {
69        *self.offset.borrow_mut() = 0;
70    }
71
72    /// Number of currently allocated elements.
73    pub fn used(&self) -> usize {
74        *self.offset.borrow()
75    }
76
77    /// Number of elements still available.
78    pub fn remaining(&self) -> usize {
79        self.capacity - self.used()
80    }
81
82    /// Total capacity of the arena.
83    pub fn capacity(&self) -> usize {
84        self.capacity
85    }
86}
87
88/// Handle to a contiguous slice within the arena.
89#[derive(Debug, Clone, Copy)]
90pub struct ArenaSlice {
91    start: usize,
92    len: usize,
93}
94
95impl ArenaSlice {
96    /// Length of the slice in number of elements.
97    pub fn len(&self) -> usize {
98        self.len
99    }
100    /// Returns `true` if the slice is empty.
101    pub fn is_empty(&self) -> bool {
102        self.len == 0
103    }
104}
105
106/// 2D matrix stored contiguously within an arena.
107pub struct ArenaMatrix {
108    slice: ArenaSlice,
109    /// Number of rows.
110    pub rows: usize,
111    /// Number of columns.
112    pub cols: usize,
113}
114
115impl ArenaMatrix {
116    /// Allocates a `rows × cols` matrix in `arena`.
117    pub fn new(arena: &Arena, rows: usize, cols: usize) -> Option<Self> {
118        let slice = arena.alloc_slice(rows * cols)?;
119        Some(Self { slice, rows, cols })
120    }
121
122    /// Returns the element at `(r, c)`, or `None` if out of bounds.
123    pub fn get(&self, arena: &Arena, r: usize, c: usize) -> Option<f64> {
124        if r >= self.rows || c >= self.cols {
125            return None;
126        }
127        arena.get(&self.slice, r * self.cols + c)
128    }
129
130    /// Sets `value` at `(r, c)`.
131    pub fn set(&self, arena: &Arena, r: usize, c: usize, value: f64) {
132        if r < self.rows && c < self.cols {
133            arena.set(&self.slice, r * self.cols + c, value);
134        }
135    }
136
137    /// Converts the matrix to `Vec<Vec<f64>>` (rows × columns).
138    pub fn to_vec(&self, arena: &Arena) -> Vec<Vec<f64>> {
139        let flat = arena.read(&self.slice);
140        flat.chunks(self.cols).map(|row| row.to_vec()).collect()
141    }
142}
143
144/// Pool of fixed-size scratch buffers stored in an arena.
145pub struct ScratchPool {
146    arena: Arena,
147    scratch_size: usize,
148    count: usize,
149}
150
151impl ScratchPool {
152    /// Creates a pool of `count` buffers with `scratch_size` elements each.
153    pub fn new(scratch_size: usize, count: usize) -> Self {
154        Self {
155            arena: Arena::new(scratch_size * count),
156            scratch_size,
157            count,
158        }
159    }
160
161    /// Returns the slice for buffer `index`, or `None` if out of bounds.
162    pub fn get(&self, index: usize) -> Option<ArenaSlice> {
163        if index >= self.count {
164            return None;
165        }
166        Some(ArenaSlice {
167            start: index * self.scratch_size,
168            len: self.scratch_size,
169        })
170    }
171
172    /// Writes `data` into buffer `index`.
173    pub fn write(&self, index: usize, data: &[f64]) {
174        if let Some(slice) = self.get(index) {
175            self.arena.write(&slice, data);
176        }
177    }
178
179    /// Reads the contents of buffer `index`.
180    pub fn read(&self, index: usize) -> Vec<f64> {
181        self.get(index)
182            .map(|s| self.arena.read(&s))
183            .unwrap_or_default()
184    }
185
186    /// Size of each scratch buffer.
187    pub fn scratch_size(&self) -> usize {
188        self.scratch_size
189    }
190    /// Number of buffers in the pool.
191    pub fn count(&self) -> usize {
192        self.count
193    }
194    /// Reference to the underlying arena.
195    pub fn arena(&self) -> &Arena {
196        &self.arena
197    }
198}