#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct SlabKey(u32);
impl SlabKey {
#[inline]
pub fn from_raw(idx: u32) -> Self {
SlabKey(idx)
}
#[inline]
pub fn into_raw(self) -> u32 {
self.0
}
}
impl std::fmt::Display for SlabKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "SlabKey({})", self.0)
}
}
enum Entry<T> {
Occupied {
value: T,
generation: u32,
},
Vacant {
next_free: u32,
freed_gen: u32,
},
}
pub struct Slab<T> {
entries: Vec<Entry<T>>,
free_head: u32,
len: usize,
}
impl<T> Slab<T> {
pub fn new() -> Self {
Slab {
entries: Vec::new(),
free_head: u32::MAX,
len: 0,
}
}
pub fn with_capacity(cap: usize) -> Self {
Slab {
entries: Vec::with_capacity(cap),
free_head: u32::MAX,
len: 0,
}
}
pub fn insert(&mut self, value: T) -> SlabKey {
if self.free_head != u32::MAX {
let idx = self.free_head as usize;
let (next_free, new_gen) = match &self.entries[idx] {
Entry::Vacant { next_free, freed_gen } => (*next_free, freed_gen.wrapping_add(1)),
Entry::Occupied { .. } => unreachable!("free list must point to vacant entries"),
};
self.free_head = next_free;
self.entries[idx] = Entry::Occupied {
value,
generation: new_gen,
};
self.len += 1;
SlabKey(idx as u32)
} else {
let idx = self.entries.len();
assert!(idx < u32::MAX as usize, "slab capacity overflow");
self.entries.push(Entry::Occupied {
value,
generation: 0,
});
self.len += 1;
SlabKey(idx as u32)
}
}
pub fn remove(&mut self, key: SlabKey) -> T {
let idx = key.0 as usize;
assert!(idx < self.entries.len(), "SlabKey out of range");
let freed_gen = match &self.entries[idx] {
Entry::Occupied { generation, .. } => *generation,
Entry::Vacant { .. } => panic!("attempted to remove already-vacant entry {key}"),
};
let old = std::mem::replace(
&mut self.entries[idx],
Entry::Vacant {
next_free: self.free_head,
freed_gen,
},
);
self.free_head = idx as u32;
self.len -= 1;
match old {
Entry::Occupied { value, .. } => value,
Entry::Vacant { .. } => unreachable!(),
}
}
pub fn get(&self, key: SlabKey) -> Option<&T> {
let idx = key.0 as usize;
match self.entries.get(idx)? {
Entry::Occupied { value, .. } => Some(value),
Entry::Vacant { .. } => None,
}
}
pub fn get_mut(&mut self, key: SlabKey) -> Option<&mut T> {
let idx = key.0 as usize;
match self.entries.get_mut(idx)? {
Entry::Occupied { value, .. } => Some(value),
Entry::Vacant { .. } => None,
}
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
pub fn len(&self) -> usize {
self.len
}
#[inline]
pub fn capacity(&self) -> usize {
self.entries.capacity()
}
pub fn iter(&self) -> SlabIter<'_, T> {
SlabIter {
entries: &self.entries,
index: 0,
}
}
pub fn iter_mut(&mut self) -> SlabIterMut<'_, T> {
SlabIterMut {
entries: &mut self.entries,
index: 0,
}
}
pub fn clear(&mut self) {
self.entries.clear();
self.free_head = u32::MAX;
self.len = 0;
}
pub fn contains(&self, key: SlabKey) -> bool {
self.get(key).is_some()
}
}
impl<T> Default for Slab<T> {
fn default() -> Self {
Slab::new()
}
}
pub struct SlabIter<'a, T> {
entries: &'a [Entry<T>],
index: usize,
}
impl<'a, T> Iterator for SlabIter<'a, T> {
type Item = (SlabKey, &'a T);
fn next(&mut self) -> Option<Self::Item> {
loop {
let idx = self.index;
if idx >= self.entries.len() {
return None;
}
self.index += 1;
if let Entry::Occupied { value, .. } = &self.entries[idx] {
return Some((SlabKey(idx as u32), value));
}
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
(0, Some(self.entries.len() - self.index))
}
}
pub struct SlabIterMut<'a, T> {
entries: &'a mut [Entry<T>],
index: usize,
}
impl<'a, T> Iterator for SlabIterMut<'a, T> {
type Item = (SlabKey, &'a mut T);
fn next(&mut self) -> Option<Self::Item> {
loop {
let idx = self.index;
if idx >= self.entries.len() {
return None;
}
self.index += 1;
let entry = unsafe { &mut *(self.entries.as_mut_ptr().add(idx)) };
if let Entry::Occupied { value, .. } = entry {
return Some((SlabKey(idx as u32), value));
}
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
(0, Some(self.entries.len() - self.index))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn slab_insert_get() {
let mut s: Slab<i32> = Slab::new();
let k1 = s.insert(10);
let k2 = s.insert(20);
let k3 = s.insert(30);
assert_eq!(s.get(k1), Some(&10));
assert_eq!(s.get(k2), Some(&20));
assert_eq!(s.get(k3), Some(&30));
assert_eq!(s.len(), 3);
}
#[test]
fn slab_remove_and_reinsert() {
let mut s: Slab<String> = Slab::new();
let k1 = s.insert("alpha".to_string());
let k2 = s.insert("beta".to_string());
assert_eq!(s.remove(k1), "alpha");
assert_eq!(s.len(), 1);
assert_eq!(s.get(k1), None); let k3 = s.insert("gamma".to_string());
assert_eq!(k3, k1, "slot k1 should have been recycled");
assert_eq!(s.get(k3), Some(&"gamma".to_string()));
assert_eq!(s.get(k2), Some(&"beta".to_string()));
}
#[test]
fn slab_get_mut() {
let mut s: Slab<i32> = Slab::new();
let k = s.insert(42);
*s.get_mut(k).expect("entry should exist") = 99;
assert_eq!(s.get(k), Some(&99));
}
#[test]
fn slab_iter() {
let mut s: Slab<u32> = Slab::new();
let k0 = s.insert(0_u32);
let k1 = s.insert(1_u32);
let k2 = s.insert(2_u32);
s.remove(k1);
let pairs: Vec<_> = s.iter().collect();
assert_eq!(pairs.len(), 2);
assert!(pairs.contains(&(k0, &0)));
assert!(pairs.contains(&(k2, &2)));
}
#[test]
fn slab_iter_mut() {
let mut s: Slab<i32> = Slab::new();
s.insert(1);
s.insert(2);
s.insert(3);
for (_k, v) in s.iter_mut() {
*v *= 10;
}
let vals: Vec<i32> = s.iter().map(|(_, v)| *v).collect();
assert_eq!(vals, vec![10, 20, 30]);
}
#[test]
fn slab_clear() {
let mut s: Slab<i32> = Slab::new();
let _k = s.insert(5);
s.clear();
assert!(s.is_empty());
assert_eq!(s.len(), 0);
}
#[test]
fn slab_contains() {
let mut s: Slab<()> = Slab::new();
let k = s.insert(());
assert!(s.contains(k));
s.remove(k);
assert!(!s.contains(k));
}
#[test]
#[should_panic(expected = "already-vacant")]
fn slab_double_remove_panics() {
let mut s: Slab<i32> = Slab::new();
let k = s.insert(1);
s.remove(k);
s.remove(k); }
#[test]
fn slab_key_raw_round_trip() {
let k = SlabKey::from_raw(42);
assert_eq!(k.into_raw(), 42);
}
#[test]
fn slab_large_sequence() {
const N: usize = 1_000;
let mut s: Slab<usize> = Slab::with_capacity(N);
let keys: Vec<SlabKey> = (0..N).map(|i| s.insert(i)).collect();
assert_eq!(s.len(), N);
for i in (0..N).step_by(2) {
let v = s.remove(keys[i]);
assert_eq!(v, i);
}
assert_eq!(s.len(), N / 2);
for i in (1..N).step_by(2) {
assert_eq!(s.get(keys[i]), Some(&i));
}
}
}