lru_slab/
lib.rs

1//! Pre-allocated storage with constant-time LRU tracking
2
3#![warn(missing_docs)]
4#![no_std]
5
6extern crate alloc;
7
8use alloc::boxed::Box;
9use core::{fmt, iter::FusedIterator, marker::PhantomData, ptr::addr_of_mut};
10
11/// A random-access table that maintains an LRU list in constant time
12#[derive(Clone)]
13pub struct LruSlab<T> {
14    slots: Box<[Slot<T>]>,
15    /// Most recently used
16    head: u32,
17    /// Least recently used
18    tail: u32,
19    /// First unused
20    free: u32,
21    /// Number of occupied slots
22    len: u32,
23}
24
25impl<T> LruSlab<T> {
26    /// Create an empty [`LruSlab`]
27    pub fn new() -> Self {
28        Self::with_capacity(0)
29    }
30
31    /// Create an [`LruSlab`] that can store at least `capacity` elements without reallocating
32    pub fn with_capacity(capacity: u32) -> Self {
33        assert!(capacity != u32::MAX, "capacity too large");
34        Self {
35            slots: (0..capacity)
36                .map(|n| Slot {
37                    value: None,
38                    prev: NONE,
39                    next: if n + 1 == capacity { NONE } else { n + 1 },
40                })
41                .collect(),
42            head: NONE,
43            tail: NONE,
44            free: if capacity == 0 { NONE } else { 0 },
45            len: 0,
46        }
47    }
48
49    /// Whether no elements are stored
50    pub fn is_empty(&self) -> bool {
51        self.len == 0
52    }
53
54    /// Number of elements stored
55    pub fn len(&self) -> u32 {
56        self.len
57    }
58
59    /// Number of elements that can be stored without reallocating
60    pub fn capacity(&self) -> u32 {
61        self.slots.len() as u32
62    }
63
64    /// The slot that will be returned by the next call to `insert`, unless `remove` is called first
65    pub fn vacant_key(&self) -> u32 {
66        match self.free {
67            NONE => self.capacity(),
68            _ => self.free,
69        }
70    }
71
72    /// Insert a value, returning the slot it was stored in
73    ///
74    /// The returned slot is marked as the most recently used.
75    pub fn insert(&mut self, value: T) -> u32 {
76        let id = match self.alloc() {
77            Some(id) => id,
78            None => {
79                let len = self.capacity();
80                // Ensure `NONE` never becomes a real slot index
81                let cap = 2u32.saturating_mul(len.max(2)).min(u32::MAX - 1);
82                if cap == len {
83                    panic!("cannot store more than 2^32-2 elements");
84                }
85                self.slots = self
86                    .slots
87                    .iter_mut()
88                    .map(|x| Slot {
89                        value: x.value.take(),
90                        next: x.next,
91                        prev: x.prev,
92                    })
93                    .chain((len..cap).map(|n| Slot {
94                        value: None,
95                        prev: NONE,
96                        next: if n + 1 == cap { NONE } else { n + 1 },
97                    }))
98                    .collect();
99                self.free = len + 1;
100                len
101            }
102        };
103        let idx = id as usize;
104
105        debug_assert!(self.slots[idx].value.is_none(), "corrupt free list");
106        self.slots[idx].value = Some(value);
107        self.link_at_head(id);
108        self.len += 1;
109
110        id
111    }
112
113    /// Get the least recently used slot, if any
114    pub fn lru(&self) -> Option<u32> {
115        if self.tail == NONE {
116            debug_assert_eq!(self.head, NONE);
117            None
118        } else {
119            Some(self.tail)
120        }
121    }
122
123    /// Remove the element stored in `slot`, returning it
124    pub fn remove(&mut self, slot: u32) -> T {
125        self.unlink(slot);
126        self.slots[slot as usize].next = self.free;
127        self.slots[slot as usize].prev = NONE;
128        self.free = slot;
129        self.len -= 1;
130        self.slots[slot as usize]
131            .value
132            .take()
133            .expect("removing empty slot")
134    }
135
136    /// Mark `slot` as the most recently used and access it uniquely
137    pub fn get_mut(&mut self, slot: u32) -> &mut T {
138        self.freshen(slot);
139        self.peek_mut(slot)
140    }
141
142    /// Access `slot` without marking it as most recently used
143    pub fn peek(&self, slot: u32) -> &T {
144        self.slots[slot as usize].value.as_ref().unwrap()
145    }
146
147    /// Access `slot` uniquely without marking it as most recently used
148    pub fn peek_mut(&mut self, slot: u32) -> &mut T {
149        self.slots[slot as usize].value.as_mut().unwrap()
150    }
151
152    /// Walk the container from most to least recently used
153    pub fn iter(&self) -> Iter<'_, T> {
154        let state = IterState::new(self);
155        Iter {
156            slots: &self.slots[..],
157            state,
158        }
159    }
160
161    /// Walk the container from most to least recently used
162    pub fn iter_mut(&mut self) -> IterMut<'_, T> {
163        let state = IterState::new(self);
164        IterMut {
165            slots: self.slots[..].as_mut_ptr().cast(),
166            state,
167            _marker: PhantomData,
168        }
169    }
170
171    /// Remove a slot from the freelist
172    fn alloc(&mut self) -> Option<u32> {
173        if self.free == NONE {
174            return None;
175        }
176        let slot = self.free;
177        self.free = self.slots[slot as usize].next;
178        Some(slot)
179    }
180
181    /// Mark `slot` as the most recently used
182    fn freshen(&mut self, slot: u32) {
183        if self.slots[slot as usize].prev == NONE {
184            // This is already the freshest slot, so we don't need to do anything
185            debug_assert_eq!(self.head, slot, "corrupt LRU list");
186            return;
187        }
188
189        self.unlink(slot);
190        self.link_at_head(slot);
191    }
192
193    /// Add a link to the head of the list
194    fn link_at_head(&mut self, slot: u32) {
195        let idx = slot as usize;
196        if self.head == NONE {
197            // List was empty
198            self.slots[idx].next = NONE;
199            self.tail = slot;
200        } else {
201            self.slots[idx].next = self.head;
202            self.slots[self.head as usize].prev = slot;
203        }
204        self.slots[idx].prev = NONE;
205        self.head = slot;
206    }
207
208    /// Remove a link from anywhere in the list
209    fn unlink(&mut self, slot: u32) {
210        let idx = slot as usize;
211        if self.slots[idx].prev != NONE {
212            self.slots[self.slots[idx].prev as usize].next = self.slots[idx].next;
213        } else {
214            self.head = self.slots[idx].next;
215        }
216        if self.slots[idx].next != NONE {
217            self.slots[self.slots[idx].next as usize].prev = self.slots[idx].prev;
218        } else {
219            // This was the tail
220            self.tail = self.slots[idx].prev;
221        }
222    }
223}
224
225impl<T> Default for LruSlab<T> {
226    fn default() -> Self {
227        Self::new()
228    }
229}
230
231impl<T> FromIterator<T> for LruSlab<T> {
232    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
233        let iter = iter.into_iter();
234        let mut slab = LruSlab::with_capacity(u32::try_from(iter.size_hint().0).unwrap());
235        for x in iter {
236            slab.insert(x);
237        }
238        slab
239    }
240}
241
242impl<'a, T> IntoIterator for &'a LruSlab<T> {
243    type Item = (u32, &'a T);
244
245    type IntoIter = Iter<'a, T>;
246
247    fn into_iter(self) -> Self::IntoIter {
248        self.iter()
249    }
250}
251
252impl<'a, T> IntoIterator for &'a mut LruSlab<T> {
253    type Item = (u32, &'a mut T);
254
255    type IntoIter = IterMut<'a, T>;
256
257    fn into_iter(self) -> Self::IntoIter {
258        self.iter_mut()
259    }
260}
261
262impl<T: fmt::Debug> fmt::Debug for LruSlab<T> {
263    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
264        f.debug_map().entries(self).finish()
265    }
266}
267
268#[derive(Clone)]
269struct Slot<T> {
270    value: Option<T>,
271    /// Next slot in the LRU or free list
272    next: u32,
273    /// Previous slot in the LRU list; NONE when free
274    prev: u32,
275}
276
277const NONE: u32 = u32::MAX;
278
279/// Iterator over elements of an [`LruSlab`], from most to least recently used
280pub struct Iter<'a, T> {
281    slots: &'a [Slot<T>],
282    state: IterState,
283}
284
285impl<'a, T> Iterator for Iter<'a, T> {
286    type Item = (u32, &'a T);
287    fn next(&mut self) -> Option<(u32, &'a T)> {
288        let idx = self.state.next(|i| self.slots[i as usize].next)?;
289        let result = self.slots[idx as usize]
290            .value
291            .as_ref()
292            .expect("corrupt LRU list");
293        Some((idx, result))
294    }
295
296    fn size_hint(&self) -> (usize, Option<usize>) {
297        (self.state.len as usize, Some(self.state.len as usize))
298    }
299}
300
301impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
302    fn next_back(&mut self) -> Option<(u32, &'a T)> {
303        let idx = self.state.next_back(|i| self.slots[i as usize].prev)?;
304        let result = self.slots[idx as usize]
305            .value
306            .as_ref()
307            .expect("corrupt LRU list");
308        Some((idx, result))
309    }
310}
311
312impl<T> ExactSizeIterator for Iter<'_, T> {
313    fn len(&self) -> usize {
314        self.state.len as usize
315    }
316}
317
318impl<T> FusedIterator for Iter<'_, T> {}
319
320/// Iterator over mutable elements of an [`LruSlab`], from most to least recently used
321pub struct IterMut<'a, T> {
322    slots: *mut Slot<T>,
323    state: IterState,
324    _marker: PhantomData<&'a mut [Slot<T>]>,
325}
326
327impl<'a, T> Iterator for IterMut<'a, T> {
328    type Item = (u32, &'a mut T);
329    fn next(&mut self) -> Option<(u32, &'a mut T)> {
330        // Safety: `next` returns unique in-bounds indices, and no live references overlap with any
331        // `next` field
332        unsafe {
333            let idx = self
334                .state
335                .next(|i| *addr_of_mut!((*self.slots.add(i as usize)).next))?;
336            let result = (*addr_of_mut!((*self.slots.add(idx as usize)).value))
337                .as_mut()
338                .expect("corrupt LRU list");
339            Some((idx, result))
340        }
341    }
342
343    fn size_hint(&self) -> (usize, Option<usize>) {
344        (self.state.len as usize, Some(self.state.len as usize))
345    }
346}
347
348impl<'a, T> DoubleEndedIterator for IterMut<'a, T> {
349    fn next_back(&mut self) -> Option<(u32, &'a mut T)> {
350        // Safety: `next_back` returns unique in-bounds indices, and no live references overlap with
351        // any `prev` field
352        unsafe {
353            let idx = self
354                .state
355                .next_back(|i| *addr_of_mut!((*self.slots.add(i as usize)).prev))?;
356            let result = (*addr_of_mut!((*self.slots.add(idx as usize)).value))
357                .as_mut()
358                .expect("corrupt LRU list");
359            Some((idx, result))
360        }
361    }
362}
363
364impl<T> ExactSizeIterator for IterMut<'_, T> {
365    fn len(&self) -> usize {
366        self.state.len as usize
367    }
368}
369
370impl<T> FusedIterator for IterMut<'_, T> {}
371
372struct IterState {
373    head: u32,
374    tail: u32,
375    len: u32,
376}
377
378impl IterState {
379    fn new<T>(slab: &LruSlab<T>) -> Self {
380        Self {
381            head: slab.head,
382            tail: slab.tail,
383            len: slab.len,
384        }
385    }
386
387    fn next(&mut self, get_next: impl Fn(u32) -> u32) -> Option<u32> {
388        if self.len == 0 {
389            return None;
390        }
391        let idx = self.head;
392        self.head = get_next(idx);
393        self.len -= 1;
394        Some(idx)
395    }
396
397    fn next_back(&mut self, get_prev: impl Fn(u32) -> u32) -> Option<u32> {
398        if self.len == 0 {
399            return None;
400        }
401        let idx = self.tail;
402        self.tail = get_prev(idx);
403        self.len -= 1;
404        Some(idx)
405    }
406}
407
408#[cfg(test)]
409mod tests {
410    use alloc::{format, string::String, vec::Vec};
411
412    use super::*;
413
414    #[test]
415    fn lru_order() {
416        let mut cache = LruSlab::new();
417        let b = cache.insert('b');
418        assert_eq!(cache.iter().map(|(_, x)| x).collect::<String>(), "b");
419        let _a = cache.insert('a');
420        assert_eq!(cache.iter().map(|(_, x)| x).collect::<String>(), "ab");
421        let d = cache.insert('d');
422        assert_eq!(cache.iter().map(|(_, x)| x).collect::<String>(), "dab");
423        let c = cache.insert('c');
424        assert_eq!(cache.iter().map(|(_, x)| x).collect::<String>(), "cdab");
425        let e = cache.insert('e');
426        assert_eq!(cache.iter().map(|(_, x)| x).collect::<String>(), "ecdab");
427
428        cache.get_mut(b);
429        cache.get_mut(c);
430        cache.get_mut(d);
431        cache.get_mut(e);
432
433        assert_eq!(cache.remove(cache.lru().unwrap()), 'a');
434        assert_eq!(cache.remove(cache.lru().unwrap()), 'b');
435        assert_eq!(cache.remove(cache.lru().unwrap()), 'c');
436        assert_eq!(cache.remove(cache.lru().unwrap()), 'd');
437        assert_eq!(cache.remove(cache.lru().unwrap()), 'e');
438        assert!(cache.lru().is_none());
439    }
440
441    #[test]
442    fn slot_reuse() {
443        let mut cache = LruSlab::new();
444        let a = cache.insert('a');
445        cache.remove(a);
446        let a_prime = cache.insert('a');
447        assert_eq!(a, a_prime);
448        assert_eq!(cache.len(), 1);
449    }
450
451    #[test]
452    fn debug() {
453        let slab = ['a', 'b'].into_iter().collect::<LruSlab<_>>();
454        assert_eq!(format!("{:?}", slab), "{1: 'b', 0: 'a'}");
455    }
456
457    #[test]
458    fn iter_reverse() {
459        let slab = ['a', 'b'].into_iter().collect::<LruSlab<_>>();
460        let mut double_reversed = slab.iter().rev().collect::<Vec<_>>();
461        double_reversed.reverse();
462        assert_eq!(slab.iter().collect::<Vec<_>>(), double_reversed);
463    }
464
465    #[test]
466    fn vacant_key() {
467        let mut slab = LruSlab::new();
468        assert_eq!(slab.vacant_key(), 0);
469        slab.insert(());
470        assert_eq!(slab.vacant_key(), 1);
471        slab.remove(0);
472        assert_eq!(slab.vacant_key(), 0);
473    }
474}
lru_slab/lib.rs

lru_slab/
lib.rs
