cranelift_entity/
set.rs

1//! Densely numbered entity references as set keys.
2
3use crate::EntityRef;
4use crate::keys::Keys;
5use core::fmt;
6use core::marker::PhantomData;
7use cranelift_bitset::CompoundBitSet;
8
9/// A set of `K` for densely indexed entity references.
10///
11/// The `EntitySet` data structure uses the dense index space to implement a set with a bitvector.
12/// Like `SecondaryMap`, an `EntitySet` is used to associate secondary information with entities.
13#[derive(Clone, PartialEq, Eq)]
14#[cfg_attr(
15    feature = "enable-serde",
16    derive(serde_derive::Serialize, serde_derive::Deserialize)
17)]
18pub struct EntitySet<K>
19where
20    K: EntityRef,
21{
22    bitset: CompoundBitSet,
23    unused: PhantomData<K>,
24}
25
26impl<K> fmt::Debug for EntitySet<K>
27where
28    K: fmt::Debug + EntityRef,
29{
30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31        f.debug_set().entries(self.keys()).finish()
32    }
33}
34
35impl<K: EntityRef> Default for EntitySet<K> {
36    fn default() -> Self {
37        Self {
38            bitset: CompoundBitSet::default(),
39            unused: PhantomData,
40        }
41    }
42}
43
44impl<K: EntityRef> Extend<K> for EntitySet<K> {
45    fn extend<T: IntoIterator<Item = K>>(&mut self, iter: T) {
46        for k in iter {
47            self.insert(k);
48        }
49    }
50}
51
52/// Shared `EntitySet` implementation for all value types.
53impl<K> EntitySet<K>
54where
55    K: EntityRef,
56{
57    /// Create a new empty set.
58    pub fn new() -> Self {
59        Self::default()
60    }
61
62    /// Creates a new empty set with the specified capacity.
63    pub fn with_capacity(capacity: usize) -> Self {
64        Self {
65            bitset: CompoundBitSet::with_capacity(capacity),
66            unused: PhantomData,
67        }
68    }
69
70    /// Ensure that the set has enough capacity to hold `capacity` total
71    /// elements.
72    pub fn ensure_capacity(&mut self, capacity: usize) {
73        self.bitset.ensure_capacity(capacity);
74    }
75
76    /// Get the element at `k` if it exists.
77    pub fn contains(&self, k: K) -> bool {
78        let index = k.index();
79        self.bitset.contains(index)
80    }
81
82    /// Is this set completely empty?
83    pub fn is_empty(&self) -> bool {
84        self.bitset.is_empty()
85    }
86
87    /// Remove all entries from this set.
88    pub fn clear(&mut self) {
89        self.bitset.clear()
90    }
91
92    /// Iterate over all the keys up to the maximum in this set.
93    ///
94    /// This will yield intermediate keys on the way up to the max key, even if
95    /// they are not contained within the set.
96    ///
97    /// ```
98    /// use cranelift_entity::{entity_impl, EntityRef, EntitySet};
99    ///
100    /// #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
101    /// struct Entity(u32);
102    /// entity_impl!(Entity);
103    ///
104    /// let mut set = EntitySet::new();
105    /// set.insert(Entity::new(2));
106    ///
107    /// let mut keys = set.keys();
108    /// assert_eq!(keys.next(), Some(Entity::new(0)));
109    /// assert_eq!(keys.next(), Some(Entity::new(1)));
110    /// assert_eq!(keys.next(), Some(Entity::new(2)));
111    /// assert!(keys.next().is_none());
112    /// ```
113    pub fn keys(&self) -> Keys<K> {
114        Keys::with_len(self.bitset.max().map_or(0, |x| x + 1))
115    }
116
117    /// Iterate over the elements of this set.
118    ///
119    /// ```
120    /// use cranelift_entity::{entity_impl, EntityRef, EntitySet};
121    ///
122    /// #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
123    /// struct Entity(u32);
124    /// entity_impl!(Entity);
125    ///
126    /// let mut set = EntitySet::new();
127    /// set.insert(Entity::new(2));
128    /// set.insert(Entity::new(3));
129    ///
130    /// let mut iter = set.iter();
131    /// assert_eq!(iter.next(), Some(Entity::new(2)));
132    /// assert_eq!(iter.next(), Some(Entity::new(3)));
133    /// assert!(iter.next().is_none());
134    /// ```
135    pub fn iter(&self) -> SetIter<'_, K> {
136        SetIter {
137            inner: self.bitset.iter(),
138            _phantom: PhantomData,
139        }
140    }
141
142    /// Insert the element at `k`.
143    ///
144    /// Returns `true` if `k` was not present in the set, i.e. this is a
145    /// newly-added element. Returns `false` otherwise.
146    pub fn insert(&mut self, k: K) -> bool {
147        let index = k.index();
148        self.bitset.insert(index)
149    }
150
151    /// Remove `k` from this bitset.
152    ///
153    /// Returns whether `k` was previously in this set or not.
154    pub fn remove(&mut self, k: K) -> bool {
155        let index = k.index();
156        self.bitset.remove(index)
157    }
158
159    /// Removes and returns the entity from the set if it exists.
160    pub fn pop(&mut self) -> Option<K> {
161        let index = self.bitset.pop()?;
162        Some(K::new(index))
163    }
164}
165
166/// An iterator over the elements in an `EntitySet`.
167pub struct SetIter<'a, K> {
168    inner: cranelift_bitset::compound::Iter<'a>,
169    _phantom: PhantomData<K>,
170}
171
172impl<K> Iterator for SetIter<'_, K>
173where
174    K: EntityRef,
175{
176    type Item = K;
177
178    #[inline]
179    fn next(&mut self) -> Option<Self::Item> {
180        let k = self.inner.next()?;
181        Some(K::new(k))
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188    use alloc::vec::Vec;
189    use core::u32;
190
191    // `EntityRef` impl for testing.
192    #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
193    struct E(u32);
194
195    impl EntityRef for E {
196        fn new(i: usize) -> Self {
197            E(i as u32)
198        }
199        fn index(self) -> usize {
200            self.0 as usize
201        }
202    }
203
204    #[test]
205    fn basic() {
206        let r0 = E(0);
207        let r1 = E(1);
208        let r2 = E(2);
209        let mut m = EntitySet::new();
210
211        let v: Vec<E> = m.keys().collect();
212        assert_eq!(v, []);
213        assert!(m.is_empty());
214
215        m.insert(r2);
216        m.insert(r1);
217
218        assert!(!m.contains(r0));
219        assert!(m.contains(r1));
220        assert!(m.contains(r2));
221        assert!(!m.contains(E(3)));
222        assert!(!m.is_empty());
223
224        let v: Vec<E> = m.keys().collect();
225        assert_eq!(v, [r0, r1, r2]);
226
227        assert!(!m.contains(E(3)));
228        assert!(!m.contains(E(4)));
229        assert!(!m.contains(E(8)));
230        assert!(!m.contains(E(15)));
231        assert!(!m.contains(E(19)));
232
233        m.insert(E(8));
234        m.insert(E(15));
235        assert!(!m.contains(E(3)));
236        assert!(!m.contains(E(4)));
237        assert!(m.contains(E(8)));
238        assert!(!m.contains(E(9)));
239        assert!(!m.contains(E(14)));
240        assert!(m.contains(E(15)));
241        assert!(!m.contains(E(16)));
242        assert!(!m.contains(E(19)));
243        assert!(!m.contains(E(20)));
244        assert!(!m.contains(E(u32::MAX)));
245
246        m.clear();
247        assert!(m.is_empty());
248    }
249
250    #[test]
251    fn pop_ordered() {
252        let r0 = E(0);
253        let r1 = E(1);
254        let r2 = E(2);
255        let mut m = EntitySet::new();
256        m.insert(r0);
257        m.insert(r1);
258        m.insert(r2);
259
260        assert_eq!(r2, m.pop().unwrap());
261        assert_eq!(r1, m.pop().unwrap());
262        assert_eq!(r0, m.pop().unwrap());
263        assert!(m.pop().is_none());
264        assert!(m.pop().is_none());
265    }
266
267    #[test]
268    fn pop_unordered() {
269        let mut blocks = [
270            E(0),
271            E(1),
272            E(6),
273            E(7),
274            E(5),
275            E(9),
276            E(10),
277            E(2),
278            E(3),
279            E(11),
280            E(12),
281        ];
282
283        let mut m = EntitySet::new();
284        for &block in &blocks {
285            m.insert(block);
286        }
287        assert_eq!(m.bitset.max(), Some(12));
288        blocks.sort();
289
290        for &block in blocks.iter().rev() {
291            assert_eq!(block, m.pop().unwrap());
292        }
293
294        assert!(m.is_empty());
295    }
296}
cranelift_entity/set.rs

cranelift_entity/
set.rs
