xref: /wasmtime-44.0.1/cranelift/entity/src/set.rs (revision bb8fa40e)
1 //! Densely numbered entity references as set keys.
2 
3 use crate::keys::Keys;
4 use crate::EntityRef;
5 use alloc::vec::Vec;
6 use core::marker::PhantomData;
7 
8 /// A set of `K` for densely indexed entity references.
9 ///
10 /// The `EntitySet` data structure uses the dense index space to implement a set with a bitvector.
11 /// Like `SecondaryMap`, an `EntitySet` is used to associate secondary information with entities.
12 #[derive(Debug, Clone)]
13 pub struct EntitySet<K>
14 where
15     K: EntityRef,
16 {
17     elems: Vec<u8>,
18     len: usize,
19     unused: PhantomData<K>,
20 }
21 
22 /// Shared `EntitySet` implementation for all value types.
23 impl<K> EntitySet<K>
24 where
25     K: EntityRef,
26 {
27     /// Create a new empty set.
28     pub fn new() -> Self {
29         Self {
30             elems: Vec::new(),
31             len: 0,
32             unused: PhantomData,
33         }
34     }
35 
36     /// Creates a new empty set with the specified capacity.
37     pub fn with_capacity(capacity: usize) -> Self {
38         Self {
39             elems: Vec::with_capacity((capacity + 7) / 8),
40             ..Self::new()
41         }
42     }
43 
44     /// Get the element at `k` if it exists.
45     pub fn contains(&self, k: K) -> bool {
46         let index = k.index();
47         if index < self.len {
48             (self.elems[index / 8] & (1 << (index % 8))) != 0
49         } else {
50             false
51         }
52     }
53 
54     /// Is this set completely empty?
55     pub fn is_empty(&self) -> bool {
56         if self.len != 0 {
57             false
58         } else {
59             self.elems.iter().all(|&e| e == 0)
60         }
61     }
62 
63     /// Returns the cardinality of the set.  More precisely, it returns the number of calls to
64     /// `insert` with different key values, that have happened since the the set was most recently
65     /// `clear`ed or created with `new`.
66     pub fn cardinality(&self) -> usize {
67         let mut n: usize = 0;
68         for byte_ix in 0..self.len / 8 {
69             n += self.elems[byte_ix].count_ones() as usize;
70         }
71         for bit_ix in (self.len / 8) * 8..self.len {
72             if (self.elems[bit_ix / 8] & (1 << (bit_ix % 8))) != 0 {
73                 n += 1;
74             }
75         }
76         n
77     }
78 
79     /// Remove all entries from this set.
80     pub fn clear(&mut self) {
81         self.len = 0;
82         self.elems.clear()
83     }
84 
85     /// Iterate over all the keys in this set.
86     pub fn keys(&self) -> Keys<K> {
87         Keys::with_len(self.len)
88     }
89 
90     /// Resize the set to have `n` entries by adding default entries as needed.
91     pub fn resize(&mut self, n: usize) {
92         self.elems.resize((n + 7) / 8, 0);
93         self.len = n
94     }
95 
96     /// Insert the element at `k`.
97     pub fn insert(&mut self, k: K) -> bool {
98         let index = k.index();
99         if index >= self.len {
100             self.resize(index + 1)
101         }
102         let result = !self.contains(k);
103         self.elems[index / 8] |= 1 << (index % 8);
104         result
105     }
106 
107     /// Removes and returns the entity from the set if it exists.
108     pub fn pop(&mut self) -> Option<K> {
109         if self.len == 0 {
110             return None;
111         }
112 
113         // Clear the last known entity in the list.
114         let last_index = self.len - 1;
115         self.elems[last_index / 8] &= !(1 << (last_index % 8));
116 
117         // Set the length to the next last stored entity or zero if we pop'ed
118         // the last entity.
119         self.len = self
120             .elems
121             .iter()
122             .enumerate()
123             .rev()
124             .find(|(_, &byte)| byte != 0)
125             // Map `i` from byte index to bit level index.
126             // `(i + 1) * 8` = Last bit in byte.
127             // `last - byte.leading_zeros()` = last set bit in byte.
128             // `as usize` won't ever truncate as the potential range is `0..=8`.
129             .map(|(i, byte)| ((i + 1) * 8) - byte.leading_zeros() as usize)
130             .unwrap_or(0);
131 
132         Some(K::new(last_index))
133     }
134 }
135 
136 #[cfg(test)]
137 mod tests {
138     use super::*;
139     use core::u32;
140 
141     // `EntityRef` impl for testing.
142     #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
143     struct E(u32);
144 
145     impl EntityRef for E {
146         fn new(i: usize) -> Self {
147             E(i as u32)
148         }
149         fn index(self) -> usize {
150             self.0 as usize
151         }
152     }
153 
154     #[test]
155     fn basic() {
156         let r0 = E(0);
157         let r1 = E(1);
158         let r2 = E(2);
159         let mut m = EntitySet::new();
160 
161         let v: Vec<E> = m.keys().collect();
162         assert_eq!(v, []);
163         assert!(m.is_empty());
164 
165         m.insert(r2);
166         m.insert(r1);
167 
168         assert!(!m.contains(r0));
169         assert!(m.contains(r1));
170         assert!(m.contains(r2));
171         assert!(!m.contains(E(3)));
172         assert!(!m.is_empty());
173 
174         let v: Vec<E> = m.keys().collect();
175         assert_eq!(v, [r0, r1, r2]);
176 
177         m.resize(20);
178         assert!(!m.contains(E(3)));
179         assert!(!m.contains(E(4)));
180         assert!(!m.contains(E(8)));
181         assert!(!m.contains(E(15)));
182         assert!(!m.contains(E(19)));
183 
184         m.insert(E(8));
185         m.insert(E(15));
186         assert!(!m.contains(E(3)));
187         assert!(!m.contains(E(4)));
188         assert!(m.contains(E(8)));
189         assert!(!m.contains(E(9)));
190         assert!(!m.contains(E(14)));
191         assert!(m.contains(E(15)));
192         assert!(!m.contains(E(16)));
193         assert!(!m.contains(E(19)));
194         assert!(!m.contains(E(20)));
195         assert!(!m.contains(E(u32::MAX)));
196 
197         m.clear();
198         assert!(m.is_empty());
199     }
200 
201     #[test]
202     fn pop_ordered() {
203         let r0 = E(0);
204         let r1 = E(1);
205         let r2 = E(2);
206         let mut m = EntitySet::new();
207         m.insert(r0);
208         m.insert(r1);
209         m.insert(r2);
210 
211         assert_eq!(r2, m.pop().unwrap());
212         assert_eq!(r1, m.pop().unwrap());
213         assert_eq!(r0, m.pop().unwrap());
214         assert!(m.pop().is_none());
215         assert!(m.pop().is_none());
216     }
217 
218     #[test]
219     fn pop_unordered() {
220         let mut ebbs = [
221             E(0),
222             E(1),
223             E(6),
224             E(7),
225             E(5),
226             E(9),
227             E(10),
228             E(2),
229             E(3),
230             E(11),
231             E(12),
232         ];
233 
234         let mut m = EntitySet::new();
235         for &ebb in &ebbs {
236             m.insert(ebb);
237         }
238         assert_eq!(m.len, 13);
239         ebbs.sort();
240 
241         for &ebb in ebbs.iter().rev() {
242             assert_eq!(ebb, m.pop().unwrap());
243         }
244 
245         assert!(m.is_empty());
246     }
247 }
248