1 //! Densely numbered entity references as set keys. 2 3 use crate::keys::Keys; 4 use crate::EntityRef; 5 use core::marker::PhantomData; 6 use 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)] 13 pub struct EntitySet<K> 14 where 15 K: EntityRef, 16 { 17 bitset: CompoundBitSet, 18 unused: PhantomData<K>, 19 } 20 21 impl<K: EntityRef> Default for EntitySet<K> { 22 fn default() -> Self { 23 Self { 24 bitset: CompoundBitSet::default(), 25 unused: PhantomData, 26 } 27 } 28 } 29 30 impl<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. 39 impl<K> EntitySet<K> 40 where 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)] 97 mod 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 } 208