use core::convert::TryInto;
use core::mem::replace;
use core::ops;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
use crate::free_pointer::FreePointer;
use crate::generation::Generation;
use crate::iter::{Drain, IntoIter, Iter, IterMut};
#[derive(Debug, Clone)]
pub struct Arena<T> {
storage: Vec<Entry<T>>,
len: u32,
first_free: Option<FreePointer>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct Index {
pub(crate) slot: u32,
pub(crate) generation: Generation,
}
impl Index {
#[allow(clippy::integer_arithmetic)]
pub const fn to_bits(self) -> u64 {
((self.generation.to_u32() as u64) << 32) | (self.slot as u64)
}
#[allow(clippy::integer_arithmetic)]
pub const fn from_bits(bits: u64) -> Option<Self> {
let generation = match Generation::from_u32((bits >> 32) as u32) {
Some(v) => v,
None => return None,
};
let slot = bits as u32;
Some(Self { generation, slot })
}
pub const fn generation(self) -> u32 {
self.generation.to_u32()
}
pub const fn slot(self) -> u32 {
self.slot
}
}
#[derive(Debug, Clone)]
pub(crate) enum Entry<T> {
Occupied(OccupiedEntry<T>),
Empty(EmptyEntry),
}
impl<T> Entry<T> {
fn into_value(self) -> Option<T> {
match self {
Entry::Occupied(occupied) => Some(occupied.value),
Entry::Empty(_) => None,
}
}
fn as_empty(&self) -> Option<&EmptyEntry> {
match self {
Entry::Empty(empty) => Some(empty),
Entry::Occupied(_) => None,
}
}
fn as_empty_mut(&mut self) -> Option<&mut EmptyEntry> {
match self {
Entry::Empty(empty) => Some(empty),
Entry::Occupied(_) => None,
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct OccupiedEntry<T> {
pub(crate) generation: Generation,
pub(crate) value: T,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct EmptyEntry {
pub(crate) generation: Generation,
pub(crate) next_free: Option<FreePointer>,
}
impl<T> Arena<T> {
pub const fn new() -> Self {
Self {
storage: Vec::new(),
len: 0,
first_free: None,
}
}
pub fn with_capacity(capacity: usize) -> Self {
Self {
storage: Vec::with_capacity(capacity),
len: 0,
first_free: None,
}
}
pub const fn len(&self) -> usize {
self.len as usize
}
pub fn capacity(&self) -> usize {
self.storage.capacity()
}
pub const fn is_empty(&self) -> bool {
self.len == 0
}
pub fn insert(&mut self, value: T) -> Index {
self.len = self
.len
.checked_add(1)
.unwrap_or_else(|| panic!("Cannot insert more than u32::MAX elements into Arena"));
if let Some(free_pointer) = self.first_free {
let slot = free_pointer.slot();
let entry = self.storage.get_mut(slot as usize).unwrap_or_else(|| {
unreachable!("first_free pointed past the end of the arena's storage")
});
let empty = entry
.as_empty()
.unwrap_or_else(|| unreachable!("first_free pointed to an occupied entry"));
self.first_free = empty.next_free;
let generation = empty.generation.next();
*entry = Entry::Occupied(OccupiedEntry { generation, value });
Index { slot, generation }
} else {
let generation = Generation::first();
let slot: u32 = self.storage.len().try_into().unwrap_or_else(|_| {
unreachable!("Arena storage exceeded what can be represented by a u32")
});
self.storage
.push(Entry::Occupied(OccupiedEntry { generation, value }));
Index { slot, generation }
}
}
fn remove_slot_from_free_list(&mut self, slot: u32, new_next_free: Option<FreePointer>) {
let mut next_fp = self
.first_free
.expect("Free entry exists but first_free is None");
let mut current_slot = None;
while next_fp.slot() != slot {
current_slot = Some(next_fp.slot());
next_fp = self
.storage
.get(next_fp.slot() as usize)
.expect("Empty entry not in storage!")
.as_empty()
.expect("Entry in free list not actually empty!")
.next_free
.expect("Hit the end of the free list without finding the target slot!");
}
match current_slot {
Some(slot_to_fix) => {
self.storage[slot_to_fix as usize]
.as_empty_mut()
.unwrap()
.next_free = new_next_free
}
None => self.first_free = new_next_free,
}
}
#[inline]
fn insert_at_inner(
&mut self,
slot: u32,
generation: Option<Generation>,
value: T,
) -> (Index, Option<T>) {
let (index, old_value) = match self.storage.get_mut(slot as usize) {
Some(Entry::Empty(empty)) => {
let generation = generation.unwrap_or_else(|| empty.generation.next());
let new_next_free = empty.next_free;
self.remove_slot_from_free_list(slot, new_next_free);
self.storage[slot as usize] = Entry::Occupied(OccupiedEntry { generation, value });
(Index { slot, generation }, None)
}
Some(Entry::Occupied(occupied)) => {
occupied.generation = generation.unwrap_or_else(|| occupied.generation.next());
let generation = occupied.generation;
let old_value = replace(&mut occupied.value, value);
(Index { slot, generation }, Some(old_value))
}
None => {
let mut first_free = self.first_free;
while self.storage.len() < slot as usize {
let new_slot: u32 = self.storage.len().try_into().unwrap_or_else(|_| {
unreachable!("Arena storage exceeded what can be represented by a u32")
});
self.storage.push(Entry::Empty(EmptyEntry {
generation: Generation::first(),
next_free: first_free,
}));
first_free = Some(FreePointer::from_slot(new_slot));
}
self.first_free = first_free;
let generation = generation.unwrap_or_else(Generation::first);
self.storage
.push(Entry::Occupied(OccupiedEntry { generation, value }));
(Index { slot, generation }, None)
}
};
if old_value.is_none() {
self.len = self
.len
.checked_add(1)
.unwrap_or_else(|| panic!("Cannot insert more than u32::MAX elements into Arena"));
}
(index, old_value)
}
pub fn insert_at(&mut self, index: Index, value: T) -> Option<T> {
self.insert_at_inner(index.slot, Some(index.generation), value)
.1
}
pub fn insert_at_slot(&mut self, slot: u32, value: T) -> (Index, Option<T>) {
self.insert_at_inner(slot, None, value)
}
pub fn contains(&self, index: Index) -> bool {
match self.storage.get(index.slot as usize) {
Some(Entry::Occupied(occupied)) if occupied.generation == index.generation => true,
_ => false,
}
}
pub fn contains_slot(&self, slot: u32) -> Option<Index> {
match self.storage.get(slot as usize) {
Some(Entry::Occupied(occupied)) => Some(Index {
slot,
generation: occupied.generation,
}),
_ => None,
}
}
pub fn get(&self, index: Index) -> Option<&T> {
match self.storage.get(index.slot as usize) {
Some(Entry::Occupied(occupied)) if occupied.generation == index.generation => {
Some(&occupied.value)
}
_ => None,
}
}
pub fn get_mut(&mut self, index: Index) -> Option<&mut T> {
match self.storage.get_mut(index.slot as usize) {
Some(Entry::Occupied(occupied)) if occupied.generation == index.generation => {
Some(&mut occupied.value)
}
_ => None,
}
}
pub fn get2_mut(&mut self, index1: Index, index2: Index) -> (Option<&mut T>, Option<&mut T>) {
if index1 == index2 {
panic!("Arena::get2_mut is called with two identical indices");
}
let item1_ptr = self.get_mut(index1).map(|x| x as *mut T);
let item2_ptr = self.get_mut(index2).map(|x| x as *mut T);
let item1 = unsafe { item1_ptr.map(|x| &mut *x) };
let item2 = unsafe { item2_ptr.map(|x| &mut *x) };
(item1, item2)
}
pub fn remove(&mut self, index: Index) -> Option<T> {
let entry = self.storage.get_mut(index.slot as usize)?;
match entry {
Entry::Occupied(occupied) if occupied.generation == index.generation => {
let new_entry = Entry::Empty(EmptyEntry {
generation: occupied.generation,
next_free: self.first_free,
});
let old_entry = replace(entry, new_entry);
let value = old_entry.into_value().unwrap_or_else(|| unreachable!());
self.first_free = Some(FreePointer::from_slot(index.slot));
self.len = self.len.checked_sub(1).unwrap_or_else(|| unreachable!());
Some(value)
}
_ => None,
}
}
pub fn invalidate(&mut self, index: Index) -> Option<Index> {
let entry = self.storage.get_mut(index.slot as usize)?;
match entry {
Entry::Occupied(occupied) if occupied.generation == index.generation => {
occupied.generation = occupied.generation.next();
Some(Index {
generation: occupied.generation,
..index
})
}
_ => None,
}
}
pub fn get_by_slot(&self, slot: u32) -> Option<(Index, &T)> {
match self.storage.get(slot as usize) {
Some(Entry::Occupied(occupied)) => {
let index = Index {
slot,
generation: occupied.generation,
};
Some((index, &occupied.value))
}
_ => None,
}
}
pub fn get_by_slot_mut(&mut self, slot: u32) -> Option<(Index, &mut T)> {
match self.storage.get_mut(slot as usize) {
Some(Entry::Occupied(occupied)) => {
let index = Index {
slot,
generation: occupied.generation,
};
Some((index, &mut occupied.value))
}
_ => None,
}
}
pub fn remove_by_slot(&mut self, slot: u32) -> Option<(Index, T)> {
let entry = self.storage.get_mut(slot as usize)?;
match entry {
Entry::Occupied(occupied) => {
let index = Index {
generation: occupied.generation,
slot,
};
let next_entry = Entry::Empty(EmptyEntry {
generation: occupied.generation,
next_free: self.first_free,
});
let old_entry = replace(entry, next_entry);
let value = old_entry.into_value().unwrap_or_else(|| unreachable!());
self.first_free = Some(FreePointer::from_slot(slot));
self.len = self.len.checked_sub(1).unwrap_or_else(|| unreachable!());
Some((index, value))
}
_ => None,
}
}
pub fn clear(&mut self) {
self.drain().for_each(drop);
}
pub fn iter(&self) -> Iter<'_, T> {
Iter {
inner: self.storage.iter(),
slot: 0,
len: self.len,
}
}
pub fn iter_mut(&mut self) -> IterMut<'_, T> {
IterMut {
inner: self.storage.iter_mut(),
slot: 0,
len: self.len,
}
}
pub fn drain(&mut self) -> Drain<'_, T> {
Drain {
arena: self,
slot: 0,
}
}
pub fn retain<F: FnMut(Index, &mut T) -> bool>(&mut self, mut f: F) {
for (i, entry) in self.storage.iter_mut().enumerate() {
if let Entry::Occupied(occupied) = entry {
let index = Index {
slot: i as u32,
generation: occupied.generation,
};
if !f(index, &mut occupied.value) {
*entry = Entry::Empty(EmptyEntry {
generation: occupied.generation,
next_free: self.first_free,
});
self.first_free = Some(FreePointer::from_slot(index.slot));
self.len = self.len.checked_sub(1).unwrap_or_else(|| unreachable!());
}
}
}
}
}
impl<T> Default for Arena<T> {
fn default() -> Self {
Arena::new()
}
}
impl<T> IntoIterator for Arena<T> {
type Item = (Index, T);
type IntoIter = IntoIter<T>;
fn into_iter(self) -> Self::IntoIter {
IntoIter {
arena: self,
slot: 0,
}
}
}
impl<'a, T> IntoIterator for &'a Arena<T> {
type Item = (Index, &'a T);
type IntoIter = Iter<'a, T>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'a, T> IntoIterator for &'a mut Arena<T> {
type Item = (Index, &'a mut T);
type IntoIter = IterMut<'a, T>;
fn into_iter(self) -> Self::IntoIter {
self.iter_mut()
}
}
impl<T> ops::Index<Index> for Arena<T> {
type Output = T;
fn index(&self, index: Index) -> &Self::Output {
self.get(index)
.unwrap_or_else(|| panic!("No entry at index {:?}", index))
}
}
impl<T> ops::IndexMut<Index> for Arena<T> {
fn index_mut(&mut self, index: Index) -> &mut Self::Output {
self.get_mut(index)
.unwrap_or_else(|| panic!("No entry at index {:?}", index))
}
}
#[cfg(test)]
mod test {
use crate::free_pointer::FreePointer;
use super::{Arena, Generation, Index};
use core::mem::size_of;
#[test]
fn size_of_index() {
assert_eq!(size_of::<Index>(), 8);
assert_eq!(size_of::<Option<Index>>(), 8);
}
#[test]
fn new() {
let arena: Arena<u32> = Arena::new();
assert_eq!(arena.len(), 0);
assert_eq!(arena.capacity(), 0);
}
#[test]
fn with_capacity() {
let arena: Arena<u32> = Arena::with_capacity(8);
assert_eq!(arena.len(), 0);
assert_eq!(arena.capacity(), 8);
}
#[test]
fn insert_and_get() {
let mut arena = Arena::new();
let one = arena.insert(1);
assert_eq!(arena.len(), 1);
assert_eq!(arena.get(one), Some(&1));
let two = arena.insert(2);
assert_eq!(arena.len(), 2);
assert_eq!(arena.get(one), Some(&1));
assert_eq!(arena.get(two), Some(&2));
}
#[test]
fn insert_remove_get() {
let mut arena = Arena::new();
let one = arena.insert(1);
let two = arena.insert(2);
assert_eq!(arena.len(), 2);
assert!(arena.contains(two));
assert_eq!(arena.remove(two), Some(2));
assert!(!arena.contains(two));
let three = arena.insert(3);
assert_eq!(arena.len(), 2);
assert_eq!(arena.get(one), Some(&1));
assert_eq!(arena.get(three), Some(&3));
assert_eq!(arena.get(two), None);
}
#[test]
fn insert_remove_get_by_slot() {
let mut arena = Arena::new();
let one = arena.insert(1);
let two = arena.insert(2);
assert_eq!(arena.len(), 2);
assert!(arena.contains(two));
assert_eq!(arena.remove_by_slot(two.slot()), Some((two, 2)));
assert!(!arena.contains(two));
assert_eq!(arena.get_by_slot(two.slot()), None);
let three = arena.insert(3);
assert_eq!(arena.len(), 2);
assert_eq!(arena.get(one), Some(&1));
assert_eq!(arena.get(three), Some(&3));
assert_eq!(arena.get(two), None);
assert_eq!(arena.get_by_slot(two.slot()), Some((three, &3)));
}
#[test]
fn insert_at() {
let mut arena = Arena::new();
let index = Index {
slot: 42,
generation: Generation::from_u32(78).unwrap(),
};
arena.insert_at(index, 5);
assert_eq!(arena.len(), 1);
assert_eq!(arena.get(index), Some(&5));
assert_eq!(arena.get_by_slot(42), Some((index, &5)));
}
#[test]
fn insert_at_first_slot() {
let mut arena = Arena::new();
let index = Index {
slot: 0,
generation: Generation::from_u32(3).unwrap(),
};
arena.insert_at(index, 5);
assert_eq!(arena.len(), 1);
assert_eq!(arena.get(index), Some(&5));
assert_eq!(arena.get_by_slot(0), Some((index, &5)));
}
#[test]
fn insert_at_slot() {
let mut arena = Arena::new();
let (index, _) = arena.insert_at_slot(42, 5);
assert_eq!(arena.len(), 1);
assert_eq!(arena.get(index), Some(&5));
assert_eq!(arena.get_by_slot(42), Some((index, &5)));
}
#[test]
fn insert_at_middle() {
let mut arena = Arena::new();
arena.insert_at_slot(4, 50);
arena.insert_at_slot(2, 40);
let empty = arena.storage.get(3).unwrap().as_empty().unwrap();
if empty.next_free != Some(FreePointer::from_slot(1)) {
panic!("Invalid free list: {:#?}", arena);
}
}
#[test]
fn get_mut() {
let mut arena = Arena::new();
let foo = arena.insert(5);
let handle = arena.get_mut(foo).unwrap();
*handle = 6;
assert_eq!(arena.get(foo), Some(&6));
}
#[test]
fn get2_mut() {
let mut arena = Arena::new();
let foo = arena.insert(100);
let bar = arena.insert(500);
let (foo_handle, bar_handle) = arena.get2_mut(foo, bar);
let foo_handle = foo_handle.unwrap();
let bar_handle = bar_handle.unwrap();
*foo_handle = 105;
*bar_handle = 505;
assert_eq!(arena.get(foo), Some(&105));
assert_eq!(arena.get(bar), Some(&505));
}
#[test]
fn get2_mut_reversed_order() {
let mut arena = Arena::new();
let foo = arena.insert(100);
let bar = arena.insert(500);
let (bar_handle, foo_handle) = arena.get2_mut(bar, foo);
let foo_handle = foo_handle.unwrap();
let bar_handle = bar_handle.unwrap();
*foo_handle = 105;
*bar_handle = 505;
assert_eq!(arena.get(foo), Some(&105));
assert_eq!(arena.get(bar), Some(&505));
}
#[test]
fn get2_mut_non_exist_handle() {
let mut arena = Arena::new();
let foo = arena.insert(100);
let bar = arena.insert(500);
arena.remove(bar);
let (bar_handle, foo_handle) = arena.get2_mut(bar, foo);
let foo_handle = foo_handle.unwrap();
assert!(bar_handle.is_none());
*foo_handle = 105;
assert_eq!(arena.get(foo), Some(&105));
}
#[test]
fn get2_mut_same_slot_different_generation() {
let mut arena = Arena::new();
let foo = arena.insert(100);
let mut foo1 = foo;
foo1.generation = foo1.generation.next();
let (foo_handle, foo1_handle) = arena.get2_mut(foo, foo1);
assert!(foo_handle.is_some());
assert!(foo1_handle.is_none());
}
#[test]
#[should_panic]
fn get2_mut_panics() {
let mut arena = Arena::new();
let foo = arena.insert(100);
arena.get2_mut(foo, foo);
}
#[test]
fn insert_remove_insert_capacity() {
let mut arena = Arena::with_capacity(2);
assert_eq!(arena.capacity(), 2);
let a = arena.insert("a");
let b = arena.insert("b");
assert_eq!(arena.len(), 2);
assert_eq!(arena.capacity(), 2);
arena.remove(a);
arena.remove(b);
assert_eq!(arena.len(), 0);
assert_eq!(arena.capacity(), 2);
let _a2 = arena.insert("a2");
let _b2 = arena.insert("b2");
assert_eq!(arena.len(), 2);
assert_eq!(arena.capacity(), 2);
}
#[test]
fn invalidate() {
let mut arena = Arena::new();
let a = arena.insert("a");
assert_eq!(arena.get(a), Some(&"a"));
let new_a = arena.invalidate(a).unwrap();
assert_eq!(arena.get(a), None);
assert_eq!(arena.get(new_a), Some(&"a"));
}
#[test]
fn retain() {
let mut arena = Arena::new();
for i in 0..100 {
arena.insert(i);
}
arena.retain(|_, &mut i| i % 2 == 1);
for (_, i) in arena.iter() {
assert_eq!(i % 2, 1);
}
assert_eq!(arena.len(), 50);
}
#[test]
fn index_bits_roundtrip() {
let index = Index::from_bits(0x1BAD_CAFE_DEAD_BEEF).unwrap();
assert_eq!(index.to_bits(), 0x1BAD_CAFE_DEAD_BEEF);
}
#[test]
fn index_bits_none_on_zero_generation() {
let index = Index::from_bits(0x0000_0000_DEAD_BEEF);
assert_eq!(index, None);
}
}