1use crate::keys::Keys;
4use crate::EntityRef;
5use core::marker::PhantomData;
6use cranelift_bitset::CompoundBitSet;
7
8#[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
42impl<K> EntitySet<K>
44where
45 K: EntityRef,
46{
47 pub fn new() -> Self {
49 Self::default()
50 }
51
52 pub fn with_capacity(capacity: usize) -> Self {
54 Self {
55 bitset: CompoundBitSet::with_capacity(capacity),
56 unused: PhantomData,
57 }
58 }
59
60 pub fn ensure_capacity(&mut self, capacity: usize) {
63 self.bitset.ensure_capacity(capacity);
64 }
65
66 pub fn contains(&self, k: K) -> bool {
68 let index = k.index();
69 self.bitset.contains(index)
70 }
71
72 pub fn is_empty(&self) -> bool {
74 self.bitset.is_empty()
75 }
76
77 pub fn clear(&mut self) {
79 self.bitset.clear()
80 }
81
82 pub fn keys(&self) -> Keys<K> {
84 Keys::with_len(self.bitset.max().map_or(0, |x| x + 1))
85 }
86
87 pub fn insert(&mut self, k: K) -> bool {
89 let index = k.index();
90 self.bitset.insert(index)
91 }
92
93 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 #[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}