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

cranelift_entity/
set.rs
