cranelift_entity/
set.rs

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

cranelift_entity/
set.rs
