cursed_collections/
lazy_array.rs

1use ::alloc::alloc;
2use core::{mem, ptr, slice};
3
4/// A collection with a size defined at creation, but where entries are initialized later.
5///
6/// It is useful when you are reading nodes of an acyclic graph, where entries can be read as you
7/// need them.
8///
9/// # Example
10///
11/// ```
12/// # use cursed_collections::LazyArray;
13/// #
14/// #[derive(Debug, Eq, PartialEq)]
15/// struct Entry<'heap> {
16///     value: i32,
17///     next: Option<&'heap Entry<'heap>>,
18/// }
19///
20/// fn do_something<'heap>(heap: &'heap LazyArray<Entry<'heap>>) {
21///     let entry_0 = heap.get_or_insert(3, Entry { value: 123, next: None });
22///     let entry_1 = heap.get_or_insert(6, Entry { value: 456, next: Some(entry_0) });
23///     assert_eq!(Some(123), entry_1.next.map(|inner| inner.value));
24///     assert_eq!(None, heap.get(2));
25/// }
26/// ```
27#[derive(Debug)]
28pub struct LazyArray<T> {
29    buffer: *mut Option<T>,
30    capacity: usize,
31    layout: alloc::Layout,
32}
33
34impl<T> LazyArray<T> {
35    pub fn new(capacity: usize) -> Self {
36        unsafe {
37            let layout = alloc::Layout::array::<Option<T>>(capacity).expect("size overflow");
38            let buffer = alloc::alloc(layout);
39            {
40                let slice =
41                    slice::from_raw_parts_mut(buffer as *mut mem::MaybeUninit<Option<T>>, capacity);
42                for i in slice {
43                    *i = mem::MaybeUninit::new(None);
44                }
45            }
46            Self {
47                buffer: buffer as *mut Option<T>,
48                capacity,
49                layout,
50            }
51        }
52    }
53
54    fn entry(&self, index: usize) -> *mut Option<T> {
55        unsafe {
56            assert!(index < self.capacity);
57            self.buffer.add(index)
58        }
59    }
60
61    pub fn get_or_insert(&self, index: usize, t: T) -> &T {
62        unsafe {
63            // We cannot use Option::get_or_insert because we need to construct a &mut, which is
64            // unsound if it is already initialized because there may be & existing!
65            let entry = self.entry(index);
66            match *entry {
67                None => {
68                    ptr::write(entry, Some(t));
69                    (*entry).as_ref().unwrap_unchecked()
70                }
71                Some(ref v) => v,
72            }
73        }
74    }
75
76    pub fn get(&self, index: usize) -> Option<&T> {
77        assert!(index < self.capacity);
78        unsafe { (*self.buffer.add(index)).as_ref() }
79    }
80}
81
82impl<T> Drop for LazyArray<T> {
83    fn drop(&mut self) {
84        unsafe {
85            alloc::dealloc(self.buffer as *mut u8, self.layout);
86        }
87    }
88}
89
90impl<T> Default for LazyArray<T> {
91    fn default() -> Self {
92        Self::new(0)
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::LazyArray;
99
100    #[test]
101    fn it_works() {
102        let lazy_array = LazyArray::<i32>::new(10);
103        for i in 0..10 {
104            assert_eq!(lazy_array.get(i), None)
105        }
106
107        assert_eq!(lazy_array.get_or_insert(7, 112233), &112233);
108
109        for i in 0..10 {
110            assert_eq!(lazy_array.get(i), if i == 7 { Some(&112233) } else { None })
111        }
112    }
113
114    #[test]
115    fn cannot_insert_twice() {
116        let lazy_array = LazyArray::<i32>::new(10);
117        assert_eq!(lazy_array.get_or_insert(7, 112233), &112233);
118        assert_eq!(lazy_array.get_or_insert(7, 445566), &112233);
119    }
120
121    #[test]
122    #[should_panic]
123    fn cannot_put_out_of_bounds() {
124        let lazy_array = LazyArray::<i32>::new(10);
125        lazy_array.get_or_insert(10, 112233);
126    }
127}
cursed_collections/lazy_array.rs

cursed_collections/
lazy_array.rs
