use std::marker::PhantomData;
#[derive(Clone)]
enum Slot<T> {
Occupied(Occupied<T>),
Free(Free),
}
#[derive(Clone)]
struct Free {
next: usize,
version: u32,
}
#[derive(Clone)]
struct Occupied<T> {
value: T,
version: u32,
}
pub trait SlabId {
fn index(&self) -> usize;
fn version(&self) -> u32;
fn from_raw(index: usize, version: u32) -> Self;
}
#[derive(Clone)]
pub struct Slab<ID, V> {
raw: Vec<Slot<V>>,
len: usize,
free_index: usize,
_phantom: PhantomData<ID>,
}
impl<Id, V> Default for Slab<Id, V>
where
Id: SlabId,
{
#[inline]
fn default() -> Self {
Self {
raw: Vec::default(),
len: 0,
free_index: 0,
_phantom: PhantomData,
}
}
}
impl<ID, V> Slab<ID, V>
where
ID: SlabId + Copy + Clone,
{
pub fn with_capacity(capacity: usize) -> Self {
Self {
raw: Vec::with_capacity(capacity),
len: 0,
free_index: 0,
_phantom: PhantomData,
}
}
#[inline]
pub fn insert(&mut self, value: V) -> ID {
let (index, version) = if self.free_index >= self.raw.len() {
debug_assert_eq!(self.free_index, self.raw.len());
let index = self.free_index;
let version = 1;
self.raw.push(Slot::Occupied(Occupied { value, version }));
self.free_index = self.raw.len();
(index, version)
} else {
let index = self.free_index;
let slot = &mut self.raw[self.free_index];
let version = match slot {
Slot::Occupied(_) => unreachable!(),
Slot::Free(free) => {
self.free_index = free.next;
free.version
}
};
*slot = Slot::Occupied(Occupied { value, version });
(index, version)
};
self.len += 1;
ID::from_raw(index, version)
}
#[allow(clippy::needless_pass_by_value)] #[inline]
pub fn remove(&mut self, id: ID) -> Option<V> {
let index = id.index();
let slot = self.raw.get_mut(index)?;
let version = match slot {
Slot::Occupied(slot) => {
if slot.version != id.version() {
return None;
}
slot.version
}
Slot::Free(_) => return None,
};
let slot = std::mem::replace(
slot,
Slot::Free(Free {
next: self.free_index,
version: version + 1,
}),
);
self.free_index = index;
self.len -= 1;
match slot {
Slot::Occupied(occupied) => Some(occupied.value),
Slot::Free(_) => unreachable!(),
}
}
#[allow(clippy::needless_pass_by_value)] #[inline]
pub fn get(&self, id: ID) -> Option<&V> {
let index = id.index();
let slot = self.raw.get(index)?;
match slot {
Slot::Occupied(slot) => {
if slot.version == id.version() {
Some(&slot.value)
} else {
None
}
}
Slot::Free(_) => None,
}
}
#[allow(clippy::needless_pass_by_value)] #[inline]
pub fn get_mut(&mut self, id: ID) -> Option<&mut V> {
let index = id.index();
let slot = self.raw.get_mut(index)?;
match slot {
Slot::Occupied(slot) => {
if slot.version == id.version() {
Some(&mut slot.value)
} else {
None
}
}
Slot::Free(_) => None,
}
}
#[inline]
pub fn contains_key(&self, id: ID) -> bool {
let index = id.index();
let slot = self.raw.get(index);
match slot {
Some(Slot::Occupied(slot)) => slot.version == id.version(),
_ => false,
}
}
#[inline]
pub fn len(&self) -> usize {
self.len
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
pub fn iter(&self) -> Iter<'_, ID, V> {
Iter {
index: 0,
inner: self.raw.iter(),
_phantom: PhantomData,
}
}
#[inline]
pub fn iter_mut(&mut self) -> IterMut<'_, ID, V> {
IterMut {
index: 0,
inner: self.raw.iter_mut(),
_phantom: PhantomData,
}
}
#[inline]
pub fn capacity(&self) -> usize {
self.raw.capacity()
}
#[inline]
pub fn clear(&mut self) {
self.free_index = 0;
self.len = 0;
for (i, slot) in &mut self.raw.iter_mut().enumerate() {
match slot {
Slot::Occupied(v) => {
let version = v.version;
*slot = Slot::Free(Free {
next: i + 1,
version: version + 1,
});
}
Slot::Free(free_v) => {
free_v.next = i + 1;
}
}
}
}
}
pub struct Iter<'a, ID, V> {
index: usize,
inner: std::slice::Iter<'a, Slot<V>>,
_phantom: PhantomData<ID>,
}
impl<'a, ID, V> Iterator for Iter<'a, ID, V>
where
ID: SlabId,
{
type Item = (ID, &'a V);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
loop {
match self.inner.next()? {
Slot::Occupied(slot) => {
let id = ID::from_raw(self.index, slot.version);
self.index += 1;
return Some((id, &slot.value));
}
Slot::Free(_) => {
self.index += 1;
continue;
}
}
}
}
}
pub struct IterMut<'a, ID, V> {
index: usize,
inner: std::slice::IterMut<'a, Slot<V>>,
_phantom: PhantomData<ID>,
}
impl<'a, ID, V> Iterator for IterMut<'a, ID, V>
where
ID: SlabId,
{
type Item = (ID, &'a mut V);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
loop {
match self.inner.next() {
Some(Slot::Occupied(slot)) => {
let id = ID::from_raw(self.index, slot.version);
self.index += 1;
return Some((id, &mut slot.value));
}
Some(Slot::Free(_)) => {
self.index += 1;
continue;
}
None => return None,
}
}
}
}
impl<ID, V> std::ops::Index<ID> for Slab<ID, V>
where
ID: SlabId + Copy + Clone,
{
type Output = V;
#[inline]
fn index(&self, index: ID) -> &Self::Output {
self.get(index).expect("index not found")
}
}
impl<ID, V> std::ops::IndexMut<ID> for Slab<ID, V>
where
ID: SlabId + Copy + Clone,
{
#[inline]
fn index_mut(&mut self, index: ID) -> &mut V {
self.get_mut(index).expect("index not found")
}
}
impl SlabId for u64 {
#[inline]
fn index(&self) -> usize {
*self as u32 as usize
}
#[inline]
fn version(&self) -> u32 {
(*self >> 32) as u32
}
#[inline]
fn from_raw(index: usize, version: u32) -> Self {
let index = u64::from(u32::try_from(index).expect("index out of range"));
let version = u64::from(version);
version << 32 | index
}
}
#[macro_export]
macro_rules! new_type_id {
(
$(#[$outer:meta])*
$name:ident
) => {
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
#[repr(C)]
$(#[$outer])*
pub struct $name(u64);
impl $crate::SlabId for $name {
#[inline]
fn index(&self) -> usize {
self.0.index()
}
#[inline]
fn version(&self) -> u32 {
self.0.version()
}
#[inline]
fn from_raw(index: usize, version: u32) -> Self {
Self(u64::from_raw(index, version))
}
}
impl std::fmt::Display for $name {
#[inline]
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
self.0.fmt(f)
}
}
impl Default for $name {
#[inline]
fn default() -> Self {
$crate::SlabId::from_raw(Self::EMPTY_INDEX, 1)
}
}
impl $name {
const EMPTY_INDEX: usize = u32::MAX as usize;
#[allow(unused)]
#[inline]
pub fn invalid() -> Self {
Self::default()
}
#[allow(unused)]
#[inline]
pub fn into_u64(self) -> u64 {
self.0
}
#[allow(unused)]
#[inline]
pub fn from_u64(id: u64) -> Self {
Self(id)
}
#[allow(unused)]
#[inline]
pub fn is_valid(self) -> bool {
self != Self::default()
}
#[allow(unused)]
#[inline]
pub fn is_null(self) -> bool {
$crate::SlabId::index(&self) == Self::EMPTY_INDEX
}
}
};
}
#[cfg(test)]
mod tests {
use wasm_bindgen_test::*;
use super::*;
wasm_bindgen_test_configure!(run_in_browser);
#[test]
#[wasm_bindgen_test]
fn test_slab() {
let mut slab = Slab::<u64, i32>::default();
let id0 = slab.insert(5);
assert_eq!(slab.len(), 1);
assert_eq!(id0.index(), 0);
assert_eq!(slab.get(id0).unwrap(), &5);
let id1 = slab.insert(6);
assert_eq!(slab.len(), 2);
assert_eq!(id1.index(), 1);
assert_eq!(slab.get_mut(id0).unwrap(), &5);
assert_eq!(slab.remove(id0).unwrap(), 5);
assert_eq!(slab.len(), 1);
assert!(slab.remove(id0).is_none());
assert!(slab.get(id0).is_none());
assert!(slab.get_mut(id0).is_none());
let id2 = slab.insert(7);
assert_eq!(id2.index(), 0);
assert_eq!(id2.version(), 2);
assert!(slab.get(id0).is_none());
assert!(slab.get_mut(id0).is_none());
assert!(slab.remove(id0).is_none());
assert_eq!(slab.remove(id2).unwrap(), 7);
assert_eq!(slab.remove(id1).unwrap(), 6);
assert!(slab.is_empty());
let invalid_id = u64::from_raw(100, 1);
assert!(slab.get(invalid_id).is_none());
assert!(slab.get_mut(invalid_id).is_none());
}
#[test]
#[wasm_bindgen_test]
fn test_slab_iter() {
let mut slab = Slab::<u64, i32>::default();
let mut id = slab.insert(8);
slab.remove(id).unwrap();
id = slab.insert(9);
slab.insert(10);
slab.insert(11);
slab.remove(id).unwrap();
slab.insert(12);
let mut iter = slab.iter();
{
let (id, value) = iter.next().unwrap();
assert_eq!(id.index(), 0);
assert_eq!(id.version(), 3);
assert_eq!(id, u64::from_raw(0, 3));
assert_eq!(value, &12);
let value = slab.get(id).unwrap();
assert_eq!(value, &12);
let value = slab[id];
assert_eq!(value, 12);
}
{
let (id, value) = iter.next().unwrap();
assert_eq!(id.index(), 1);
assert_eq!(id.version(), 1);
assert_eq!(id, u64::from_raw(1, 1));
assert_eq!(value, &10);
let value = slab.get(id).unwrap();
assert_eq!(value, &10);
let value = slab[id];
assert_eq!(value, 10);
}
{
let (id, value) = iter.next().unwrap();
assert_eq!(id.index(), 2);
assert_eq!(id.version(), 1);
assert_eq!(id, u64::from_raw(2, 1));
assert_eq!(value, &11);
let value = slab.get(id).unwrap();
assert_eq!(value, &11);
let value = slab[id];
assert_eq!(value, 11);
}
assert!(iter.next().is_none());
let mut iter = slab.iter_mut();
{
let (id, value) = iter.next().unwrap();
assert_eq!(id.index(), 0);
assert_eq!(id.version(), 3);
assert_eq!(id, u64::from_raw(0, 3));
assert_eq!(value, &mut 12);
}
{
let (id, value) = iter.next().unwrap();
assert_eq!(id.index(), 1);
assert_eq!(id.version(), 1);
assert_eq!(id, u64::from_raw(1, 1));
assert_eq!(value, &mut 10);
}
{
let (id, value) = iter.next().unwrap();
assert_eq!(id.index(), 2);
assert_eq!(id.version(), 1);
assert_eq!(id, u64::from_raw(2, 1));
assert_eq!(value, &mut 11);
}
assert!(iter.next().is_none());
assert_eq!(slab.get(u64::from_raw(0, 3)), Some(&12));
assert_eq!(slab.get(u64::from_raw(1, 1)), Some(&10));
assert!(slab.contains_key(u64::from_raw(0, 3)));
assert!(slab.contains_key(u64::from_raw(1, 1)));
assert!(!slab.contains_key(u64::from_raw(0, 1)));
assert!(!slab.contains_key(u64::from_raw(1, 3)));
assert!(slab.contains_key(u64::from_raw(2, 1)));
assert!(!slab.contains_key(u64::from_raw(3, 1)));
{
let slab = &mut slab;
let value = slab[u64::from_raw(0, 3)];
assert_eq!(value, 12);
}
slab.remove(u64::from_raw(3, 1));
}
#[test]
#[wasm_bindgen_test]
#[allow(clippy::cast_sign_loss)]
fn test_slab_capacity() {
let mut slab = Slab::<u64, i32>::with_capacity(10);
assert_eq!(slab.capacity(), 10);
assert_eq!(slab.len(), 0);
for i in 0..10 {
slab.insert(i);
assert_eq!(slab.capacity(), 10);
assert_eq!(slab.len(), (i + 1) as usize);
}
slab.insert(10);
assert_eq!(slab.capacity(), 20);
assert_eq!(slab.len(), 11);
}
#[test]
#[wasm_bindgen_test]
fn test_slab_remove() {
let mut slab = Slab::<u64, i32>::with_capacity(10);
assert_eq!(slab.capacity(), 10);
assert_eq!(slab.len(), 0);
let id = slab.insert(1);
assert_eq!(slab.capacity(), 10);
assert_eq!(slab.len(), 1);
slab.remove(id);
assert_eq!(slab.capacity(), 10);
assert_eq!(slab.len(), 0);
slab.remove(id);
assert_eq!(slab.capacity(), 10);
assert_eq!(slab.len(), 0);
}
new_type_id!(TestId);
#[test]
#[wasm_bindgen_test]
fn test_type_id() {
let id = TestId::default();
assert_eq!(id.index(), u32::MAX as usize);
assert_eq!(id.version(), 1);
assert_eq!(id, TestId::default());
assert_eq!(id.to_string(), "8589934591");
let id = TestId::from_raw(1, 2);
assert_eq!(id.index(), 1);
assert_eq!(id.version(), 2);
assert_eq!(id, TestId::from_raw(1, 2));
assert_eq!(id.to_string(), "8589934593");
let id = TestId::from_raw(1, 2);
assert_eq!(id.into_u64(), 8589934593_u64);
assert_eq!(TestId::from_u64(8589934593_u64), TestId::from_raw(1, 2));
assert!(id.is_valid());
assert!(!TestId::default().is_valid());
assert!(TestId::default().is_null());
}
#[test]
#[wasm_bindgen_test]
fn test_slab_clear() {
let mut slab = Slab::<u64, i32>::with_capacity(10);
assert_eq!(slab.capacity(), 10);
assert_eq!(slab.len(), 0);
for i in 0..10 {
slab.insert(i);
}
assert_eq!(slab.capacity(), 10);
assert_eq!(slab.len(), 10);
slab.clear();
assert_eq!(slab.capacity(), 10);
assert_eq!(slab.len(), 0);
for (i, v) in slab.raw.iter().enumerate() {
match v {
Slot::Occupied(_) => panic!("Occupied slot found"),
Slot::Free(free_v) => {
assert_eq!(free_v.next, i + 1)
}
}
}
for i in 0..5 {
slab.insert(i);
}
slab.clear();
for (i, v) in slab.raw.iter().enumerate() {
match v {
Slot::Occupied(_) => panic!("Occupied slot found"),
Slot::Free(free_v) => {
assert_eq!(free_v.next, i + 1)
}
}
}
}
}