#![allow(clippy::cast_possible_truncation)]
use std::fmt;
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct NodeId {
index: u32,
generation: u32,
}
impl NodeId {
#[inline]
pub(crate) const fn index(self) -> usize {
self.index as usize
}
#[inline]
#[allow(dead_code)]
pub(crate) const fn generation(self) -> u32 {
self.generation
}
#[inline]
#[must_use]
pub const fn from_raw_parts(index: u32, generation: u32) -> Self {
Self { index, generation }
}
#[inline]
#[must_use]
pub const fn raw_parts(self) -> (u32, u32) {
(self.index, self.generation)
}
}
impl fmt::Debug for NodeId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "NodeId({}g{})", self.index, self.generation)
}
}
#[derive(Debug)]
#[allow(clippy::cast_possible_truncation)] enum Slot<T> {
Occupied(T),
Vacant(u32),
}
#[derive(Debug)]
#[allow(clippy::cast_possible_truncation)] pub struct Arena<T> {
slots: Vec<Slot<T>>,
generations: Vec<u32>,
free_head: u32,
len: usize,
}
impl<T> Arena<T> {
#[must_use]
pub fn new() -> Self {
Self {
slots: Vec::new(),
generations: Vec::new(),
free_head: u32::MAX,
len: 0,
}
}
#[must_use]
pub fn with_capacity(capacity: usize) -> Self {
Self {
slots: Vec::with_capacity(capacity),
generations: Vec::with_capacity(capacity),
free_head: u32::MAX,
len: 0,
}
}
pub fn alloc(&mut self, value: T) -> NodeId {
self.len += 1;
if self.free_head == u32::MAX {
let index = self.slots.len();
self.slots.push(Slot::Occupied(value));
self.generations.push(0);
NodeId {
index: index as u32,
generation: 0,
}
} else {
let index = self.free_head as usize;
let next_free = match self.slots[index] {
Slot::Vacant(next) => next,
Slot::Occupied(_) => unreachable!("free list pointed to occupied slot"),
};
self.slots[index] = Slot::Occupied(value);
self.free_head = next_free;
NodeId {
index: index as u32,
generation: self.generations[index],
}
}
}
pub fn dealloc(&mut self, id: NodeId) -> Option<T> {
let index = id.index();
if index >= self.slots.len() || self.generations[index] != id.generation {
return None;
}
match &self.slots[index] {
Slot::Vacant(_) => return None,
Slot::Occupied(_) => {}
}
let old_slot = std::mem::replace(&mut self.slots[index], Slot::Vacant(self.free_head));
self.free_head = index as u32;
self.generations[index] = self.generations[index].wrapping_add(1);
self.len -= 1;
match old_slot {
Slot::Occupied(value) => Some(value),
Slot::Vacant(_) => unreachable!(),
}
}
#[must_use]
pub fn get(&self, id: NodeId) -> Option<&T> {
let index = id.index();
if index >= self.slots.len() || self.generations[index] != id.generation {
return None;
}
match &self.slots[index] {
Slot::Occupied(value) => Some(value),
Slot::Vacant(_) => None,
}
}
#[must_use]
pub fn get_mut(&mut self, id: NodeId) -> Option<&mut T> {
let index = id.index();
if index >= self.slots.len() || self.generations[index] != id.generation {
return None;
}
match &mut self.slots[index] {
Slot::Occupied(value) => Some(value),
Slot::Vacant(_) => None,
}
}
#[inline]
#[must_use]
pub fn len(&self) -> usize {
self.len
}
#[inline]
#[must_use]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
#[must_use]
pub fn capacity(&self) -> usize {
self.slots.capacity()
}
pub fn clear(&mut self) {
self.slots.clear();
self.generations.clear();
self.free_head = u32::MAX;
self.len = 0;
}
pub fn iter(&self) -> ArenaIter<'_, T> {
ArenaIter {
arena: self,
index: 0,
}
}
pub fn iter_mut(&mut self) -> ArenaIterMut<'_, T> {
ArenaIterMut {
inner: self.slots.iter_mut().enumerate(),
generations: &self.generations,
}
}
#[must_use]
pub fn contains(&self, id: NodeId) -> bool {
self.get(id).is_some()
}
}
impl<T> Default for Arena<T> {
fn default() -> Self {
Self::new()
}
}
impl<'a, T> IntoIterator for &'a Arena<T> {
type Item = (NodeId, &'a T);
type IntoIter = ArenaIter<'a, T>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'a, T> IntoIterator for &'a mut Arena<T> {
type Item = (NodeId, &'a mut T);
type IntoIter = ArenaIterMut<'a, T>;
fn into_iter(self) -> Self::IntoIter {
self.iter_mut()
}
}
pub struct ArenaIter<'a, T> {
arena: &'a Arena<T>,
index: usize,
}
impl<'a, T> Iterator for ArenaIter<'a, T> {
type Item = (NodeId, &'a T);
fn next(&mut self) -> Option<Self::Item> {
while self.index < self.arena.slots.len() {
let i = self.index;
self.index += 1;
if let Slot::Occupied(ref value) = self.arena.slots[i] {
let id = NodeId {
index: i as u32,
generation: self.arena.generations[i],
};
return Some((id, value));
}
}
None
}
fn size_hint(&self) -> (usize, Option<usize>) {
(0, Some(self.arena.slots.len() - self.index))
}
}
pub struct ArenaIterMut<'a, T> {
inner: std::iter::Enumerate<std::slice::IterMut<'a, Slot<T>>>,
generations: &'a [u32],
}
impl<'a, T> Iterator for ArenaIterMut<'a, T> {
type Item = (NodeId, &'a mut T);
fn next(&mut self) -> Option<Self::Item> {
loop {
let (i, slot) = self.inner.next()?;
if let Slot::Occupied(value) = slot {
let id = NodeId {
index: i as u32,
generation: self.generations[i],
};
return Some((id, value));
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn alloc_and_get() {
let mut arena = Arena::new();
let id = arena.alloc(42);
assert_eq!(arena.get(id), Some(&42));
assert_eq!(arena.len(), 1);
}
#[test]
fn alloc_multiple() {
let mut arena = Arena::new();
let id1 = arena.alloc("a");
let id2 = arena.alloc("b");
let id3 = arena.alloc("c");
assert_eq!(arena.get(id1), Some(&"a"));
assert_eq!(arena.get(id2), Some(&"b"));
assert_eq!(arena.get(id3), Some(&"c"));
assert_eq!(arena.len(), 3);
}
#[test]
fn dealloc_returns_value() {
let mut arena = Arena::new();
let id = arena.alloc("hello".to_string());
let removed = arena.dealloc(id);
assert_eq!(removed, Some("hello".to_string()));
assert_eq!(arena.len(), 0);
}
#[test]
fn stale_id_returns_none() {
let mut arena = Arena::new();
let id = arena.alloc(42);
arena.dealloc(id);
assert_eq!(arena.get(id), None);
}
#[test]
fn dealloc_stale_id_returns_none() {
let mut arena = Arena::new();
let id = arena.alloc(42);
arena.dealloc(id);
assert_eq!(arena.dealloc(id), None);
}
#[test]
fn free_list_reuse() {
let mut arena = Arena::new();
let id1 = arena.alloc(1);
let id2 = arena.alloc(2);
arena.dealloc(id1);
let id3 = arena.alloc(3);
assert_eq!(id3.index(), id1.index());
assert_ne!(id3.generation(), id1.generation());
assert_eq!(arena.get(id3), Some(&3));
assert_eq!(arena.get(id1), None); assert_eq!(arena.get(id2), Some(&2));
assert_eq!(arena.len(), 2);
}
#[test]
fn get_mut() {
let mut arena = Arena::new();
let id = arena.alloc(10);
*arena.get_mut(id).unwrap() = 20;
assert_eq!(arena.get(id), Some(&20));
}
#[test]
fn get_mut_stale() {
let mut arena = Arena::new();
let id = arena.alloc(10);
arena.dealloc(id);
assert_eq!(arena.get_mut(id), None);
}
#[test]
fn clear() {
let mut arena = Arena::new();
let id1 = arena.alloc(1);
let id2 = arena.alloc(2);
arena.clear();
assert_eq!(arena.len(), 0);
assert!(arena.is_empty());
assert_eq!(arena.get(id1), None);
assert_eq!(arena.get(id2), None);
}
#[test]
fn is_empty() {
let mut arena: Arena<i32> = Arena::new();
assert!(arena.is_empty());
let id = arena.alloc(1);
assert!(!arena.is_empty());
arena.dealloc(id);
assert!(arena.is_empty());
}
#[test]
fn contains() {
let mut arena = Arena::new();
let id = arena.alloc(42);
assert!(arena.contains(id));
arena.dealloc(id);
assert!(!arena.contains(id));
}
#[test]
fn with_capacity() {
let arena: Arena<i32> = Arena::with_capacity(100);
assert!(arena.capacity() >= 100);
assert!(arena.is_empty());
}
#[test]
fn iter_all_occupied() {
let mut arena = Arena::new();
let id1 = arena.alloc(10);
let id2 = arena.alloc(20);
let id3 = arena.alloc(30);
let items: Vec<_> = arena.iter().collect();
assert_eq!(items.len(), 3);
assert!(items.contains(&(id1, &10)));
assert!(items.contains(&(id2, &20)));
assert!(items.contains(&(id3, &30)));
}
#[test]
fn iter_with_holes() {
let mut arena = Arena::new();
let id1 = arena.alloc(10);
let id2 = arena.alloc(20);
let id3 = arena.alloc(30);
arena.dealloc(id2);
let items: Vec<_> = arena.iter().collect();
assert_eq!(items.len(), 2);
assert!(items.contains(&(id1, &10)));
assert!(items.contains(&(id3, &30)));
}
#[test]
fn iter_empty() {
let arena: Arena<i32> = Arena::new();
assert_eq!(arena.iter().count(), 0);
}
#[test]
fn out_of_bounds_id() {
let arena: Arena<i32> = Arena::new();
let fake_id = NodeId {
index: 999,
generation: 0,
};
assert_eq!(arena.get(fake_id), None);
}
#[test]
fn generation_wrapping() {
let mut arena = Arena::new();
let mut last_id = arena.alloc(0);
for i in 1..10 {
arena.dealloc(last_id);
last_id = arena.alloc(i);
}
assert_eq!(arena.get(last_id), Some(&9));
assert_eq!(arena.len(), 1);
}
#[test]
fn node_id_debug() {
let id = NodeId {
index: 5,
generation: 3,
};
assert_eq!(format!("{id:?}"), "NodeId(5g3)");
}
#[test]
fn default_arena() {
let arena: Arena<i32> = Arena::default();
assert!(arena.is_empty());
}
#[test]
fn stress_alloc_dealloc() {
let mut arena = Arena::new();
let mut ids = Vec::new();
for i in 0..1000 {
ids.push(arena.alloc(i));
}
assert_eq!(arena.len(), 1000);
for i in (0..1000).step_by(2) {
arena.dealloc(ids[i]);
}
assert_eq!(arena.len(), 500);
for i in 0..500 {
arena.alloc(i + 1000);
}
assert_eq!(arena.len(), 1000);
}
#[test]
fn node_id_is_copy_and_eq() {
let id1 = NodeId {
index: 0,
generation: 0,
};
let id2 = id1; assert_eq!(id1, id2);
}
#[test]
fn node_id_hash() {
use std::collections::HashSet;
let mut set = HashSet::new();
let id1 = NodeId {
index: 0,
generation: 0,
};
let id2 = NodeId {
index: 1,
generation: 0,
};
set.insert(id1);
set.insert(id2);
assert_eq!(set.len(), 2);
assert!(set.contains(&id1));
}
}