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