playit_agent_core/utils/
id_slab.rs

1use std::{mem::{ManuallyDrop, MaybeUninit}, u32};
2
3pub struct IdSlab<T> {
4    entries: Vec<Entry<T>>,
5    free_slots: Vec<usize>,
6}
7
8struct Entry<T> {
9    id: u64,
10    value: MaybeUninit<T>,
11}
12
13pub struct IdSlabVacantEntry<'a, T> {
14    slab: &'a mut IdSlab<T>,
15    id: u64,
16    slot: usize,
17}
18
19const EMPTY_BIT: u64 = 1u64 << 63;
20const EMPTY_BIT_NEG: u64 = !EMPTY_BIT;
21const USE_NUM: u64 = (u32::MAX as u64) + 1;
22const SLOT_MASK: u64 = 0x00000000FFFFFFFF;
23
24impl<T> IdSlab<T> {
25    pub fn with_capacity(capacity: usize) -> Self {
26        let mut slab = IdSlab {
27            entries: Vec::with_capacity(capacity),
28            free_slots: Vec::with_capacity(capacity),
29        };
30
31        for pos in 0..capacity {
32            slab.entries.push(Entry {
33                id: EMPTY_BIT | (pos as u64),
34                value: MaybeUninit::uninit(),
35            });
36
37            slab.free_slots.push(capacity - (pos + 1));
38        }
39
40        slab
41    }
42
43    pub fn capacity(&self) -> usize {
44        self.entries.len()
45    }
46
47    pub fn len(&self) -> usize {
48        self.entries.len() - self.free_slots.len()
49    }
50
51    pub fn available(&self) -> usize {
52        self.free_slots.len()
53    }
54
55    pub fn get(&self, id: u64) -> Option<&T> {
56        let slot = self.slot(id)?;
57
58        let entry = &self.entries[slot];
59        if (entry.id & EMPTY_BIT) == EMPTY_BIT {
60            return None;
61        }
62
63        unsafe { Some(entry.value.assume_init_ref()) }
64    }
65
66    pub fn get_mut(&mut self, id: u64) -> Option<&mut T> {
67        let slot = self.slot(id)?;
68
69        let entry = &mut self.entries[slot];
70        if (entry.id & EMPTY_BIT) == EMPTY_BIT {
71            return None;
72        }
73
74        unsafe { Some(entry.value.assume_init_mut()) }
75    }
76
77    pub fn remove(&mut self, id: u64) -> Option<T> {
78        let slot = self.slot(id)?;
79
80        let entry = &mut self.entries[slot];
81        if (entry.id & EMPTY_BIT) == EMPTY_BIT {
82            return None;
83        }
84
85        entry.id = EMPTY_BIT | (entry.id + USE_NUM);
86        assert_eq!(entry.id & EMPTY_BIT, EMPTY_BIT);
87
88        self.free_slots.push(slot);
89
90        Some(unsafe {
91            std::mem::replace(&mut entry.value, MaybeUninit::uninit()).assume_init()
92        })
93    }
94
95    fn slot(&self, id: u64) -> Option<usize> {
96        let slot = (id & SLOT_MASK) as usize;
97        if self.entries.len() <= slot {
98            return None;
99        }
100        Some(slot)
101    }
102
103    pub fn insert(&mut self, value: T) -> Result<u64, T> {
104        let slot = match self.free_slots.pop() {
105            Some(v) => v,
106            None => return Err(value),
107        };
108
109        let entry = &mut self.entries[slot];
110        assert!((entry.id & EMPTY_BIT) == EMPTY_BIT);
111
112        entry.id = EMPTY_BIT_NEG & entry.id;
113        assert!((entry.id & EMPTY_BIT) == 0);
114
115        entry.value.write(value);
116        Ok(entry.id)
117    }
118
119    pub fn vacant_entry(&mut self) -> Option<IdSlabVacantEntry<T>> {
120        let slot = self.free_slots.pop()?;
121        let id = self.entries[slot].id & EMPTY_BIT_NEG;
122
123        Some(IdSlabVacantEntry {
124            slab: self,
125            id,
126            slot,
127        })
128    }
129
130    pub fn iter(&self) -> IdSlabIter<T> {
131        IdSlabIter {
132            slab: self,
133            slot: 0,
134            remaining: self.len(),
135        }
136    }
137
138    pub fn iter_mut(&mut self) -> IdSlabIterMut<T> {
139        let remaining = self.len();
140        IdSlabIterMut { slab: self, slot: 0, remaining }
141    }
142}
143
144impl<'a, T> IdSlabVacantEntry<'a, T> {
145    pub fn id(&self) -> u64 {
146        self.id
147    }
148
149    pub fn insert(self, value: T) -> u64 {
150        let entry = &mut self.slab.entries[self.slot];
151        assert!(entry.id & EMPTY_BIT == EMPTY_BIT);
152
153        let id = EMPTY_BIT_NEG & entry.id;
154        assert!((id & EMPTY_BIT) == 0);
155
156        entry.id = id;
157        entry.value.write(value);
158
159        let _ = ManuallyDrop::new(self);
160        id
161    }
162}
163
164impl<'a, T> Drop for IdSlabVacantEntry<'a, T> {
165    fn drop(&mut self) {
166        self.slab.free_slots.push(self.slot);
167    }
168}
169
170pub struct IdSlabIter<'a, T> {
171    slab: &'a IdSlab<T>,
172    slot: usize,
173    remaining: usize,
174}
175
176impl<'a, T> Iterator for IdSlabIter<'a, T> {
177    type Item = &'a T;
178
179    fn next(&mut self) -> Option<Self::Item> {
180        while self.slot < self.slab.entries.len() && self.remaining > 0 {
181            let entry = &self.slab.entries[self.slot];
182            self.slot += 1;
183
184            if (entry.id & EMPTY_BIT) == EMPTY_BIT {
185                continue;
186            }
187
188            self.remaining -= 1;
189            return Some(unsafe { entry.value.assume_init_ref() });
190        }
191
192        None
193    }
194}
195
196pub struct IdSlabIterMut<'a, T> {
197    slab: &'a mut IdSlab<T>,
198    slot: usize,
199    remaining: usize,
200}
201
202impl<'a, T> Iterator for IdSlabIterMut<'a, T> {
203    type Item = &'a mut T;
204
205    fn next(&mut self) -> Option<Self::Item> {
206        let len = self.slab.entries.len();
207
208        while self.slot < len && self.remaining > 0 {
209            let entry = &mut self.slab.entries[self.slot];
210            self.slot += 1;
211
212            if (entry.id & EMPTY_BIT) == EMPTY_BIT {
213                continue;
214            }
215
216            self.remaining -= 1;
217            return Some(unsafe { std::mem::transmute(entry.value.assume_init_mut()) });
218        }
219
220        None
221    }
222}
223
224#[cfg(test)]
225mod test {
226    use std::collections::HashSet;
227
228    use rand::{seq::SliceRandom, thread_rng};
229
230    use super::IdSlab;
231
232    #[test]
233    fn test() {
234        let mut slab = IdSlab::<String>::with_capacity(16);
235        let world_id = slab.insert("hello world".to_string()).unwrap();
236
237        let mut old_ids = HashSet::new();
238        let mut ids = Vec::new();
239
240        for i in 0..100 {
241            for j in 0..8 {
242                let entry = slab.vacant_entry().unwrap();
243                assert!(old_ids.insert(entry.id()));
244
245                ids.push(entry.insert(format!("{} - {}", i, j)));
246                assert_eq!(slab.len(), ids.len() + 1);
247            }
248
249            ids.shuffle(&mut thread_rng());
250            while 7 < ids.len() {
251                let id = ids.pop().unwrap();
252
253                slab.remove(id).unwrap();
254                assert_eq!(slab.len(), ids.len() + 1);
255            }
256        }
257
258        assert_eq!(slab.get(world_id).unwrap(), "hello world");
259        for id in ids {
260            slab.remove(id).unwrap();
261        }
262    }
263}
playit_agent_core/utils/id_slab.rs

playit_agent_core/utils/
id_slab.rs
