regalloc2/
indexset.rs

1/*
2 * Released under the terms of the Apache 2.0 license with LLVM
3 * exception. See `LICENSE` for details.
4 */
5
6//! Index sets: sets of integers that represent indices into a space.
7
8use alloc::vec::Vec;
9use core::cell::Cell;
10
11use crate::FxHashMap;
12
13const SMALL_ELEMS: usize = 12;
14
15/// A hybrid large/small-mode sparse mapping from integer indices to
16/// elements.
17///
18/// The trailing `(u32, u64)` elements in each variant is a one-item
19/// cache to allow fast access when streaming through.
20#[derive(Clone, Debug)]
21enum AdaptiveMap {
22    Small {
23        len: u32,
24        keys: [u32; SMALL_ELEMS],
25        values: [u64; SMALL_ELEMS],
26    },
27    Large(FxHashMap<u32, u64>),
28}
29
30const INVALID: u32 = 0xffff_ffff;
31
32impl AdaptiveMap {
33    fn new() -> Self {
34        Self::Small {
35            len: 0,
36            keys: [INVALID; SMALL_ELEMS],
37            values: [0; SMALL_ELEMS],
38        }
39    }
40
41    #[inline(always)]
42    fn get_or_insert<'a>(&'a mut self, key: u32) -> &'a mut u64 {
43        // Check whether the key is present and we are in small mode;
44        // if no to both, we need to expand first.
45        let small_mode_idx = match self {
46            &mut Self::Small {
47                len,
48                ref mut keys,
49                ref values,
50            } => {
51                // Perform this scan but do not return right away;
52                // doing so runs into overlapping-borrow issues
53                // because the current non-lexical lifetimes
54                // implementation is not able to see that the `self`
55                // mutable borrow on return is only on the
56                // early-return path.
57                if let Some(i) = keys[..len as usize].iter().position(|&k| k == key) {
58                    Some(i)
59                } else if len != SMALL_ELEMS as u32 {
60                    debug_assert!(len < SMALL_ELEMS as u32);
61                    None
62                } else if let Some(i) = values.iter().position(|&v| v == 0) {
63                    // If an existing value is zero, reuse that slot.
64                    keys[i] = key;
65                    Some(i)
66                } else {
67                    *self = Self::Large(keys.iter().copied().zip(values.iter().copied()).collect());
68                    None
69                }
70            }
71            _ => None,
72        };
73
74        match self {
75            Self::Small { len, keys, values } => {
76                // If we found the key already while checking whether
77                // we need to expand above, use that index to return
78                // early.
79                if let Some(i) = small_mode_idx {
80                    return &mut values[i];
81                }
82                // Otherwise, the key must not be present; add a new
83                // entry.
84                debug_assert!(*len < SMALL_ELEMS as u32);
85                let idx = *len as usize;
86                *len += 1;
87                keys[idx] = key;
88                values[idx] = 0;
89                &mut values[idx]
90            }
91            Self::Large(map) => map.entry(key).or_insert(0),
92        }
93    }
94
95    #[inline(always)]
96    fn get_mut(&mut self, key: u32) -> Option<&mut u64> {
97        match self {
98            &mut Self::Small {
99                len,
100                ref keys,
101                ref mut values,
102            } => {
103                for i in 0..len {
104                    if keys[i as usize] == key {
105                        return Some(&mut values[i as usize]);
106                    }
107                }
108                None
109            }
110            &mut Self::Large(ref mut map) => map.get_mut(&key),
111        }
112    }
113    #[inline(always)]
114    fn get(&self, key: u32) -> Option<u64> {
115        match self {
116            &Self::Small {
117                len,
118                ref keys,
119                ref values,
120            } => {
121                for i in 0..len {
122                    if keys[i as usize] == key {
123                        let value = values[i as usize];
124                        return Some(value);
125                    }
126                }
127                None
128            }
129            &Self::Large(ref map) => {
130                let value = map.get(&key).cloned();
131                value
132            }
133        }
134    }
135    fn iter<'a>(&'a self) -> AdaptiveMapIter<'a> {
136        match self {
137            &Self::Small {
138                len,
139                ref keys,
140                ref values,
141            } => AdaptiveMapIter::Small(&keys[0..len as usize], &values[0..len as usize]),
142            &Self::Large(ref map) => AdaptiveMapIter::Large(map.iter()),
143        }
144    }
145
146    fn is_empty(&self) -> bool {
147        match self {
148            AdaptiveMap::Small { values, .. } => values.iter().all(|&value| value == 0),
149            AdaptiveMap::Large(m) => m.values().all(|&value| value == 0),
150        }
151    }
152}
153
154enum AdaptiveMapIter<'a> {
155    Small(&'a [u32], &'a [u64]),
156    Large(hashbrown::hash_map::Iter<'a, u32, u64>),
157}
158
159impl<'a> core::iter::Iterator for AdaptiveMapIter<'a> {
160    type Item = (u32, u64);
161
162    #[inline]
163    fn next(&mut self) -> Option<Self::Item> {
164        match self {
165            &mut Self::Small(ref mut keys, ref mut values) => {
166                if keys.is_empty() {
167                    None
168                } else {
169                    let (k, v) = ((*keys)[0], (*values)[0]);
170                    *keys = &(*keys)[1..];
171                    *values = &(*values)[1..];
172                    Some((k, v))
173                }
174            }
175            &mut Self::Large(ref mut it) => it.next().map(|(&k, &v)| (k, v)),
176        }
177    }
178}
179
180/// A conceptually infinite-length set of indices that allows union
181/// and efficient iteration over elements.
182#[derive(Clone)]
183pub struct IndexSet {
184    elems: AdaptiveMap,
185    cache: Cell<(u32, u64)>,
186}
187
188const BITS_PER_WORD: usize = 64;
189
190impl IndexSet {
191    pub fn new() -> Self {
192        Self {
193            elems: AdaptiveMap::new(),
194            cache: Cell::new((INVALID, 0)),
195        }
196    }
197
198    #[inline(always)]
199    fn elem(&mut self, bit_index: usize) -> &mut u64 {
200        let word_index = (bit_index / BITS_PER_WORD) as u32;
201        if self.cache.get().0 == word_index {
202            self.cache.set((INVALID, 0));
203        }
204        self.elems.get_or_insert(word_index)
205    }
206
207    #[inline(always)]
208    fn maybe_elem_mut(&mut self, bit_index: usize) -> Option<&mut u64> {
209        let word_index = (bit_index / BITS_PER_WORD) as u32;
210        if self.cache.get().0 == word_index {
211            self.cache.set((INVALID, 0));
212        }
213        self.elems.get_mut(word_index)
214    }
215
216    #[inline(always)]
217    fn maybe_elem(&self, bit_index: usize) -> Option<u64> {
218        let word_index = (bit_index / BITS_PER_WORD) as u32;
219        if self.cache.get().0 == word_index {
220            Some(self.cache.get().1)
221        } else {
222            self.elems.get(word_index)
223        }
224    }
225
226    #[inline(always)]
227    pub fn set(&mut self, idx: usize, val: bool) {
228        let bit = idx % BITS_PER_WORD;
229        if val {
230            *self.elem(idx) |= 1 << bit;
231        } else if let Some(word) = self.maybe_elem_mut(idx) {
232            *word &= !(1 << bit);
233        }
234    }
235
236    pub fn assign(&mut self, other: &Self) {
237        self.elems = other.elems.clone();
238        self.cache = other.cache.clone();
239    }
240
241    #[inline(always)]
242    pub fn get(&self, idx: usize) -> bool {
243        let bit = idx % BITS_PER_WORD;
244        if let Some(word) = self.maybe_elem(idx) {
245            (word & (1 << bit)) != 0
246        } else {
247            false
248        }
249    }
250
251    pub fn union_with(&mut self, other: &Self) -> bool {
252        let mut changed = 0;
253        for (word_idx, bits) in other.elems.iter() {
254            if bits == 0 {
255                continue;
256            }
257            let word_idx = word_idx as usize;
258            let self_word = self.elem(word_idx * BITS_PER_WORD);
259            changed |= bits & !*self_word;
260            *self_word |= bits;
261        }
262        changed != 0
263    }
264
265    pub fn iter<'a>(&'a self) -> impl Iterator<Item = usize> + 'a {
266        self.elems.iter().flat_map(|(word_idx, bits)| {
267            let word_idx = word_idx as usize;
268            SetBitsIter(bits).map(move |i| BITS_PER_WORD * word_idx + i)
269        })
270    }
271
272    /// Is the adaptive data structure in "small" mode? This is meant
273    /// for testing assertions only.
274    pub(crate) fn is_small(&self) -> bool {
275        match &self.elems {
276            &AdaptiveMap::Small { .. } => true,
277            _ => false,
278        }
279    }
280
281    /// Is the set empty?
282    pub(crate) fn is_empty(&self) -> bool {
283        self.elems.is_empty()
284    }
285}
286
287pub struct SetBitsIter(u64);
288
289impl Iterator for SetBitsIter {
290    type Item = usize;
291
292    #[inline]
293    fn next(&mut self) -> Option<usize> {
294        // Build an `Option<NonZeroU64>` so that on the nonzero path,
295        // the compiler can optimize the trailing-zeroes operator
296        // using that knowledge.
297        core::num::NonZeroU64::new(self.0).map(|nz| {
298            let bitidx = nz.trailing_zeros();
299            self.0 &= self.0 - 1; // clear highest set bit
300            bitidx as usize
301        })
302    }
303}
304
305impl core::fmt::Debug for IndexSet {
306    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
307        let vals = self.iter().collect::<Vec<_>>();
308        write!(f, "{:?}", vals)
309    }
310}
311
312#[cfg(test)]
313mod test {
314    use super::IndexSet;
315
316    #[test]
317    fn test_set_bits_iter() {
318        let mut vec = IndexSet::new();
319        let mut sum = 0;
320        for i in 0..1024 {
321            if i % 17 == 0 {
322                vec.set(i, true);
323                sum += i;
324            }
325        }
326
327        let mut checksum = 0;
328        for bit in vec.iter() {
329            debug_assert!(bit % 17 == 0);
330            checksum += bit;
331        }
332
333        debug_assert_eq!(sum, checksum);
334    }
335
336    #[test]
337    fn test_expand_remove_zero_elems() {
338        let mut vec = IndexSet::new();
339        // Set 12 different words (this is the max small-mode size).
340        for i in 0..12 {
341            vec.set(64 * i, true);
342        }
343        // Now clear a bit, and set a bit in a different word. We
344        // should still be in small mode.
345        vec.set(64 * 5, false);
346        vec.set(64 * 100, true);
347        debug_assert!(vec.is_small());
348    }
349}
regalloc2/indexset.rs

regalloc2/
indexset.rs
