/src/regalloc2/src/indexset.rs
Line | Count | Source |
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 | | |
8 | | use alloc::vec::Vec; |
9 | | use core::cell::Cell; |
10 | | |
11 | | use crate::FxHashMap; |
12 | | |
13 | | const 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)] |
21 | | enum AdaptiveMap { |
22 | | Small { |
23 | | len: u32, |
24 | | keys: [u32; SMALL_ELEMS], |
25 | | values: [u64; SMALL_ELEMS], |
26 | | }, |
27 | | Large(FxHashMap<u32, u64>), |
28 | | } |
29 | | |
30 | | const INVALID: u32 = 0xffff_ffff; |
31 | | |
32 | | impl AdaptiveMap { |
33 | 0 | fn new() -> Self { |
34 | 0 | Self::Small { |
35 | 0 | len: 0, |
36 | 0 | keys: [INVALID; SMALL_ELEMS], |
37 | 0 | values: [0; SMALL_ELEMS], |
38 | 0 | } |
39 | 0 | } |
40 | | |
41 | | #[inline(always)] |
42 | 0 | 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 | 0 | let small_mode_idx = match self { |
46 | | &mut Self::Small { |
47 | 0 | len, |
48 | 0 | ref mut keys, |
49 | 0 | 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 | 0 | if let Some(i) = keys[..len as usize].iter().position(|&k| k == key) { |
58 | 0 | Some(i) |
59 | 0 | } else if len != SMALL_ELEMS as u32 { |
60 | 0 | debug_assert!(len < SMALL_ELEMS as u32); |
61 | 0 | None |
62 | 0 | } else if let Some(i) = values.iter().position(|&v| v == 0) { |
63 | | // If an existing value is zero, reuse that slot. |
64 | 0 | keys[i] = key; |
65 | 0 | Some(i) |
66 | | } else { |
67 | 0 | *self = Self::Large(keys.iter().copied().zip(values.iter().copied()).collect()); |
68 | 0 | None |
69 | | } |
70 | | } |
71 | 0 | _ => None, |
72 | | }; |
73 | | |
74 | 0 | match self { |
75 | 0 | 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 | 0 | if let Some(i) = small_mode_idx { |
80 | 0 | return &mut values[i]; |
81 | 0 | } |
82 | | // Otherwise, the key must not be present; add a new |
83 | | // entry. |
84 | 0 | debug_assert!(*len < SMALL_ELEMS as u32); |
85 | 0 | let idx = *len as usize; |
86 | 0 | *len += 1; |
87 | 0 | keys[idx] = key; |
88 | 0 | values[idx] = 0; |
89 | 0 | &mut values[idx] |
90 | | } |
91 | 0 | Self::Large(map) => map.entry(key).or_insert(0), |
92 | | } |
93 | 0 | } |
94 | | |
95 | | #[inline(always)] |
96 | 0 | fn get_mut(&mut self, key: u32) -> Option<&mut u64> { |
97 | 0 | match self { |
98 | | &mut Self::Small { |
99 | 0 | len, |
100 | 0 | ref keys, |
101 | 0 | ref mut values, |
102 | | } => { |
103 | 0 | for i in 0..len { |
104 | 0 | if keys[i as usize] == key { |
105 | 0 | return Some(&mut values[i as usize]); |
106 | 0 | } |
107 | | } |
108 | 0 | None |
109 | | } |
110 | 0 | &mut Self::Large(ref mut map) => map.get_mut(&key), |
111 | | } |
112 | 0 | } |
113 | | #[inline(always)] |
114 | 0 | fn get(&self, key: u32) -> Option<u64> { |
115 | 0 | match self { |
116 | | &Self::Small { |
117 | 0 | len, |
118 | 0 | ref keys, |
119 | 0 | ref values, |
120 | | } => { |
121 | 0 | for i in 0..len { |
122 | 0 | if keys[i as usize] == key { |
123 | 0 | let value = values[i as usize]; |
124 | 0 | return Some(value); |
125 | 0 | } |
126 | | } |
127 | 0 | None |
128 | | } |
129 | 0 | &Self::Large(ref map) => { |
130 | 0 | let value = map.get(&key).cloned(); |
131 | 0 | value |
132 | | } |
133 | | } |
134 | 0 | } |
135 | 0 | fn iter<'a>(&'a self) -> AdaptiveMapIter<'a> { |
136 | 0 | match self { |
137 | | &Self::Small { |
138 | 0 | len, |
139 | 0 | ref keys, |
140 | 0 | ref values, |
141 | 0 | } => AdaptiveMapIter::Small(&keys[0..len as usize], &values[0..len as usize]), |
142 | 0 | &Self::Large(ref map) => AdaptiveMapIter::Large(map.iter()), |
143 | | } |
144 | 0 | } |
145 | | |
146 | 0 | fn is_empty(&self) -> bool { |
147 | 0 | match self { |
148 | 0 | AdaptiveMap::Small { values, .. } => values.iter().all(|&value| value == 0), |
149 | 0 | AdaptiveMap::Large(m) => m.values().all(|&value| value == 0), |
150 | | } |
151 | 0 | } |
152 | | } |
153 | | |
154 | | enum AdaptiveMapIter<'a> { |
155 | | Small(&'a [u32], &'a [u64]), |
156 | | Large(hashbrown::hash_map::Iter<'a, u32, u64>), |
157 | | } |
158 | | |
159 | | impl<'a> core::iter::Iterator for AdaptiveMapIter<'a> { |
160 | | type Item = (u32, u64); |
161 | | |
162 | | #[inline] |
163 | 0 | fn next(&mut self) -> Option<Self::Item> { |
164 | 0 | match self { |
165 | 0 | &mut Self::Small(ref mut keys, ref mut values) => { |
166 | 0 | if keys.is_empty() { |
167 | 0 | None |
168 | | } else { |
169 | 0 | let (k, v) = ((*keys)[0], (*values)[0]); |
170 | 0 | *keys = &(*keys)[1..]; |
171 | 0 | *values = &(*values)[1..]; |
172 | 0 | Some((k, v)) |
173 | | } |
174 | | } |
175 | 0 | &mut Self::Large(ref mut it) => it.next().map(|(&k, &v)| (k, v)), |
176 | | } |
177 | 0 | } |
178 | | } |
179 | | |
180 | | /// A conceptually infinite-length set of indices that allows union |
181 | | /// and efficient iteration over elements. |
182 | | #[derive(Clone)] |
183 | | pub struct IndexSet { |
184 | | elems: AdaptiveMap, |
185 | | cache: Cell<(u32, u64)>, |
186 | | } |
187 | | |
188 | | const BITS_PER_WORD: usize = 64; |
189 | | |
190 | | impl IndexSet { |
191 | 0 | pub fn new() -> Self { |
192 | 0 | Self { |
193 | 0 | elems: AdaptiveMap::new(), |
194 | 0 | cache: Cell::new((INVALID, 0)), |
195 | 0 | } |
196 | 0 | } |
197 | | |
198 | | #[inline(always)] |
199 | 0 | fn elem(&mut self, bit_index: usize) -> &mut u64 { |
200 | 0 | let word_index = (bit_index / BITS_PER_WORD) as u32; |
201 | 0 | if self.cache.get().0 == word_index { |
202 | 0 | self.cache.set((INVALID, 0)); |
203 | 0 | } |
204 | 0 | self.elems.get_or_insert(word_index) |
205 | 0 | } |
206 | | |
207 | | #[inline(always)] |
208 | 0 | fn maybe_elem_mut(&mut self, bit_index: usize) -> Option<&mut u64> { |
209 | 0 | let word_index = (bit_index / BITS_PER_WORD) as u32; |
210 | 0 | if self.cache.get().0 == word_index { |
211 | 0 | self.cache.set((INVALID, 0)); |
212 | 0 | } |
213 | 0 | self.elems.get_mut(word_index) |
214 | 0 | } |
215 | | |
216 | | #[inline(always)] |
217 | 0 | fn maybe_elem(&self, bit_index: usize) -> Option<u64> { |
218 | 0 | let word_index = (bit_index / BITS_PER_WORD) as u32; |
219 | 0 | if self.cache.get().0 == word_index { |
220 | 0 | Some(self.cache.get().1) |
221 | | } else { |
222 | 0 | self.elems.get(word_index) |
223 | | } |
224 | 0 | } |
225 | | |
226 | | #[inline(always)] |
227 | 0 | pub fn set(&mut self, idx: usize, val: bool) { |
228 | 0 | let bit = idx % BITS_PER_WORD; |
229 | 0 | if val { |
230 | 0 | *self.elem(idx) |= 1 << bit; |
231 | 0 | } else if let Some(word) = self.maybe_elem_mut(idx) { |
232 | 0 | *word &= !(1 << bit); |
233 | 0 | } |
234 | 0 | } |
235 | | |
236 | 0 | pub fn assign(&mut self, other: &Self) { |
237 | 0 | self.elems = other.elems.clone(); |
238 | 0 | self.cache = other.cache.clone(); |
239 | 0 | } |
240 | | |
241 | | #[inline(always)] |
242 | 0 | pub fn get(&self, idx: usize) -> bool { |
243 | 0 | let bit = idx % BITS_PER_WORD; |
244 | 0 | if let Some(word) = self.maybe_elem(idx) { |
245 | 0 | (word & (1 << bit)) != 0 |
246 | | } else { |
247 | 0 | false |
248 | | } |
249 | 0 | } |
250 | | |
251 | 0 | pub fn union_with(&mut self, other: &Self) -> bool { |
252 | 0 | let mut changed = 0; |
253 | 0 | for (word_idx, bits) in other.elems.iter() { |
254 | 0 | if bits == 0 { |
255 | 0 | continue; |
256 | 0 | } |
257 | 0 | let word_idx = word_idx as usize; |
258 | 0 | let self_word = self.elem(word_idx * BITS_PER_WORD); |
259 | 0 | changed |= bits & !*self_word; |
260 | 0 | *self_word |= bits; |
261 | | } |
262 | 0 | changed != 0 |
263 | 0 | } |
264 | | |
265 | 0 | pub fn iter<'a>(&'a self) -> impl Iterator<Item = usize> + 'a { |
266 | 0 | self.elems.iter().flat_map(|(word_idx, bits)| { |
267 | 0 | let word_idx = word_idx as usize; |
268 | 0 | SetBitsIter(bits).map(move |i| BITS_PER_WORD * word_idx + i) |
269 | 0 | }) |
270 | 0 | } |
271 | | |
272 | | /// Is the adaptive data structure in "small" mode? This is meant |
273 | | /// for testing assertions only. |
274 | 0 | pub(crate) fn is_small(&self) -> bool { |
275 | 0 | match &self.elems { |
276 | 0 | &AdaptiveMap::Small { .. } => true, |
277 | 0 | _ => false, |
278 | | } |
279 | 0 | } |
280 | | |
281 | | /// Is the set empty? |
282 | 0 | pub(crate) fn is_empty(&self) -> bool { |
283 | 0 | self.elems.is_empty() |
284 | 0 | } |
285 | | } |
286 | | |
287 | | pub struct SetBitsIter(u64); |
288 | | |
289 | | impl Iterator for SetBitsIter { |
290 | | type Item = usize; |
291 | | |
292 | | #[inline] |
293 | 0 | 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 | 0 | core::num::NonZeroU64::new(self.0).map(|nz| { |
298 | 0 | let bitidx = nz.trailing_zeros(); |
299 | 0 | self.0 &= self.0 - 1; // clear highest set bit |
300 | 0 | bitidx as usize |
301 | 0 | }) |
302 | 0 | } |
303 | | } |
304 | | |
305 | | impl core::fmt::Debug for IndexSet { |
306 | 0 | fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result { |
307 | 0 | let vals = self.iter().collect::<Vec<_>>(); |
308 | 0 | write!(f, "{:?}", vals) |
309 | 0 | } |
310 | | } |
311 | | |
312 | | #[cfg(test)] |
313 | | mod 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 | | } |