wasm/core/
slotmap.rs

1use core::{marker::PhantomData, mem};
2
3use alloc::vec::Vec;
4
5enum SlotContent<T> {
6    Unoccupied { prev_unoccupied: Option<usize> },
7    Occupied { item: T },
8}
9
10struct Slot<T> {
11    generation: u64,
12    content: SlotContent<T>,
13}
14
15/// A contigious data structure that never shrinks, but keeps track of lazily deleted elements so that when a new item is inserted, the lazily deleted place is reused.
16/// Insertion, removal, and update are all of O(1) time complexity. Inserted elements can be (mutably) accessed or removed with the key returned when they were inserted.
17/// Note: might make a slot permanently unusable per `u64::MAX` insert calls during runtime.
18pub struct SlotMap<T> {
19    slots: Vec<Slot<T>>,
20    last_unoccupied: Option<usize>,
21}
22
23impl<T> Default for SlotMap<T> {
24    fn default() -> Self {
25        Self {
26            slots: Vec::new(),
27            last_unoccupied: None,
28        }
29    }
30}
31
32pub struct SlotMapKey<T> {
33    index: usize,
34    generation: u64,
35    phantom: PhantomData<T>,
36}
37
38impl<T> SlotMap<T> {
39    pub fn insert(&mut self, item: T) -> SlotMapKey<T> {
40        match self.last_unoccupied {
41            Some(last_unoccupied) => {
42                let slot = &mut self.slots[last_unoccupied];
43                let SlotContent::Unoccupied { prev_unoccupied } = slot.content else {
44                    unreachable!("last unoccupied slot in slotmap must be unoccupied")
45                };
46                self.last_unoccupied = prev_unoccupied;
47
48                slot.content = SlotContent::Occupied { item };
49                SlotMapKey {
50                    index: last_unoccupied,
51                    generation: slot.generation,
52                    phantom: PhantomData,
53                }
54            }
55            None => {
56                let index = self.slots.len();
57                let generation = 0;
58                self.slots.push(Slot {
59                    generation,
60                    content: SlotContent::Occupied { item },
61                });
62                SlotMapKey {
63                    index,
64                    generation,
65                    phantom: PhantomData,
66                }
67            }
68        }
69    }
70
71    #[allow(unused)]
72    pub fn get(&self, key: &SlotMapKey<T>) -> Option<&T> {
73        let slot = self.slots.get(key.index)?;
74        if slot.generation != key.generation {
75            return None;
76        }
77        match &slot.content {
78            SlotContent::Occupied { item } => Some(item),
79            SlotContent::Unoccupied { .. } => None,
80        }
81    }
82
83    pub fn get_mut(&mut self, key: &SlotMapKey<T>) -> Option<&mut T> {
84        let slot = self.slots.get_mut(key.index)?;
85        if slot.generation != key.generation {
86            return None;
87        }
88        match &mut slot.content {
89            SlotContent::Occupied { item } => Some(item),
90            SlotContent::Unoccupied { .. } => None,
91        }
92    }
93
94    pub fn remove(&mut self, key: &SlotMapKey<T>) -> Option<T> {
95        let slot = self.slots.get_mut(key.index)?;
96        if slot.generation != key.generation
97            || matches!(slot.content, SlotContent::Unoccupied { .. })
98        {
99            return None;
100        }
101
102        let new_slot = if let Some(generation) = slot.generation.checked_add(1) {
103            let prev_unoccupied = self.last_unoccupied;
104            self.last_unoccupied = Some(key.index);
105            Slot {
106                generation,
107                content: SlotContent::Unoccupied { prev_unoccupied },
108            }
109        } else {
110            // no self.last_unoccupied update here, permanently make slot unusable if its generation overflows.
111            Slot {
112                generation: 0,
113                content: SlotContent::Unoccupied {
114                    prev_unoccupied: None,
115                },
116            }
117        };
118        let previous_slot = mem::replace(slot, new_slot);
119
120        let SlotContent::Occupied { item } = previous_slot.content else {
121            unreachable!("slot was full")
122        };
123
124        Some(item)
125    }
126}
127
128#[test]
129fn test() {
130    let mut slotmap = SlotMap::<i32>::default();
131    let key = slotmap.insert(5);
132    assert_eq!(slotmap.get(&key), Some(&5));
133    slotmap.remove(&key);
134    assert_eq!(slotmap.get(&key), None);
135    let key2 = slotmap.insert(10);
136    assert_eq!(slotmap.get(&key), None);
137    assert_eq!(slotmap.get(&key2), Some(&10));
138    let n = slotmap.get_mut(&key2).unwrap();
139    *n = 42;
140    assert_eq!(slotmap.get(&key2), Some(&42));
141}
wasm/core/slotmap.rs

wasm/core/
slotmap.rs
