netsblox_vm/
slotmap.rs

1//! A vector that allows for generational indexing.
2//! 
3//! This is essentially a rewrite of a subset of the `slotmap` crate since we can't implement [`Collect`] on foreign types.
4
5use alloc::vec::{self, Vec};
6use core::marker::PhantomData;
7use core::{slice, iter};
8
9use crate::gc::*;
10
11/// A key that can be used with a [`SlotMap`].
12/// 
13/// Instead of implementing this trait manually, it is recommended to use the [`new_key`](crate::new_key) macro.
14pub trait Key: Copy + Eq + Ord + 'static {
15    fn new(slot: usize, generation: usize) -> Self;
16    fn get_slot(&self) -> usize;
17    fn get_generation(&self) -> usize;
18}
19
20/// Defines a new key type for use in [`SlotMap`].
21#[macro_export]
22macro_rules! new_key {
23    ($($(#[doc = $doc:expr])? $vis:vis struct $name:ident;)*) => {$(
24        #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
25        $(#[doc = $doc])?
26        $vis struct $name(usize, usize);
27        impl $crate::slotmap::Key for $name {
28            fn new(slot: usize, generation: usize) -> Self { Self(slot, generation) }
29            fn get_slot(&self) -> usize { self.0 }
30            fn get_generation(&self) -> usize { self.1 }
31        }
32    )*}
33}
34
35#[derive(Clone, Collect)]
36#[collect(no_drop)]
37struct Slot<T> {
38                               value: Option<T>,
39    #[collect(require_static)] generation: usize,
40}
41
42/// A dense, resizable array that supports generational indexing.
43/// 
44/// You can use the [`new_key`](crate::new_key) macro to create a new key type to use.
45/// It is recommended to use different key types for different collections to avoid accidentally using a key from a different map.
46#[derive(Clone, Collect)]
47#[collect(no_drop)]
48pub struct SlotMap<K: Key, T> {
49                               slots: Vec<Slot<T>>,
50    #[collect(require_static)] empty_slots: Vec<usize>,
51    #[collect(require_static)] num_values: usize,
52    #[collect(require_static)] _key: PhantomData<K>
53}
54impl<K: Key, T> Default for SlotMap<K, T> {
55    fn default() -> Self {
56        Self::new()
57    }
58}
59impl<K: Key, T> SlotMap<K, T> {
60    /// Creates a new empty container.
61    pub fn new() -> Self {
62        SlotMap {
63            slots: vec![],
64            empty_slots: vec![],
65            num_values: 0,
66            _key: PhantomData,
67        }
68    }
69    #[cfg(test)]
70    fn invariant(&self) -> bool {
71        self.num_values == self.slots.iter().filter(|x| x.value.is_some()).count()
72        &&
73        self.num_values + self.empty_slots.len() == self.slots.len()
74    }
75    /// Adds a new value to the map and returns a new key that references it.
76    pub fn insert(&mut self, value: T) -> K {
77        #[cfg(test)] assert!(self.invariant());
78
79        self.num_values += 1;
80        let key = match self.empty_slots.pop() {
81            Some(slot) => {
82                debug_assert!(self.slots[slot].value.is_none());
83                self.slots[slot].value = Some(value);
84                K::new(slot, self.slots[slot].generation)
85            }
86            None => {
87                debug_assert!(self.slots.iter().all(|x| x.value.is_some()));
88                let slot = self.slots.len();
89                self.slots.push(Slot { value: Some(value), generation: 0 });
90                K::new(slot, 0)
91            }
92        };
93
94        #[cfg(test)] assert!(self.invariant());
95        #[allow(clippy::let_and_return)] key
96    }
97    /// Removes a value from the map and returns it (if it existed).
98    /// It is guaranteed that all future accesses with the removed key will return [`None`].
99    pub fn remove(&mut self, key: K) -> Option<T> {
100        #[cfg(test)] assert!(self.invariant());
101
102        let slot = self.slots.get_mut(key.get_slot())?;
103        let res = if slot.generation == key.get_generation() { slot.value.take() } else { None };
104        if res.is_some() {
105            slot.generation += 1;
106            self.num_values -= 1;
107            self.empty_slots.push(key.get_slot());
108        }
109
110        #[cfg(test)] assert!(self.invariant());
111        res
112    }
113    /// Removes all values from the slotmap.
114    pub fn clear(&mut self) {
115        #[cfg(test)] assert!(self.invariant());
116
117        for (i, slot) in self.slots.iter_mut().enumerate() {
118            if slot.value.take().is_some() {
119                slot.generation += 1;
120                self.empty_slots.push(i);
121            }
122        }
123        self.num_values = 0;
124
125        #[cfg(test)] assert!(self.invariant());
126    }
127    /// Retains only the values for which the predicate returns true.
128    /// For all retained values, any existing keys are still valid after this operation.
129    pub fn retain_mut<F: FnMut(K, &mut T) -> bool>(&mut self, mut f: F) {
130        #[cfg(test)] assert!(self.invariant());
131
132        for (i, slot) in self.slots.iter_mut().enumerate() {
133            if let Some(value) = &mut slot.value {
134                if !f(K::new(i, slot.generation), value) {
135                    slot.value = None;
136                    slot.generation += 1;
137                    self.num_values -= 1;
138                    self.empty_slots.push(i);
139                }
140            }
141        }
142
143        #[cfg(test)] assert!(self.invariant());
144    }
145    /// Get a reference to a value in the map.
146    pub fn get(&self, key: K) -> Option<&T> {
147        let slot = self.slots.get(key.get_slot())?;
148        if slot.generation == key.get_generation() { slot.value.as_ref() } else { None }
149    }
150    /// Get a mutable reference to a value in the map.
151    pub fn get_mut(&mut self, key: K) -> Option<&mut T> {
152        let slot = self.slots.get_mut(key.get_slot())?;
153        if slot.generation == key.get_generation() { slot.value.as_mut() } else { None }
154    }
155    /// Get the number of values stored in the map.
156    pub fn len(&self) -> usize {
157        self.num_values
158    }
159    /// Checks if the map is currently empty.
160    pub fn is_empty(&self) -> bool {
161        self.len() == 0
162    }
163    /// Iterates over the keys and values currently stored in the map.
164    pub fn iter(&self) -> Iter<K, T> {
165        Iter(self.slots.iter().enumerate(), PhantomData)
166    }
167    /// Mutably iterates over the keys and values currently stored in the map.
168    pub fn iter_mut(&mut self) -> IterMut<K, T> {
169        IterMut(self.slots.iter_mut().enumerate(), PhantomData)
170    }
171}
172impl<K: Key, T> IntoIterator for SlotMap<K, T> {
173    type IntoIter = IntoIter<K, T>;
174    type Item = (K, T);
175    fn into_iter(self) -> IntoIter<K, T> {
176        IntoIter(self.slots.into_iter().enumerate(), PhantomData)
177    }
178}
179
180pub struct IntoIter<K: Key, T>(iter::Enumerate<vec::IntoIter<Slot<T>>>, PhantomData<K>);
181pub struct Iter<'a, K: Key, T>(iter::Enumerate<slice::Iter<'a, Slot<T>>>, PhantomData<K>);
182pub struct IterMut<'a, K: Key, T>(iter::Enumerate<slice::IterMut<'a, Slot<T>>>, PhantomData<K>);
183
184impl<K: Key, T> Iterator for IntoIter<K, T> {
185    type Item = (K, T);
186    fn next(&mut self) -> Option<Self::Item> {
187        loop {
188            let (i, slot) = self.0.next()?;
189            if let Some(x) = slot.value { return Some((K::new(i, slot.generation), x)) }
190        }
191    }
192}
193impl<'a, K: Key, T> Iterator for Iter<'a, K, T> {
194    type Item = (K, &'a T);
195    fn next(&mut self) -> Option<Self::Item> {
196        loop {
197            let (i, slot) = self.0.next()?;
198            if let Some(x) = slot.value.as_ref() { return Some((K::new(i, slot.generation), x)) }
199        }
200    }
201}
202impl<'a, K: Key, T> Iterator for IterMut<'a, K, T> {
203    type Item = (K, &'a mut T);
204    fn next(&mut self) -> Option<Self::Item> {
205        loop {
206            let (i, slot) = self.0.next()?;
207            if let Some(x) = slot.value.as_mut() { return Some((K::new(i, slot.generation), x)) }
208        }
209    }
210}
211
212#[test]
213fn test_slotmap() {
214    use alloc::collections::BTreeSet;
215    new_key! {
216        struct TestKey;
217    }
218    let mut map: SlotMap<TestKey, i32> = SlotMap::new();
219    assert_eq!(map.len(), 0);
220    assert_eq!(map.slots.len(), 0);
221
222    let key = map.insert(56);
223    assert_eq!(*map.get(key).unwrap(), 56);
224    *map.get_mut(key).unwrap() = 23;
225    assert_eq!(*map.get(key).unwrap(), 23);
226    assert_eq!(map.slots.len(), 1);
227
228    let key2 = map.insert(11);
229    assert_eq!(*map.get(key).unwrap(), 23);
230    assert_eq!(*map.get(key2).unwrap(), 11);
231    assert_eq!(map.slots.len(), 2);
232
233    assert_eq!(map.remove(key), Some(23));
234    assert!(map.get(key).is_none());
235    assert_eq!(*map.get(key2).unwrap(), 11);
236    assert_eq!(map.slots.len(), 2);
237
238    assert_eq!(map.remove(key), None);
239    assert!(map.get(key).is_none());
240    assert_eq!(*map.get(key2).unwrap(), 11);
241    assert_eq!(map.slots.len(), 2);
242
243    for (_, v) in map.iter_mut() { *v += 1; }
244    assert!(map.get(key).is_none());
245    assert_eq!(*map.get(key2).unwrap(), 12);
246    assert_eq!(map.slots.len(), 2);
247
248    assert_eq!(map.remove(key2), Some(12));
249    assert!(map.get(key).is_none());
250    assert!(map.get(key2).is_none());
251    assert_eq!(map.slots.len(), 2);
252
253    assert_eq!(map.remove(key2), None);
254    assert!(map.get(key).is_none());
255    assert!(map.get(key2).is_none());
256    assert_eq!(map.slots.len(), 2);
257
258    for _ in 0..4 {
259        let mut keys = vec![];
260        for i in 0..128 {
261            keys.push((map.insert(i), i));
262        }
263        assert_eq!(map.slots.len(), 128);
264        keys[20..].reverse();
265        keys[..100].reverse();
266        keys[40..70].reverse();
267
268        for i in 0..keys.len() {
269            assert_eq!(map.len(), 128 - i);
270            for j in 0..i {
271                assert!(map.get(keys[j].0).is_none());
272                assert!(map.get_mut(keys[j].0).is_none());
273                assert!(map.remove(keys[j].0).is_none());
274            }
275            assert_eq!(*map.get(keys[i].0).unwrap(), keys[i].1);
276            assert_eq!(*map.get_mut(keys[i].0).unwrap(), keys[i].1);
277            assert_eq!(map.remove(keys[i].0), Some(keys[i].1));
278            assert_eq!(map.remove(keys[i].0), None);
279            assert_eq!(map.remove(keys[i].0), None);
280            assert!(map.get(keys[i].0).is_none());
281            assert!(map.get_mut(keys[i].0).is_none());
282            for j in i+1..keys.len() {
283                assert_eq!(*map.get(keys[j].0).unwrap(), keys[j].1);
284                assert_eq!(*map.get_mut(keys[j].0).unwrap(), keys[j].1);
285            }
286            assert_eq!(map.len(), 128 - i - 1);
287
288            let mut cache = BTreeSet::default();
289            let cpy = map.clone();
290
291            cache.clear();
292            for (key, val) in map.iter() {
293                assert!(cache.insert(*val));
294                assert_eq!(map.get(key).unwrap(), val);
295                assert_eq!(cpy.get(key).unwrap(), val);
296            }
297            cache.clear();
298            for (key, val) in map.iter_mut() {
299                assert!(cache.insert(*val));
300                assert_eq!(cpy.get(key).unwrap(), val);
301            }
302            cache.clear();
303            for (key, val) in map.clone() {
304                assert!(cache.insert(val));
305                assert_eq!(*map.get(key).unwrap(), val);
306                assert_eq!(*cpy.get(key).unwrap(), val);
307            }
308        }
309    }
310
311    assert_eq!(map.slots.len(), 128);
312    assert_eq!(map.len(), 0);
313    assert!(map.is_empty());
314
315    let mut keys = vec![];
316    for i in 0..32 {
317        keys.push((i, map.insert(i)));
318    }
319    keys[..25].reverse();
320    keys[7..].reverse();
321    keys[11..19].reverse();
322
323    assert_eq!(map.slots.len(), 128);
324    assert_eq!(map.len(), 32);
325    assert!(!map.is_empty());
326    for (i, key) in keys.iter().copied() {
327        assert_eq!(*map.get(key).unwrap(), i);
328        assert_eq!(*map.get_mut(key).unwrap(), i);
329    }
330
331    map.clear();
332    assert_eq!(map.slots.len(), 128);
333    assert_eq!(map.len(), 0);
334    assert!(map.is_empty());
335    for (_, key) in keys.iter().copied() {
336        assert_eq!(map.get(key).copied(), None);
337        assert_eq!(map.get_mut(key).copied(), None);
338    }
339
340    map.clear();
341    let k1 = map.insert(12);
342    assert_eq!(map.remove(k1), Some(12));
343    assert_eq!(map.remove(k1), None);
344    assert_eq!(map.remove(k1), None);
345    let k2 = map.insert(13);
346    assert_eq!(k1.0, k2.0);
347    assert_eq!(k1.1 + 1, k2.1);
348    assert_eq!(map.remove(k1), None);
349    assert_eq!(map.remove(k1), None);
350    assert_eq!(map.remove(k2), Some(13));
351    assert_eq!(map.remove(k2), None);
352    assert_eq!(map.remove(k2), None);
353
354    map.clear();
355    let kv1 = map.insert(54);
356    let kv2 = map.insert(-4);
357    let kv3 = map.insert(51);
358    let kv4 = map.insert(-53);
359    let kv5 = map.insert(52);
360    let kv6 = map.insert(12);
361    let kv7 = map.insert(2);
362    assert_eq!(map.get(kv1).copied(), Some(54));
363    assert_eq!(map.get(kv2).copied(), Some(-4));
364    assert_eq!(map.get(kv3).copied(), Some(51));
365    assert_eq!(map.get(kv4).copied(), Some(-53));
366    assert_eq!(map.get(kv5).copied(), Some(52));
367    assert_eq!(map.get(kv6).copied(), Some(12));
368    assert_eq!(map.get(kv7).copied(), Some(2));
369    map.retain_mut(|k, v| k == kv6 || *v % 2 != 0);
370    assert_eq!(map.get(kv1).copied(), None);
371    assert_eq!(map.get(kv2).copied(), None);
372    assert_eq!(map.get(kv3).copied(), Some(51));
373    assert_eq!(map.get(kv4).copied(), Some(-53));
374    assert_eq!(map.get(kv5).copied(), None);
375    assert_eq!(map.get(kv6).copied(), Some(12));
376    assert_eq!(map.get(kv7).copied(), None);
377}
netsblox_vm/slotmap.rs

netsblox_vm/
slotmap.rs
