cranelift_entity/
set.rs

1//! Densely numbered entity references as set keys.
2
3use crate::keys::Keys;
4use crate::EntityRef;
5use core::marker::PhantomData;
6use cranelift_bitset::CompoundBitSet;
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, PartialEq, Eq)]
13#[cfg_attr(
14    feature = "enable-serde",
15    derive(serde_derive::Serialize, serde_derive::Deserialize)
16)]
17pub struct EntitySet<K>
18where
19    K: EntityRef,
20{
21    bitset: CompoundBitSet,
22    unused: PhantomData<K>,
23}
24
25impl<K: EntityRef> Default for EntitySet<K> {
26    fn default() -> Self {
27        Self {
28            bitset: CompoundBitSet::default(),
29            unused: PhantomData,
30        }
31    }
32}
33
34impl<K: EntityRef> Extend<K> for EntitySet<K> {
35    fn extend<T: IntoIterator<Item = K>>(&mut self, iter: T) {
36        for k in iter {
37            self.insert(k);
38        }
39    }
40}
41
42/// Shared `EntitySet` implementation for all value types.
43impl<K> EntitySet<K>
44where
45    K: EntityRef,
46{
47    /// Create a new empty set.
48    pub fn new() -> Self {
49        Self::default()
50    }
51
52    /// Creates a new empty set with the specified capacity.
53    pub fn with_capacity(capacity: usize) -> Self {
54        Self {
55            bitset: CompoundBitSet::with_capacity(capacity),
56            unused: PhantomData,
57        }
58    }
59
60    /// Ensure that the set has enough capacity to hold `capacity` total
61    /// elements.
62    pub fn ensure_capacity(&mut self, capacity: usize) {
63        self.bitset.ensure_capacity(capacity);
64    }
65
66    /// Get the element at `k` if it exists.
67    pub fn contains(&self, k: K) -> bool {
68        let index = k.index();
69        self.bitset.contains(index)
70    }
71
72    /// Is this set completely empty?
73    pub fn is_empty(&self) -> bool {
74        self.bitset.is_empty()
75    }
76
77    /// Remove all entries from this set.
78    pub fn clear(&mut self) {
79        self.bitset.clear()
80    }
81
82    /// Iterate over all the keys in this set.
83    pub fn keys(&self) -> Keys<K> {
84        Keys::with_len(self.bitset.max().map_or(0, |x| x + 1))
85    }
86
87    /// Insert the element at `k`.
88    pub fn insert(&mut self, k: K) -> bool {
89        let index = k.index();
90        self.bitset.insert(index)
91    }
92
93    /// Removes and returns the entity from the set if it exists.
94    pub fn pop(&mut self) -> Option<K> {
95        let index = self.bitset.pop()?;
96        Some(K::new(index))
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use alloc::vec::Vec;
104    use core::u32;
105
106    // `EntityRef` impl for testing.
107    #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
108    struct E(u32);
109
110    impl EntityRef for E {
111        fn new(i: usize) -> Self {
112            E(i as u32)
113        }
114        fn index(self) -> usize {
115            self.0 as usize
116        }
117    }
118
119    #[test]
120    fn basic() {
121        let r0 = E(0);
122        let r1 = E(1);
123        let r2 = E(2);
124        let mut m = EntitySet::new();
125
126        let v: Vec<E> = m.keys().collect();
127        assert_eq!(v, []);
128        assert!(m.is_empty());
129
130        m.insert(r2);
131        m.insert(r1);
132
133        assert!(!m.contains(r0));
134        assert!(m.contains(r1));
135        assert!(m.contains(r2));
136        assert!(!m.contains(E(3)));
137        assert!(!m.is_empty());
138
139        let v: Vec<E> = m.keys().collect();
140        assert_eq!(v, [r0, r1, r2]);
141
142        assert!(!m.contains(E(3)));
143        assert!(!m.contains(E(4)));
144        assert!(!m.contains(E(8)));
145        assert!(!m.contains(E(15)));
146        assert!(!m.contains(E(19)));
147
148        m.insert(E(8));
149        m.insert(E(15));
150        assert!(!m.contains(E(3)));
151        assert!(!m.contains(E(4)));
152        assert!(m.contains(E(8)));
153        assert!(!m.contains(E(9)));
154        assert!(!m.contains(E(14)));
155        assert!(m.contains(E(15)));
156        assert!(!m.contains(E(16)));
157        assert!(!m.contains(E(19)));
158        assert!(!m.contains(E(20)));
159        assert!(!m.contains(E(u32::MAX)));
160
161        m.clear();
162        assert!(m.is_empty());
163    }
164
165    #[test]
166    fn pop_ordered() {
167        let r0 = E(0);
168        let r1 = E(1);
169        let r2 = E(2);
170        let mut m = EntitySet::new();
171        m.insert(r0);
172        m.insert(r1);
173        m.insert(r2);
174
175        assert_eq!(r2, m.pop().unwrap());
176        assert_eq!(r1, m.pop().unwrap());
177        assert_eq!(r0, m.pop().unwrap());
178        assert!(m.pop().is_none());
179        assert!(m.pop().is_none());
180    }
181
182    #[test]
183    fn pop_unordered() {
184        let mut blocks = [
185            E(0),
186            E(1),
187            E(6),
188            E(7),
189            E(5),
190            E(9),
191            E(10),
192            E(2),
193            E(3),
194            E(11),
195            E(12),
196        ];
197
198        let mut m = EntitySet::new();
199        for &block in &blocks {
200            m.insert(block);
201        }
202        assert_eq!(m.bitset.max(), Some(12));
203        blocks.sort();
204
205        for &block in blocks.iter().rev() {
206            assert_eq!(block, m.pop().unwrap());
207        }
208
209        assert!(m.is_empty());
210    }
211}