use std::prelude::v1::*;
use std::marker::PhantomData;
use std::{vec, slice, iter};
use crate::gc::*;
pub trait Key: Copy + Eq + Ord + 'static {
fn new(slot: usize, generation: usize) -> Self;
fn get_slot(&self) -> usize;
fn get_generation(&self) -> usize;
}
#[macro_export]
macro_rules! new_key {
($($(#[doc = $doc:expr])? $vis:vis struct $name:ident;)*) => {$(
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, $crate::gc::Collect)]
#[collect(require_static)]
$(#[doc = $doc])?
$vis struct $name(usize, usize);
impl $crate::slotmap::Key for $name {
fn new(slot: usize, generation: usize) -> Self { Self(slot, generation) }
fn get_slot(&self) -> usize { self.0 }
fn get_generation(&self) -> usize { self.1 }
}
)*}
}
#[derive(Clone, Collect)]
#[collect(no_drop)]
struct Slot<T> {
value: Option<T>,
generation: usize,
}
#[derive(Clone, Collect)]
#[collect(no_drop)]
pub struct SlotMap<K: Key, T> {
slots: Vec<Slot<T>>,
empty_slots: Vec<usize>,
num_values: usize,
_key: PhantomData<K>
}
impl<K: Key, T> Default for SlotMap<K, T> {
fn default() -> Self {
Self::new()
}
}
impl<K: Key, T> SlotMap<K, T> {
pub fn new() -> Self {
SlotMap {
slots: vec![],
empty_slots: vec![],
num_values: 0,
_key: PhantomData,
}
}
#[cfg(test)]
fn invariant(&self) -> bool {
self.num_values == self.slots.iter().filter(|x| x.value.is_some()).count()
&&
self.num_values + self.empty_slots.len() == self.slots.len()
}
pub fn insert(&mut self, value: T) -> K {
#[cfg(test)] assert!(self.invariant());
self.num_values += 1;
let key = match self.empty_slots.pop() {
Some(slot) => {
debug_assert!(self.slots[slot].value.is_none());
self.slots[slot].value = Some(value);
K::new(slot, self.slots[slot].generation)
}
None => {
debug_assert!(self.slots.iter().all(|x| x.value.is_some()));
let slot = self.slots.len();
self.slots.push(Slot { value: Some(value), generation: 0 });
K::new(slot, 0)
}
};
#[cfg(test)] assert!(self.invariant());
#[allow(clippy::let_and_return)] key
}
pub fn remove(&mut self, key: K) -> Option<T> {
#[cfg(test)] assert!(self.invariant());
let slot = self.slots.get_mut(key.get_slot())?;
let res = slot.value.take();
if res.is_some() {
slot.generation += 1;
self.num_values -= 1;
self.empty_slots.push(key.get_slot());
}
#[cfg(test)] assert!(self.invariant());
res
}
pub fn clear(&mut self) {
#[cfg(test)] assert!(self.invariant());
for (i, slot) in self.slots.iter_mut().enumerate() {
if slot.value.take().is_some() {
slot.generation += 1;
self.empty_slots.push(i);
}
}
self.num_values = 0;
#[cfg(test)] assert!(self.invariant());
}
pub fn get(&self, key: K) -> Option<&T> {
let slot = self.slots.get(key.get_slot())?;
if slot.generation == key.get_generation() { slot.value.as_ref() } else { None }
}
pub fn get_mut(&mut self, key: K) -> Option<&mut T> {
let slot = self.slots.get_mut(key.get_slot())?;
if slot.generation == key.get_generation() { slot.value.as_mut() } else { None }
}
pub fn len(&self) -> usize {
self.num_values
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn iter(&self) -> Iter<K, T> {
Iter(self.slots.iter().enumerate(), PhantomData)
}
pub fn iter_mut(&mut self) -> IterMut<K, T> {
IterMut(self.slots.iter_mut().enumerate(), PhantomData)
}
}
impl<K: Key, T> IntoIterator for SlotMap<K, T> {
type IntoIter = IntoIter<K, T>;
type Item = (K, T);
fn into_iter(self) -> IntoIter<K, T> {
IntoIter(self.slots.into_iter().enumerate(), PhantomData)
}
}
pub struct IntoIter<K: Key, T>(iter::Enumerate<vec::IntoIter<Slot<T>>>, PhantomData<K>);
pub struct Iter<'a, K: Key, T>(iter::Enumerate<slice::Iter<'a, Slot<T>>>, PhantomData<K>);
pub struct IterMut<'a, K: Key, T>(iter::Enumerate<slice::IterMut<'a, Slot<T>>>, PhantomData<K>);
impl<K: Key, T> Iterator for IntoIter<K, T> {
type Item = (K, T);
fn next(&mut self) -> Option<Self::Item> {
loop {
let (i, slot) = self.0.next()?;
if let Some(x) = slot.value { return Some((K::new(i, slot.generation), x)) }
}
}
}
impl<'a, K: Key, T> Iterator for Iter<'a, K, T> {
type Item = (K, &'a T);
fn next(&mut self) -> Option<Self::Item> {
loop {
let (i, slot) = self.0.next()?;
if let Some(x) = slot.value.as_ref() { return Some((K::new(i, slot.generation), x)) }
}
}
}
impl<'a, K: Key, T> Iterator for IterMut<'a, K, T> {
type Item = (K, &'a mut T);
fn next(&mut self) -> Option<Self::Item> {
loop {
let (i, slot) = self.0.next()?;
if let Some(x) = slot.value.as_mut() { return Some((K::new(i, slot.generation), x)) }
}
}
}
#[test]
fn test_slotmap() {
use std::collections::BTreeSet;
new_key! {
struct TestKey;
}
let mut map: SlotMap<TestKey, i32> = SlotMap::new();
assert_eq!(map.len(), 0);
assert_eq!(map.slots.len(), 0);
let key = map.insert(56);
assert_eq!(*map.get(key).unwrap(), 56);
*map.get_mut(key).unwrap() = 23;
assert_eq!(*map.get(key).unwrap(), 23);
assert_eq!(map.slots.len(), 1);
let key2 = map.insert(11);
assert_eq!(*map.get(key).unwrap(), 23);
assert_eq!(*map.get(key2).unwrap(), 11);
assert_eq!(map.slots.len(), 2);
assert_eq!(map.remove(key), Some(23));
assert!(map.get(key).is_none());
assert_eq!(*map.get(key2).unwrap(), 11);
assert_eq!(map.slots.len(), 2);
assert_eq!(map.remove(key), None);
assert!(map.get(key).is_none());
assert_eq!(*map.get(key2).unwrap(), 11);
assert_eq!(map.slots.len(), 2);
for (_, v) in map.iter_mut() { *v += 1; }
assert!(map.get(key).is_none());
assert_eq!(*map.get(key2).unwrap(), 12);
assert_eq!(map.slots.len(), 2);
assert_eq!(map.remove(key2), Some(12));
assert!(map.get(key).is_none());
assert!(map.get(key2).is_none());
assert_eq!(map.slots.len(), 2);
assert_eq!(map.remove(key2), None);
assert!(map.get(key).is_none());
assert!(map.get(key2).is_none());
assert_eq!(map.slots.len(), 2);
for _ in 0..4 {
let mut keys = vec![];
for i in 0..128 {
keys.push((map.insert(i), i));
}
assert_eq!(map.slots.len(), 128);
keys[20..].reverse();
keys[..100].reverse();
keys[40..70].reverse();
for i in 0..keys.len() {
assert_eq!(map.len(), 128 - i);
for j in 0..i {
assert!(map.get(keys[j].0).is_none());
assert!(map.get_mut(keys[j].0).is_none());
assert!(map.remove(keys[j].0).is_none());
}
assert_eq!(*map.get(keys[i].0).unwrap(), keys[i].1);
assert_eq!(*map.get_mut(keys[i].0).unwrap(), keys[i].1);
assert_eq!(map.remove(keys[i].0), Some(keys[i].1));
assert_eq!(map.remove(keys[i].0), None);
assert_eq!(map.remove(keys[i].0), None);
assert!(map.get(keys[i].0).is_none());
assert!(map.get_mut(keys[i].0).is_none());
for j in i+1..keys.len() {
assert_eq!(*map.get(keys[j].0).unwrap(), keys[j].1);
assert_eq!(*map.get_mut(keys[j].0).unwrap(), keys[j].1);
}
assert_eq!(map.len(), 128 - i - 1);
let mut cache = BTreeSet::default();
let cpy = map.clone();
cache.clear();
for (key, val) in map.iter() {
assert!(cache.insert(*val));
assert_eq!(map.get(key).unwrap(), val);
assert_eq!(cpy.get(key).unwrap(), val);
}
cache.clear();
for (key, val) in map.iter_mut() {
assert!(cache.insert(*val));
assert_eq!(cpy.get(key).unwrap(), val);
}
cache.clear();
for (key, val) in map.clone() {
assert!(cache.insert(val));
assert_eq!(*map.get(key).unwrap(), val);
assert_eq!(*cpy.get(key).unwrap(), val);
}
}
}
assert_eq!(map.slots.len(), 128);
assert_eq!(map.len(), 0);
assert!(map.is_empty());
let mut keys = vec![];
for i in 0..32 {
keys.push((i, map.insert(i)));
}
keys[..25].reverse();
keys[7..].reverse();
keys[11..19].reverse();
assert_eq!(map.slots.len(), 128);
assert_eq!(map.len(), 32);
assert!(!map.is_empty());
for (i, key) in keys.iter().copied() {
assert_eq!(*map.get(key).unwrap(), i);
assert_eq!(*map.get_mut(key).unwrap(), i);
}
map.clear();
assert_eq!(map.slots.len(), 128);
assert_eq!(map.len(), 0);
assert!(map.is_empty());
for (_, key) in keys.iter().copied() {
assert_eq!(map.get(key).copied(), None);
assert_eq!(map.get_mut(key).copied(), None);
}
}