use super::{ManagedSlice as Slice};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub struct Slot {
generation_id: GenerationOrFreelink,
}
pub struct SlotMap<'a, T> {
elements: Slice<'a, T>,
slots: Partial<'a, Slot>,
generation: Generation,
free_top: usize,
indices: IndexComputer,
}
struct Partial<'a, T> {
slice: Slice<'a, T>,
next_idx: usize,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub struct Key {
idx: usize,
generation: Generation,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
struct GenerationOrFreelink(isize);
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
struct FreeIndex(usize);
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
struct Generation(isize);
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct Offset(isize);
struct IndexComputer(usize);
impl<T> SlotMap<'_, T> {
pub fn get(&self, index: Key) -> Option<&T> {
let slot_generation = self.slots
.get(index.idx)?
.generation_id
.generation().ok()?;
if slot_generation != index.generation {
return None;
}
self.elements.get(index.idx)
}
pub fn get_mut(&mut self, index: Key) -> Option<&mut T> {
let slot_generation = self.slots
.get(index.idx)?
.generation_id
.generation().ok()?;
if slot_generation != index.generation {
return None;
}
self.elements.get_mut(index.idx)
}
pub fn reserve(&mut self) -> Option<(Key, &mut T)> {
let index = self.next_free_slot()?;
let slot = self.slots.get_mut(index.0).unwrap();
let element = &mut self.elements[index.0];
let offset = slot.generation_id
.free_link()
.expect("Free link should be free");
slot.generation_id = self.generation.into();
let key = Key {
idx: index.0,
generation: self.generation,
};
self.free_top = self.indices.free_list_next(index, offset);
self.generation.advance();
Some((key, element))
}
pub fn insert(&mut self, value: T) -> Option<Key> {
let (index, element) = self.reserve()?;
*element = value;
Some(index)
}
pub fn remove(&mut self, index: Key) -> Option<&mut T> {
if self.get(index).is_none() {
return None
}
let free = FreeIndex(index.idx);
let slot = self.slots.get_mut(index.idx).unwrap();
assert!(slot.generation_id.generation().is_ok());
let offset = self.indices.free_list_offset(free, self.free_top);
slot.generation_id = offset.into();
self.free_top = index.idx;
Some(&mut self.elements[index.idx])
}
fn next_free_slot(&mut self) -> Option<FreeIndex> {
let free = match self.slots.get_mut(self.free_top) {
Some(_) => {
let _= self.elements.get_mut(self.free_top)?;
FreeIndex(self.free_top)
},
None => {
let new_index = self.slots.len();
let _ = self.elements.get_mut(new_index)?;
let free_slot = self.slots.try_reserve()?;
let free_index = FreeIndex(new_index);
let new_top = new_index.checked_add(1).unwrap();
let offset = self.indices.free_list_offset(free_index, new_top);
free_slot.generation_id = offset.into();
self.free_top = free_index.0;
free_index
}
};
Some(free)
}
}
impl<'a, T> SlotMap<'a, T> {
pub fn new(elements: Slice<'a, T>, slots: Slice<'a, Slot>) -> Self {
let capacity = elements.len().min(slots.len());
SlotMap {
elements,
slots: Partial {
slice: slots,
next_idx: 0,
},
generation: Generation::default(),
free_top: 0,
indices: IndexComputer::from_capacity(capacity),
}
}
}
impl<'a, T> Partial<'a, T> {
fn get(&self, idx: usize) -> Option<&T> {
if idx >= self.next_idx {
None
} else {
Some(&self.slice[idx])
}
}
fn get_mut(&mut self, idx: usize) -> Option<&mut T> {
if idx >= self.next_idx {
None
} else {
Some(&mut self.slice[idx])
}
}
fn len(&self) -> usize {
self.next_idx
}
fn try_reserve(&mut self) -> Option<&mut T> {
if self.next_idx == self.slice.len() {
None
} else {
let idx = self.next_idx;
self.next_idx += 1;
Some(&mut self.slice[idx])
}
}
}
impl GenerationOrFreelink {
pub(crate) fn free_link(self) -> Result<Offset, Generation> {
if self.0 > 0 {
Err(Generation(self.0))
} else {
Ok(Offset(self.0))
}
}
pub(crate) fn generation(self) -> Result<Generation, Offset> {
match self.free_link() {
Ok(offset) => Err(offset),
Err(generation) => Ok(generation),
}
}
}
impl IndexComputer {
pub(crate) fn from_capacity(capacity: usize) -> Self {
assert!(capacity < isize::max_value() as usize);
IndexComputer(capacity)
}
fn free_list_next(&self, FreeIndex(base): FreeIndex, offset: Offset)
-> usize
{
let capacity = self.0;
let offset = offset.int_offset();
assert!(base < capacity);
assert!(offset <= capacity);
let base = base + 1;
if capacity - offset >= base {
offset + base
} else {
offset
.wrapping_add(base)
.wrapping_sub(capacity + 1)
}
}
fn free_list_offset(&self, FreeIndex(base): FreeIndex, to: usize)
-> Offset
{
let capacity = self.0;
assert!(base != to, "Cant offset element to itself");
assert!(base < capacity, "Should never have to offset the end-of-list marker");
assert!(to <= capacity, "Can only offset to the end-of-list marker");
let base = base + 1;
Offset::from_int_offset(if base <= to {
to - base
} else {
to
.wrapping_add(capacity + 1)
.wrapping_sub(base)
})
}
}
impl Generation {
pub(crate) fn advance(&mut self) {
assert!(self.0 > 0);
self.0 = self.0.wrapping_add(1).max(1)
}
}
impl Offset {
pub(crate) fn from_int_offset(offset: usize) -> Self {
assert!(offset < isize::max_value() as usize);
Offset((offset as isize).checked_neg().unwrap())
}
pub(crate) fn int_offset(self) -> usize {
self.0.checked_neg().unwrap() as usize
}
}
impl Default for Generation {
fn default() -> Self {
Generation(1)
}
}
impl From<Generation> for GenerationOrFreelink {
fn from(gen: Generation) -> GenerationOrFreelink {
GenerationOrFreelink(gen.0)
}
}
impl From<Offset> for GenerationOrFreelink {
fn from(offset: Offset) -> GenerationOrFreelink {
GenerationOrFreelink(offset.0)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::slice::ManagedSlice as Slice;
#[test]
fn simple() {
let mut elements = [0u32; 2];
let mut slots = [Slot::default(); 2];
let mut map = SlotMap::new(
Slice::Borrowed(&mut elements[..]),
Slice::Borrowed(&mut slots[..]));
let key42 = map.insert(42).unwrap();
let keylo = map.insert('K' as _).unwrap();
assert_eq!(map.insert(0x9999), None);
assert_eq!(map.get(key42).cloned(), Some(42));
assert_eq!(map.get(keylo).cloned(), Some('K' as _));
assert!(map.remove(key42).is_some());
assert_eq!(map.get(key42), None);
let lastkey = map.insert(0x9999).unwrap();
assert_eq!(map.get(lastkey).cloned(), Some(0x9999));
*map.remove(keylo).unwrap() = 0;
assert_eq!(map.get(lastkey).cloned(), Some(0x9999));
assert!(map.remove(lastkey).is_some());
}
#[test]
fn retained() {
let mut elements = [0u32; 1];
let mut slots = [Slot::default(); 1];
let mut map = SlotMap::new(
Slice::Borrowed(&mut elements[..]),
Slice::Borrowed(&mut slots[..]));
let key = map.insert(0xde).unwrap();
map.remove(key).unwrap();
assert_eq!(map.get(key), None);
let new_key = map.insert(0xad).unwrap();
assert_eq!(map.get(key), None);
assert_eq!(map.get(new_key).cloned(), Some(0xad));
assert_eq!(map.remove(key), None);
map.remove(new_key).unwrap();
assert_eq!(map.get(key), None);
assert_eq!(map.get(new_key), None);
}
#[test]
fn non_simple_free_list() {
let mut elements = [0u32; 3];
let mut slots = [Slot::default(); 3];
let mut map = SlotMap::new(
Slice::Borrowed(&mut elements[..]),
Slice::Borrowed(&mut slots[..]));
let key0 = map.insert(0).unwrap();
let key1 = map.insert(1).unwrap();
let key2 = map.insert(2).unwrap();
*map.remove(key1).unwrap() = 0xF;
assert_eq!(map.free_top, 1);
assert_eq!(map.get(key0).cloned(), Some(0));
assert_eq!(map.get(key2).cloned(), Some(2));
*map.remove(key2).unwrap() = 0xF;
assert_eq!(map.free_top, 2);
assert_eq!(map.get(key0).cloned(), Some(0));
*map.remove(key0).unwrap() = 0xF;
assert_eq!(map.free_top, 0);
let key0 = map.insert(0).unwrap();
assert_eq!(map.free_top, 2);
let key1 = map.insert(1).unwrap();
let key2 = map.insert(2).unwrap();
assert_eq!(map.get(key0).cloned(), Some(0));
assert_eq!(map.get(key1).cloned(), Some(1));
assert_eq!(map.get(key2).cloned(), Some(2));
}
}