1 //! Densely numbered entity references as set keys. 2 3 use crate::keys::Keys; 4 use crate::EntityRef; 5 use alloc::vec::Vec; 6 use core::marker::PhantomData; 7 8 // How many bits are used to represent a single element in `EntitySet`. 9 const BITS: usize = core::mem::size_of::<usize>() * 8; 10 11 /// A set of `K` for densely indexed entity references. 12 /// 13 /// The `EntitySet` data structure uses the dense index space to implement a set with a bitvector. 14 /// Like `SecondaryMap`, an `EntitySet` is used to associate secondary information with entities. 15 #[derive(Debug, Clone)] 16 pub struct EntitySet<K> 17 where 18 K: EntityRef, 19 { 20 elems: Vec<usize>, 21 len: usize, 22 unused: PhantomData<K>, 23 } 24 25 impl<K: EntityRef> Default for EntitySet<K> { 26 fn default() -> Self { 27 Self { 28 elems: Vec::new(), 29 len: 0, 30 unused: PhantomData, 31 } 32 } 33 } 34 35 /// Shared `EntitySet` implementation for all value types. 36 impl<K> EntitySet<K> 37 where 38 K: EntityRef, 39 { 40 /// Create a new empty set. 41 pub fn new() -> Self { 42 Self::default() 43 } 44 45 /// Creates a new empty set with the specified capacity. 46 pub fn with_capacity(capacity: usize) -> Self { 47 Self { 48 elems: Vec::with_capacity((capacity + (BITS - 1)) / BITS), 49 ..Self::new() 50 } 51 } 52 53 /// Get the element at `k` if it exists. 54 pub fn contains(&self, k: K) -> bool { 55 let index = k.index(); 56 if index < self.len { 57 (self.elems[index / BITS] & (1 << (index % BITS))) != 0 58 } else { 59 false 60 } 61 } 62 63 /// Is this set completely empty? 64 pub fn is_empty(&self) -> bool { 65 if self.len != 0 { 66 false 67 } else { 68 self.elems.iter().all(|&e| e == 0) 69 } 70 } 71 72 /// Returns the cardinality of the set. More precisely, it returns the number of calls to 73 /// `insert` with different key values, that have happened since the the set was most recently 74 /// `clear`ed or created with `new`. 75 pub fn cardinality(&self) -> usize { 76 self.elems[..(self.len + (BITS - 1)) / BITS] 77 .iter() 78 .map(|x| x.count_ones() as usize) 79 .sum() 80 } 81 82 /// Remove all entries from this set. 83 pub fn clear(&mut self) { 84 self.len = 0; 85 self.elems.clear() 86 } 87 88 /// Iterate over all the keys in this set. 89 pub fn keys(&self) -> Keys<K> { 90 Keys::with_len(self.len) 91 } 92 93 /// Resize the set to have `n` entries by adding default entries as needed. 94 pub fn resize(&mut self, n: usize) { 95 self.elems.resize((n + (BITS - 1)) / BITS, 0); 96 self.len = n 97 } 98 99 /// Insert the element at `k`. 100 pub fn insert(&mut self, k: K) -> bool { 101 let index = k.index(); 102 if index >= self.len { 103 self.resize(index + 1) 104 } 105 let result = !self.contains(k); 106 self.elems[index / BITS] |= 1 << (index % BITS); 107 result 108 } 109 110 /// Removes and returns the entity from the set if it exists. 111 pub fn pop(&mut self) -> Option<K> { 112 if self.len == 0 { 113 return None; 114 } 115 116 // Clear the last known entity in the list. 117 let last_index = self.len - 1; 118 self.elems[last_index / BITS] &= !(1 << (last_index % BITS)); 119 120 // Set the length to the next last stored entity or zero if we pop'ed 121 // the last entity. 122 self.len = self 123 .elems 124 .iter() 125 .enumerate() 126 .rev() 127 .find(|(_, &elem)| elem != 0) 128 // Map `i` from `elem` index to bit level index. 129 // `(i + 1) * BITS` = Last bit in `elem`. 130 // `last - elem.leading_zeros()` = last set bit in `elem`. 131 // `as usize` won't ever truncate as the potential range is `0..=8`. 132 .map_or(0, |(i, elem)| { 133 ((i + 1) * BITS) - elem.leading_zeros() as usize 134 }); 135 136 Some(K::new(last_index)) 137 } 138 } 139 140 #[cfg(test)] 141 mod tests { 142 use super::*; 143 use core::u32; 144 145 // `EntityRef` impl for testing. 146 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] 147 struct E(u32); 148 149 impl EntityRef for E { 150 fn new(i: usize) -> Self { 151 E(i as u32) 152 } 153 fn index(self) -> usize { 154 self.0 as usize 155 } 156 } 157 158 #[test] 159 fn basic() { 160 let r0 = E(0); 161 let r1 = E(1); 162 let r2 = E(2); 163 let mut m = EntitySet::new(); 164 165 let v: Vec<E> = m.keys().collect(); 166 assert_eq!(v, []); 167 assert!(m.is_empty()); 168 169 m.insert(r2); 170 m.insert(r1); 171 172 assert!(!m.contains(r0)); 173 assert!(m.contains(r1)); 174 assert!(m.contains(r2)); 175 assert!(!m.contains(E(3))); 176 assert!(!m.is_empty()); 177 178 let v: Vec<E> = m.keys().collect(); 179 assert_eq!(v, [r0, r1, r2]); 180 181 m.resize(20); 182 assert!(!m.contains(E(3))); 183 assert!(!m.contains(E(4))); 184 assert!(!m.contains(E(8))); 185 assert!(!m.contains(E(15))); 186 assert!(!m.contains(E(19))); 187 188 m.insert(E(8)); 189 m.insert(E(15)); 190 assert!(!m.contains(E(3))); 191 assert!(!m.contains(E(4))); 192 assert!(m.contains(E(8))); 193 assert!(!m.contains(E(9))); 194 assert!(!m.contains(E(14))); 195 assert!(m.contains(E(15))); 196 assert!(!m.contains(E(16))); 197 assert!(!m.contains(E(19))); 198 assert!(!m.contains(E(20))); 199 assert!(!m.contains(E(u32::MAX))); 200 201 m.clear(); 202 assert!(m.is_empty()); 203 } 204 205 #[test] 206 fn pop_ordered() { 207 let r0 = E(0); 208 let r1 = E(1); 209 let r2 = E(2); 210 let mut m = EntitySet::new(); 211 m.insert(r0); 212 m.insert(r1); 213 m.insert(r2); 214 215 assert_eq!(r2, m.pop().unwrap()); 216 assert_eq!(r1, m.pop().unwrap()); 217 assert_eq!(r0, m.pop().unwrap()); 218 assert!(m.pop().is_none()); 219 assert!(m.pop().is_none()); 220 } 221 222 #[test] 223 fn pop_unordered() { 224 let mut blocks = [ 225 E(0), 226 E(1), 227 E(6), 228 E(7), 229 E(5), 230 E(9), 231 E(10), 232 E(2), 233 E(3), 234 E(11), 235 E(12), 236 ]; 237 238 let mut m = EntitySet::new(); 239 for &block in &blocks { 240 m.insert(block); 241 } 242 assert_eq!(m.len, 13); 243 blocks.sort(); 244 245 for &block in blocks.iter().rev() { 246 assert_eq!(block, m.pop().unwrap()); 247 } 248 249 assert!(m.is_empty()); 250 } 251 252 #[test] 253 fn cardinality() { 254 let mut m = EntitySet::new(); 255 256 m.insert(E(1)); 257 assert!(m.cardinality() == 1); 258 259 m.insert(E(0)); 260 assert!(m.cardinality() == 2); 261 262 m.insert(E(1)); 263 assert!(m.cardinality() == 2); 264 265 m.insert(E(BITS as u32 - 1)); 266 assert!(m.cardinality() == 3); 267 268 m.insert(E(BITS as u32)); 269 assert!(m.cardinality() == 4); 270 271 m.insert(E(BITS as u32 - 1)); 272 assert!(m.cardinality() == 4); 273 274 assert!(m.pop() == Some(E(BITS as u32))); 275 assert!(m.cardinality() == 3); 276 assert!(m.pop() == Some(E(BITS as u32 - 1))); 277 assert!(m.cardinality() == 2); 278 279 m.insert(E(100)); 280 assert!(m.cardinality() == 3); 281 282 assert!(m.pop() == Some(E(100))); 283 assert!(m.cardinality() == 2); 284 285 m.insert(E(100)); 286 m.insert(E(101)); 287 m.insert(E(102)); 288 assert!(m.cardinality() == 5); 289 } 290 } 291