simple_arena 0.2.0

//! # `simple-arena`
//!
//! A simple but fast arena allocator for a single type.
//! Using it allows you to allocate structs that are guaranteed to have the same lifetime.
//! This is useful for data structures that need to be able to cyclically reference each other, like trees.
//! The downside is that you can't deallocate individual entries, you can deallocate the whole arena at once.
//!
//! ## Examples
//!
//! ```
//! use simple_arena::Arena;
//!
//! let mut arena = Arena::new();
//! let node1 = arena.alloc(1);
//! let node2 = arena.alloc(2);
//! assert_eq!(*node1, 1);
//! assert_eq!(*node2, 2);
//! ```
//!
//! The arena can also be used to allocate structs that contain references to other structs in the arena using `Cell`.
//!
//! ```
//! use simple_arena::Arena;
//! use std::cell::Cell;
//!
//! struct Node<'a> {
//!     next: Cell<Option<&'a Node<'a>>>,
//!     value: u32,
//! }
//!
//! let mut arena = Arena::new();
//! let node1 = arena.alloc(Node { next: Cell::new(None), value: 1 });
//! let node2 = arena.alloc(Node { next: Cell::new(Some(node1)), value: 2 });
//! node1.next.set(Some(node2));
//!
//! assert_eq!(node1.next.get().unwrap().value, 2);
//! ```

use core::cell::RefCell;

pub struct Arena<T> {
	chunks: RefCell<Vec<Vec<T>>>,
	chunk_capacity: usize,
}

#[allow(dead_code)]
impl<T> Arena<T> {
	const INITIAL_CHUNK_CAPACITY: usize = 1024;

	/// Creates a new empty arena.
	#[must_use]
	pub const fn new() -> Self {
		Self {
			chunks: RefCell::new(Vec::new()),
			chunk_capacity: Self::INITIAL_CHUNK_CAPACITY,
		}
	}

	/// Creates a new empty arena with the specified initial capacity.
	#[must_use]
	pub const fn with_capacity(initial_capacity: usize) -> Self {
		Self {
			chunks: RefCell::new(Vec::new()),
			chunk_capacity: if initial_capacity == 0 { // use if instead of max to allow const fn
				1
			} else {
				initial_capacity
			},
		}
	}

	#[allow(clippy::mut_from_ref)]
	/// Allocates a new value in the arena and returns a mutable reference to it.
	///
	/// ## Example
	/// ```
	/// use simple_arena::Arena;
	///
	/// let arena = Arena::new();
	/// let a = arena.alloc(1);
	/// assert_eq!(*a, 1);
	/// ```
	pub fn alloc(&self, value: T) -> &mut T {
		let mut chunks = self.chunks.borrow_mut();

		if chunks.is_empty() {
			chunks.push(Vec::with_capacity(self.chunk_capacity));
		} else if let Some(last) = chunks.last() {
			if last.len() == last.capacity() {
				let new_capacity = last.len() * 2;
				chunks.push(Vec::with_capacity(new_capacity));
			}
		};
		let current = chunks.last_mut().expect("chunks should not be empty");
		let last = current.len();
		current.push(value);

		unsafe { &mut *current.as_mut_ptr().add(last) }
	}

	/// Returns the number of elements in the arena.
	/// O(n) time complexity.
	///
	/// ## Example
	/// ```
	/// use simple_arena::Arena;
	///
	/// let arena = Arena::new();
	/// let a = arena.alloc(1);
	/// let b = arena.alloc(2);
	/// let c = arena.alloc(3);
	/// assert_eq!(arena.len(), 3);
	/// ```
	pub fn len(&self) -> usize {
		let mut len = 0;
		for chunk in self.chunks.borrow().iter() {
			len += chunk.len();
		}
		len
	}

	/// Returns true if the arena is empty.
	/// This is faster than checking if len() == 0.
	/// O(1) time complexity.
	#[inline]
	pub fn is_empty(&self) -> bool {
		self.chunks.borrow().is_empty() || self.chunks.borrow()[0].is_empty()
	}

	/// Converts the arena into a vector consuming the arena.
	///
	/// ## Example
	/// ```
	/// use simple_arena::Arena;
	///
	/// let arena = Arena::new();
	/// let a = arena.alloc(1);
	/// let b = arena.alloc(2);
	/// let vec = arena.into_vec();
	///
	/// assert_eq!(vec, vec![1, 2]);
	/// ```
	pub fn into_vec(self) -> Vec<T> {
		self.chunks.into_inner().into_iter().flatten().collect()
	}

	/// Returns the number of chunks in the arena.
	fn chunks(&self) -> usize {
		self.chunks.borrow().len()
	}
}

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

	#[derive(Debug, PartialEq)]
	struct TestStruct {
		a: i32,
		b: i32,
	}

	#[test]
	pub fn test() {
		let arena = Arena::<TestStruct>::with_capacity(2);

		assert_eq!(arena.len(), 0);
		assert_eq!(arena.chunks(), 0);
		let a = arena.alloc(TestStruct { a: 1, b: 2 });
		assert_eq!(arena.len(), 1);
		assert_eq!(arena.chunks(), 1);
		let b = arena.alloc(TestStruct { a: 3, b: 4 });
		assert_eq!(arena.len(), 2);
		assert_eq!(arena.chunks(), 1);
		let c = arena.alloc(TestStruct { a: 5, b: 6 });
		assert_eq!(arena.len(), 3);
		assert_eq!(arena.chunks(), 2);
		let d = arena.alloc(TestStruct { a: 7, b: 8 });
		assert_eq!(arena.len(), 4);
		assert_eq!(arena.chunks(), 2);

		c.a = 10;
		c.b = 11;

		assert_eq!(a, &TestStruct { a: 1, b: 2 });
		assert_eq!(b, &TestStruct { a: 3, b: 4 });
		assert_eq!(c, &TestStruct { a: 10, b: 11 });
		assert_eq!(d, &TestStruct { a: 7, b: 8 });

		let vec = arena.into_vec();
		assert_eq!(
			vec,
			vec![
				TestStruct { a: 1, b: 2 },
				TestStruct { a: 3, b: 4 },
				TestStruct { a: 10, b: 11 },
				TestStruct { a: 7, b: 8 },
			]
		);
	}

	#[test]
	fn vec_order() {
		let arena = Arena::with_capacity(1);
		for &s in &["h", "e", "l", "l", "o"] {
			arena.alloc(String::from(s));
		}
		assert_eq!(arena.into_vec(), vec!["h", "e", "l", "l", "o"]);
	}

	#[test]
	fn test_zero_capacity() {
		let arena = Arena::with_capacity(0);
		assert_eq!(arena.chunk_capacity, 1);
		let a = arena.alloc(1);
		let b = arena.alloc(2);
		assert_eq!(arena.len(), 2);
		assert_eq!(arena.chunks(), 2);
		assert_eq!(*a, 1);
		assert_eq!(*b, 2);
	}
}