Nothing to see here, move along meow
1use crate::error::KernelError;
2use crate::static_vec::StaticVec;
3use core::mem::MaybeUninit;
4
5const END_OF_LIST: u32 = u32::MAX;
6const ACTIVE_MARKER: u32 = u32::MAX - 1;
7
8#[repr(C)]
9pub struct SlotEntry<T> {
10 payload: MaybeUninit<T>,
11 generation: u32,
12 next_free: u32,
13}
14
15impl<T> SlotEntry<T> {
16 const fn is_active(&self) -> bool {
17 self.next_free == ACTIVE_MARKER
18 }
19
20 const fn empty_free(next: u32) -> Self {
21 Self {
22 payload: MaybeUninit::uninit(),
23 generation: 0,
24 next_free: next,
25 }
26 }
27}
28
29struct MagazineLayer<const M: usize> {
30 loaded: StaticVec<u32, M>,
31 previous: StaticVec<u32, M>,
32}
33
34impl<const M: usize> MagazineLayer<M> {
35 const fn new() -> Self {
36 Self {
37 loaded: StaticVec::new(),
38 previous: StaticVec::new(),
39 }
40 }
41
42 fn alloc(&mut self) -> Option<u32> {
43 self.loaded.pop().or_else(|| {
44 core::mem::swap(&mut self.loaded, &mut self.previous);
45 self.loaded.pop()
46 })
47 }
48
49 fn free(&mut self, id: u32) -> bool {
50 self.loaded.push(id).is_ok() || {
51 core::mem::swap(&mut self.loaded, &mut self.previous);
52 self.loaded.push(id).is_ok()
53 }
54 }
55
56 fn drain_matching(&mut self, leaf_base: u32, leaf_end: u32) {
57 Self::retain_outside(&mut self.loaded, leaf_base, leaf_end);
58 Self::retain_outside(&mut self.previous, leaf_base, leaf_end);
59 }
60
61 fn retain_outside(vec: &mut StaticVec<u32, M>, base: u32, end: u32) {
62 let mut kept = StaticVec::<u32, M>::new();
63 vec.iter()
64 .filter(|&&id| id < base || id >= end)
65 .for_each(|&id| {
66 let _ = kept.push(id);
67 });
68 *vec = kept;
69 }
70}
71
72pub struct SlotMap<T, const LEAF: usize, const MAX_LEAVES: usize> {
73 root: [Option<*mut [SlotEntry<T>; LEAF]>; MAX_LEAVES],
74 free_head: u32,
75 leaf_count: u16,
76 active_count: u32,
77 capacity: u32,
78 magazine: MagazineLayer<32>,
79 leaf_active: [u16; MAX_LEAVES],
80}
81
82unsafe impl<T: Send, const LEAF: usize, const MAX_LEAVES: usize> Send
83 for SlotMap<T, LEAF, MAX_LEAVES>
84{
85}
86
87impl<T, const LEAF: usize, const MAX_LEAVES: usize> SlotMap<T, LEAF, MAX_LEAVES> {
88 pub const fn new() -> Self {
89 Self {
90 root: [None; MAX_LEAVES],
91 free_head: END_OF_LIST,
92 leaf_count: 0,
93 active_count: 0,
94 capacity: 0,
95 magazine: MagazineLayer::new(),
96 leaf_active: [0; MAX_LEAVES],
97 }
98 }
99
100 #[allow(clippy::missing_safety_doc)]
101 pub unsafe fn expand(&mut self, leaf: *mut [SlotEntry<T>; LEAF]) {
102 let leaf_idx = self.leaf_count as usize;
103 assert!(leaf_idx < MAX_LEAVES, "SlotMap: too many leaves");
104 let base = (leaf_idx as u32) * (LEAF as u32);
105
106 (0..LEAF).rev().for_each(|slot| {
107 let next = match slot + 1 < LEAF {
108 true => base + (slot as u32 + 1),
109 false => self.free_head,
110 };
111 unsafe {
112 (*leaf)[slot] = SlotEntry::empty_free(next);
113 }
114 });
115
116 self.free_head = base;
117 self.root[leaf_idx] = Some(leaf);
118 self.leaf_active[leaf_idx] = 0;
119 self.leaf_count += 1;
120 self.capacity += LEAF as u32;
121 }
122
123 fn entry_ptr(&self, index: u32) -> Option<*mut SlotEntry<T>> {
124 let leaf_idx = (index as usize) / LEAF;
125 let slot_idx = (index as usize) % LEAF;
126 self.root
127 .get(leaf_idx)
128 .and_then(|opt| *opt)
129 .map(|leaf| unsafe { &raw mut (*leaf)[slot_idx] })
130 }
131
132 fn resolve(&self, index: u32) -> Option<&SlotEntry<T>> {
133 self.entry_ptr(index).map(|p| unsafe { &*p })
134 }
135
136 fn resolve_mut(&mut self, index: u32) -> Option<&mut SlotEntry<T>> {
137 self.entry_ptr(index).map(|p| unsafe { &mut *p })
138 }
139
140 pub fn allocate(&mut self, value: T) -> Option<(u32, u32)> {
141 let idx = self.magazine.alloc().or_else(|| match self.free_head {
142 END_OF_LIST => None,
143 head => {
144 let ptr = self.entry_ptr(head)?;
145 self.free_head = unsafe { (*ptr).next_free };
146 Some(head)
147 }
148 })?;
149
150 let ptr = self.entry_ptr(idx)?;
151 unsafe {
152 (*ptr).payload = MaybeUninit::new(value);
153 (*ptr).next_free = ACTIVE_MARKER;
154 let generation = (*ptr).generation;
155 self.active_count += 1;
156 self.leaf_active[(idx as usize) / LEAF] += 1;
157 Some((idx, generation))
158 }
159 }
160
161 pub fn free(&mut self, index: u32, expected_gen: u32) -> Result<T, KernelError> {
162 let ptr = self.entry_ptr(index).ok_or(KernelError::InvalidObject)?;
163
164 let entry = unsafe { &mut *ptr };
165
166 match (entry.generation == expected_gen, entry.is_active()) {
167 (false, _) => return Err(KernelError::StaleGeneration),
168 (_, false) => return Err(KernelError::InvalidObject),
169 _ => {}
170 }
171
172 let value = unsafe { entry.payload.assume_init_read() };
173 entry.payload = MaybeUninit::uninit();
174 entry.generation = entry.generation.wrapping_add(1);
175
176 match self.magazine.free(index) {
177 true => {
178 entry.next_free = END_OF_LIST;
179 }
180 false => {
181 entry.next_free = self.free_head;
182 self.free_head = index;
183 }
184 }
185
186 debug_assert!(self.active_count > 0, "active_count underflow in free");
187 self.active_count = self.active_count.saturating_sub(1);
188 let leaf_idx = (index as usize) / LEAF;
189 debug_assert!(self.leaf_active[leaf_idx] > 0, "leaf_active underflow");
190 self.leaf_active[leaf_idx] = self.leaf_active[leaf_idx].saturating_sub(1);
191 Ok(value)
192 }
193
194 pub fn try_reclaim_trailing(&mut self) -> Option<*mut [SlotEntry<T>; LEAF]> {
195 let last = self.leaf_count.checked_sub(1)? as usize;
196 match self.leaf_active[last] {
197 0 => {}
198 _ => return None,
199 }
200 let leaf_ptr = self.root[last]?;
201
202 let base = (last as u32) * (LEAF as u32);
203 let end = base + LEAF as u32;
204
205 self.magazine.drain_matching(base, end);
206
207 let root_snapshot: [Option<*mut [SlotEntry<T>; LEAF]>; MAX_LEAVES] = self.root;
208 let resolve_raw = |index: u32| -> *mut SlotEntry<T> {
209 let leaf_idx = (index as usize) / LEAF;
210 let slot_idx = (index as usize) % LEAF;
211 let leaf = root_snapshot[leaf_idx].expect("free list references missing leaf");
212 unsafe { &raw mut (*leaf)[slot_idx] }
213 };
214
215 let mut cursor = self.free_head;
216 self.free_head = END_OF_LIST;
217 let mut new_tail: u32 = END_OF_LIST;
218
219 core::iter::from_fn(|| {
220 (cursor != END_OF_LIST).then(|| {
221 let idx = cursor;
222 let ptr = resolve_raw(idx);
223 cursor = unsafe { (*ptr).next_free };
224 unsafe { (*ptr).next_free = END_OF_LIST };
225 idx
226 })
227 })
228 .filter(|&idx| idx < base || idx >= end)
229 .for_each(|idx| {
230 match new_tail {
231 END_OF_LIST => self.free_head = idx,
232 tail => unsafe { (*resolve_raw(tail)).next_free = idx },
233 }
234 new_tail = idx;
235 });
236
237 self.root[last] = None;
238 self.leaf_count -= 1;
239 self.capacity -= LEAF as u32;
240
241 Some(leaf_ptr)
242 }
243
244 pub fn leaf_active_count(&self, leaf_idx: usize) -> u16 {
245 self.leaf_active[leaf_idx]
246 }
247
248 pub fn get(&self, index: u32) -> Option<&T> {
249 self.resolve(index)
250 .filter(|e| e.is_active())
251 .map(|e| unsafe { e.payload.assume_init_ref() })
252 }
253
254 pub fn get_checked(&self, index: u32, expected_gen: u32) -> Option<&T> {
255 self.resolve(index)
256 .filter(|e| e.is_active() && e.generation == expected_gen)
257 .map(|e| unsafe { e.payload.assume_init_ref() })
258 }
259
260 pub fn get_mut(&mut self, index: u32) -> Option<&mut T> {
261 self.resolve_mut(index)
262 .filter(|e| e.is_active())
263 .map(|e| unsafe { e.payload.assume_init_mut() })
264 }
265
266 pub fn get_mut_checked(&mut self, index: u32, expected_gen: u32) -> Option<&mut T> {
267 self.resolve_mut(index)
268 .filter(|e| e.is_active() && e.generation == expected_gen)
269 .map(|e| unsafe { e.payload.assume_init_mut() })
270 }
271
272 pub fn generation_of(&self, index: u32) -> Option<u32> {
273 self.resolve(index).map(|e| e.generation)
274 }
275
276 pub fn for_each_active(&self, mut f: impl FnMut(u32, &T)) {
277 (0..self.leaf_count as usize).for_each(|leaf_idx| {
278 let Some(leaf_ptr) = self.root[leaf_idx] else {
279 return;
280 };
281 let base = (leaf_idx as u32) * (LEAF as u32);
282 (0..LEAF).for_each(|slot| {
283 let entry = unsafe { &(*leaf_ptr)[slot] };
284 if entry.is_active() {
285 f(base + slot as u32, unsafe {
286 entry.payload.assume_init_ref()
287 });
288 }
289 });
290 });
291 }
292
293 pub const fn active_count(&self) -> u32 {
294 self.active_count
295 }
296
297 pub const fn capacity(&self) -> u32 {
298 self.capacity
299 }
300
301 pub const fn leaf_count(&self) -> u16 {
302 self.leaf_count
303 }
304}