use alloc::vec::Vec;
use core::marker::PhantomData;
pub struct Id<T> {
index: u32,
generation: u32,
_marker: PhantomData<fn() -> T>,
}
impl<T> Id<T> {
#[inline]
#[must_use]
pub fn index(self) -> u32 {
self.index
}
#[inline]
#[must_use]
pub fn generation(self) -> u32 {
self.generation
}
}
impl<T> Clone for Id<T> {
#[inline]
fn clone(&self) -> Self {
*self
}
}
impl<T> Copy for Id<T> {}
impl<T> PartialEq for Id<T> {
#[inline]
fn eq(&self, other: &Self) -> bool {
self.index == other.index && self.generation == other.generation
}
}
impl<T> Eq for Id<T> {}
impl<T> PartialOrd for Id<T> {
#[inline]
fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl<T> Ord for Id<T> {
#[inline]
fn cmp(&self, other: &Self) -> core::cmp::Ordering {
(self.index, self.generation).cmp(&(other.index, other.generation))
}
}
impl<T> core::hash::Hash for Id<T> {
#[inline]
fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
core::hash::Hash::hash(&self.index, state);
core::hash::Hash::hash(&self.generation, state);
}
}
impl<T> core::fmt::Debug for Id<T> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "Id({}@g{})", self.index, self.generation)
}
}
#[derive(Clone, Debug)]
struct Slot<T> {
generation: u32,
value: Option<T>,
}
#[derive(Clone, Debug)]
pub struct Arena<T> {
slots: Vec<Slot<T>>,
free: Vec<u32>,
len: usize,
}
impl<T> Arena<T> {
#[must_use]
pub fn new() -> Self {
Self {
slots: Vec::new(),
free: Vec::new(),
len: 0,
}
}
#[must_use]
pub fn with_capacity(capacity: usize) -> Self {
Self {
slots: Vec::with_capacity(capacity),
free: Vec::new(),
len: 0,
}
}
pub fn insert(&mut self, value: T) -> Id<T> {
self.len += 1;
if let Some(index) = self.free.pop() {
let slot = &mut self.slots[index as usize];
slot.generation = slot.generation.wrapping_add(1);
slot.value = Some(value);
return Id {
index,
generation: slot.generation,
_marker: PhantomData,
};
}
let index = u32::try_from(self.slots.len()).expect("arena index space exhausted");
self.slots.push(Slot {
generation: 0,
value: Some(value),
});
Id {
index,
generation: 0,
_marker: PhantomData,
}
}
#[must_use]
pub fn get(&self, id: Id<T>) -> Option<&T> {
let slot = self.slots.get(id.index as usize)?;
if slot.generation != id.generation {
return None;
}
slot.value.as_ref()
}
#[must_use]
pub fn get_mut(&mut self, id: Id<T>) -> Option<&mut T> {
let slot = self.slots.get_mut(id.index as usize)?;
if slot.generation != id.generation {
return None;
}
slot.value.as_mut()
}
#[must_use]
pub fn contains(&self, id: Id<T>) -> bool {
self.get(id).is_some()
}
#[must_use]
pub fn remove(&mut self, id: Id<T>) -> Option<T> {
let slot = self.slots.get_mut(id.index as usize)?;
if slot.generation != id.generation {
return None;
}
let value = slot.value.take()?;
slot.generation = slot.generation.wrapping_add(1);
self.free.push(id.index);
self.len -= 1;
Some(value)
}
#[must_use]
pub fn len(&self) -> usize {
self.len
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len == 0
}
pub fn iter(&self) -> impl Iterator<Item = (Id<T>, &T)> {
self.slots.iter().enumerate().filter_map(|(index, slot)| {
let value = slot.value.as_ref()?;
let index = u32::try_from(index).ok()?;
Some((
Id {
index,
generation: slot.generation,
_marker: PhantomData,
},
value,
))
})
}
pub fn iter_mut(&mut self) -> impl Iterator<Item = (Id<T>, &mut T)> {
self
.slots
.iter_mut()
.enumerate()
.filter_map(|(index, slot)| {
let value = slot.value.as_mut()?;
let index = u32::try_from(index).ok()?;
Some((
Id {
index,
generation: slot.generation,
_marker: PhantomData,
},
value,
))
})
}
}
impl<T> Default for Arena<T> {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use alloc::vec::Vec;
use super::Arena;
#[test]
fn insert_and_get() {
let mut arena = Arena::new();
let id = arena.insert(42u32);
assert_eq!(arena.get(id), Some(&42));
assert!(arena.contains(id));
assert_eq!(arena.len(), 1);
}
#[test]
fn removed_id_becomes_stale() {
let mut arena = Arena::new();
let id = arena.insert(7u32);
assert_eq!(arena.remove(id), Some(7));
assert!(!arena.contains(id));
assert_eq!(arena.get(id), None);
assert!(arena.is_empty());
}
#[test]
fn reused_slot_invalidates_previous_handle() {
let mut arena = Arena::new();
let first = arena.insert(1u32);
let _ = arena.remove(first);
let second = arena.insert(2u32);
assert_eq!(first.index(), second.index());
assert_ne!(first, second);
assert_eq!(arena.get(first), None);
assert_eq!(arena.get(second), Some(&2));
}
#[test]
fn iter_yields_live_pairs() {
let mut arena = Arena::new();
let a = arena.insert("a");
let b = arena.insert("b");
let c = arena.insert("c");
let _ = arena.remove(b);
let mut seen: Vec<&str> = arena.iter().map(|(_, v)| *v).collect();
seen.sort_unstable();
assert_eq!(seen, ["a", "c"]);
assert!(arena.contains(a) && arena.contains(c));
}
}