binius_compute/
alloc.rs

1// Copyright 2025 Irreducible Inc.
2
3use std::cell::Cell;
4
5use super::memory::{ComputeMemory, DevSlice};
6
7/// Basic bump allocator that allocates slices from an underlying memory buffer provided at
8/// construction.
9pub struct BumpAllocator<'a, F, Mem: ComputeMemory<F>> {
10	buffer: Cell<Option<Mem::FSliceMut<'a>>>,
11}
12
13impl<'a, F, Mem> BumpAllocator<'a, F, Mem>
14where
15	F: 'static,
16	Mem: ComputeMemory<F> + 'a,
17{
18	pub fn new(buffer: Mem::FSliceMut<'a>) -> Self {
19		Self {
20			buffer: Cell::new(Some(buffer)),
21		}
22	}
23
24	/// Allocates a slice of elements.
25	///
26	/// This method operates on an immutable self reference so that multiple allocator references
27	/// can co-exist. This follows how the `bumpalo` crate's `Bump` interface works. It may not be
28	/// necessary actually (since this partitions a borrowed slice, whereas `Bump` owns its memory).
29	///
30	/// ## Pre-conditions
31	///
32	/// - `n` must be a multiple of `Mem::MIN_SLICE_LEN`
33	pub fn alloc(&self, n: usize) -> Result<Mem::FSliceMut<'a>, Error> {
34		let buffer = self
35			.buffer
36			.take()
37			.expect("buffer is always Some by invariant");
38		// buffer temporarily contains None
39		if buffer.len() < n {
40			self.buffer.set(Some(buffer));
41			// buffer contains Some, invariant restored
42			Err(Error::OutOfMemory)
43		} else {
44			let (lhs, rhs) = Mem::split_at_mut(buffer, n.max(Mem::MIN_SLICE_LEN));
45			self.buffer.set(Some(rhs));
46			// buffer contains Some, invariant restored
47			Ok(lhs)
48		}
49	}
50}
51
52#[derive(Debug, thiserror::Error)]
53pub enum Error {
54	#[error("allocator is out of memory")]
55	OutOfMemory,
56}
57
58#[cfg(test)]
59mod tests {
60	use assert_matches::assert_matches;
61
62	use super::*;
63	use crate::cpu::memory::CpuMemory;
64
65	#[test]
66	fn test_alloc() {
67		let mut data = (0..256u128).collect::<Vec<_>>();
68
69		{
70			let bump = BumpAllocator::<u128, CpuMemory>::new(&mut data);
71			assert_eq!(bump.alloc(100).unwrap().len(), 100);
72			assert_eq!(bump.alloc(100).unwrap().len(), 100);
73			assert_matches!(bump.alloc(100), Err(Error::OutOfMemory));
74			// Release memory all at once.
75		}
76
77		// Reuse memory
78		let bump = BumpAllocator::<u128, CpuMemory>::new(&mut data);
79		let data = bump.alloc(100).unwrap();
80		assert_eq!(data.len(), 100);
81	}
82}
binius_compute/alloc.rs

binius_compute/
alloc.rs
