1use crate::keys::Keys;
4use crate::EntityRef;
5use core::marker::PhantomData;
6use cranelift_bitset::CompoundBitSet;
7
8#[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
38impl<K> EntitySet<K>
40where
41 K: EntityRef,
42{
43 pub fn new() -> Self {
45 Self::default()
46 }
47
48 pub fn with_capacity(capacity: usize) -> Self {
50 Self {
51 bitset: CompoundBitSet::with_capacity(capacity),
52 unused: PhantomData,
53 }
54 }
55
56 pub fn ensure_capacity(&mut self, capacity: usize) {
59 self.bitset.ensure_capacity(capacity);
60 }
61
62 pub fn contains(&self, k: K) -> bool {
64 let index = k.index();
65 self.bitset.contains(index)
66 }
67
68 pub fn is_empty(&self) -> bool {
70 self.bitset.is_empty()
71 }
72
73 pub fn clear(&mut self) {
75 self.bitset.clear()
76 }
77
78 pub fn keys(&self) -> Keys<K> {
80 Keys::with_len(self.bitset.max().map_or(0, |x| x + 1))
81 }
82
83 pub fn insert(&mut self, k: K) -> bool {
85 let index = k.index();
86 self.bitset.insert(index)
87 }
88
89 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 #[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}