rune/compile/v1/
slots.rs

Help
1//! This is a specialized slab used to allocate slots of memory for the compiler.
2
3#![allow(clippy::bool_assert_comparison)]
4
5#[cfg(test)]
6mod tests;
7
8use core::fmt;
9use core::slice;
10
11use crate::alloc::{self, Vec};
12
13pub(super) struct Slots {
14    storage: Vec<u128>,
15    head: usize,
16}
17
18impl Slots {
19    /// Construct a new empty slab.
20    pub(super) const fn new() -> Self {
21        Self {
22            storage: Vec::new(),
23            head: 0,
24        }
25    }
26
27    /// Allocate the first free variable.
28    #[tracing::instrument(ret(level = tracing::Level::TRACE), skip(self))]
29    pub(super) fn insert(&mut self) -> alloc::Result<usize> {
30        let mut key = (u128::BITS as usize) * self.head;
31
32        for bits in self
33            .storage
34            .get_mut(self.head..)
35            .unwrap_or_default()
36            .iter_mut()
37        {
38            if *bits == u128::MAX {
39                key += u128::BITS as usize;
40                self.head += 1;
41                continue;
42            }
43
44            let o = bits.trailing_ones();
45            key += o as usize;
46            *bits |= 1 << o;
47            return Ok(key);
48        }
49
50        self.head = self.storage.len();
51        self.storage.try_push(1)?;
52        Ok(key)
53    }
54
55    #[tracing::instrument(ret(level = tracing::Level::TRACE), skip(self))]
56    pub(super) fn push(&mut self) -> alloc::Result<usize> {
57        let mut last = None;
58
59        let key = 'key: {
60            for (n, bits) in self.storage.iter_mut().enumerate().rev() {
61                let o = bits.leading_zeros();
62
63                // Whole segment is free, skip over it.
64                if o == u128::BITS {
65                    last = Some((n, bits));
66                    continue;
67                }
68
69                let key = (u128::BITS as usize) * n;
70
71                // There is no space in this segment.
72                if o == 0 {
73                    break 'key key + u128::BITS as usize;
74                }
75
76                let o = u128::BITS - o;
77                *bits |= 1 << o;
78                return Ok(key + o as usize);
79            }
80
81            0
82        };
83
84        if let Some((n, bits)) = last {
85            *bits = 1;
86            self.storage.truncate(n + 1);
87        } else {
88            self.storage.try_push(1)?;
89        }
90
91        Ok(key)
92    }
93
94    #[tracing::instrument(ret(level = tracing::Level::TRACE), skip(self))]
95    pub(super) fn remove(&mut self, key: usize) -> bool {
96        let index = key / (u128::BITS as usize);
97
98        let Some(bits) = self.storage.get_mut(index) else {
99            return false;
100        };
101
102        self.head = index;
103        let o = key % (u128::BITS as usize);
104        let existed = *bits & (1 << o) != 0;
105        *bits &= !(1 << o);
106        existed
107    }
108
109    fn iter(&self) -> Iter<'_> {
110        let (current, rest) = match &self.storage[..] {
111            [first, rest @ ..] => (*first, rest),
112            [] => (0, &[][..]),
113        };
114
115        Iter {
116            storage: rest.iter(),
117            current,
118            key: 0,
119        }
120    }
121}
122
123impl fmt::Debug for Slots {
124    #[inline]
125    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
126        f.debug_set().entries(self.iter()).finish()
127    }
128}
129
130struct Iter<'a> {
131    storage: slice::Iter<'a, u128>,
132    current: u128,
133    key: usize,
134}
135
136impl Iterator for Iter<'_> {
137    type Item = usize;
138
139    fn next(&mut self) -> Option<Self::Item> {
140        loop {
141            let o = self.current.trailing_zeros();
142
143            if o != u128::BITS {
144                self.current ^= 1 << o;
145                return Some(self.key + o as usize);
146            }
147
148            self.current = match self.storage.next() {
149                Some(current) => *current,
150                None => return None,
151            };
152
153            self.key += u128::BITS as usize;
154        }
155    }
156}
rune/compile/v1/slots.rs

rune/compile/v1/
slots.rs
