bones_ecs/components/
iterator.rs

1use std::rc::Rc;
2
3use crate::prelude::*;
4
5/// Read-only iterator over components matching a given bitset
6pub type ComponentBitsetIterator<'a, T> =
7    std::iter::Map<UntypedComponentBitsetIterator<'a>, for<'b> fn(SchemaRef<'b>) -> &'b T>;
8
9/// Read-only iterator over components matching a given bitset.
10/// Returns None for entities matching bitset but not in this ComponentStore.
11pub type ComponentBitsetOptionalIterator<'a, T> = std::iter::Map<
12    UntypedComponentOptionalBitsetIterator<'a>,
13    for<'b> fn(Option<SchemaRef<'b>>) -> Option<&'b T>,
14>;
15
16/// Mutable iterator over components matching a given bitset
17pub type ComponentBitsetIteratorMut<'a, T> = std::iter::Map<
18    UntypedComponentBitsetIteratorMut<'a>,
19    for<'b> fn(SchemaRefMut<'b>) -> &'b mut T,
20>;
21
22/// Mutable iterator over components matching a given bitset.
23/// Returns None for entities matching bitset but not in this ComponentStore.
24pub type ComponentBitsetOptionalIteratorMut<'a, T> = std::iter::Map<
25    UntypedComponentOptionalBitsetIteratorMut<'a>,
26    for<'b> fn(Option<SchemaRefMut<'b>>) -> Option<&'b mut T>,
27>;
28
29/// Iterates over components using a provided bitset. Each time the bitset has a 1 in index i, the
30/// iterator will fetch data from the storage at index i and return it.
31pub struct UntypedComponentBitsetIterator<'a> {
32    pub(crate) current_id: usize,
33    pub(crate) components: &'a UntypedComponentStore,
34    pub(crate) bitset: Rc<BitSetVec>,
35}
36
37impl<'a> Iterator for UntypedComponentBitsetIterator<'a> {
38    type Item = SchemaRef<'a>;
39    fn next(&mut self) -> Option<Self::Item> {
40        let max_id = self.components.max_id;
41        while self.current_id < max_id
42            && !(self.bitset.bit_test(self.current_id)
43                && self.components.bitset.bit_test(self.current_id))
44        {
45            self.current_id += 1;
46        }
47        let ret = if self.current_id < max_id {
48            // SAFE: Here we are just getting a pointer, not doing anything unsafe with it.
49            Some(unsafe {
50                SchemaRef::from_ptr_schema(
51                    self.components.storage.unchecked_idx(self.current_id),
52                    self.components.schema,
53                )
54            })
55        } else {
56            None
57        };
58        self.current_id += 1;
59        ret
60    }
61}
62
63/// Iterate over component store returning `Option<SchemaRef<'a>>`,
64/// filtered by bitset of iterator, but not bitset of own ComponentStore. Returns None on
65/// bitset entries that do not have this Component.
66pub struct UntypedComponentOptionalBitsetIterator<'a> {
67    /// The component bitset iterator.
68    pub inner: UntypedComponentBitsetIterator<'a>,
69    /// The number of components (enabled bits in the `inner` iterator).
70    pub components_count: usize,
71    /// The number of enabled bits in the input bitset.
72    pub query_count: usize,
73    /// The number of enabled bits discovered in the query bitset. Iteration is complete when this
74    /// reaches `query_count`.
75    pub found: usize,
76}
77
78impl<'a> Iterator for UntypedComponentOptionalBitsetIterator<'a> {
79    type Item = Option<SchemaRef<'a>>;
80    fn next(&mut self) -> Option<Self::Item> {
81        if self.found >= self.query_count {
82            // All enabled bits from the query bitset have been discovered. The rest is just zeros,
83            // we are done iterating.
84            return None;
85        }
86
87        if self.components_count == 0 {
88            // The component store is empty, no need to test all of the bits.
89            self.found += 1;
90            return Some(None);
91        }
92
93        // We stop iterating at bitset length, not component store length, as we want to iterate over
94        // whole bitset and return None for entities that don't have this optional component.
95        let max_id = self.inner.bitset.bit_len();
96        while self.inner.current_id < max_id && !self.inner.bitset.bit_test(self.inner.current_id) {
97            self.inner.current_id += 1;
98        }
99
100        let ret = if self.inner.components.bitset.bit_test(self.inner.current_id) {
101            self.found += 1;
102            // SAFE: Here we are just getting a pointer, not doing anything unsafe with it.
103            Some(Some(unsafe {
104                SchemaRef::from_ptr_schema(
105                    self.inner
106                        .components
107                        .storage
108                        .unchecked_idx(self.inner.current_id),
109                    self.inner.components.schema,
110                )
111            }))
112        } else {
113            // Component at current_id is not in store, however we are still iterating,
114            // later ids in self.bitset may have components in store.
115            self.found += 1;
116            Some(None)
117        };
118
119        self.inner.current_id += 1;
120
121        ret
122    }
123}
124
125/// Iterate mutably over component store returning `Option<SchemaRef<'a>>`,
126/// filtered by bitset of iterator, but not bitset of own ComponentStore. Returns None on
127/// bitset entries that do not have this Component.
128pub struct UntypedComponentOptionalBitsetIteratorMut<'a> {
129    /// The component bitset iterator.
130    pub inner: UntypedComponentBitsetIteratorMut<'a>,
131    /// The number of components (enabled bits in the `inner` iterator).
132    pub components_count: usize,
133    /// The number of enabled bits in the input bitset.
134    pub query_count: usize,
135    /// The number of enabled bits discovered in the query bitset. Iteration is complete when this
136    /// reaches `query_count`.
137    pub found: usize,
138}
139
140impl<'a> Iterator for UntypedComponentOptionalBitsetIteratorMut<'a> {
141    type Item = Option<SchemaRefMut<'a>>;
142    fn next(&mut self) -> Option<Self::Item> {
143        if self.found >= self.query_count {
144            // All enabled bits from the query bitset have been discovered. The rest is just zeros,
145            // we are done iterating.
146            return None;
147        }
148
149        if self.components_count == 0 {
150            // The component store is empty, no need to test all of the bits.
151            self.found += 1;
152            return Some(None);
153        }
154
155        // We do not stop iterating at component store length, as we want to iterate over
156        // whole bitset and return None for entities that don't have this optional component.
157        let max_id = self.inner.bitset.bit_len();
158        while self.inner.current_id < max_id && !self.inner.bitset.bit_test(self.inner.current_id) {
159            self.inner.current_id += 1;
160        }
161
162        let ret = if self.inner.components.bitset.bit_test(self.inner.current_id) {
163            self.found += 1;
164            // SAFE: Here we are just getting a pointer, not doing anything unsafe with it.
165            Some(Some(unsafe {
166                SchemaRefMut::from_ptr_schema(
167                    self.inner
168                        .components
169                        .storage
170                        .unchecked_idx(self.inner.current_id),
171                    self.inner.components.schema,
172                )
173            }))
174        } else {
175            // Component at current_id is not in store, however we are still iterating,
176            // later ids in self.bitset may have components in store.
177            self.found += 1;
178            Some(None)
179        };
180
181        self.inner.current_id += 1;
182
183        ret
184    }
185}
186
187/// Iterates over components using a provided bitset. Each time the bitset has a 1 in index i, the
188/// iterator will fetch data from the storage at index i.
189pub struct UntypedComponentBitsetIteratorMut<'a> {
190    pub(crate) current_id: usize,
191    pub(crate) components: &'a mut UntypedComponentStore,
192    pub(crate) bitset: Rc<BitSetVec>,
193}
194
195impl<'a> Iterator for UntypedComponentBitsetIteratorMut<'a> {
196    type Item = SchemaRefMut<'a>;
197    fn next(&mut self) -> Option<Self::Item> {
198        let max_id = self.components.max_id;
199        while self.current_id < max_id
200            && !(self.bitset.bit_test(self.current_id)
201                && self.components.bitset.bit_test(self.current_id))
202        {
203            self.current_id += 1;
204        }
205        let ret = if self.current_id < max_id {
206            // SAFE: We know that the index is within bounds, and we know that the pointer will be
207            // valid for the new lifetime.
208            Some(unsafe {
209                SchemaRefMut::from_ptr_schema(
210                    self.components.storage.unchecked_idx(self.current_id),
211                    self.components.schema,
212                )
213            })
214        } else {
215            None
216        };
217        self.current_id += 1;
218        ret
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    #![allow(non_snake_case)]
225
226    use super::*;
227
228    #[derive(Clone, HasSchema, Default)]
229    struct A;
230
231    #[derive(Clone, HasSchema, Default)]
232    struct B;
233
234    #[test]
235    fn iter_with_empty_bitset() {
236        let mut entities = Entities::default();
237        let e = entities.create();
238        let mut components = ComponentStore::<A>::default();
239
240        components.insert(e, A);
241
242        let bitset = Rc::new(BitSetVec::default());
243        assert_eq!(components.iter_with_bitset(bitset.clone()).count(), 0);
244        assert_eq!(components.iter_mut_with_bitset(bitset).count(), 0);
245    }
246
247    #[test]
248    /// Test that iterating with optional components does not filter entities.
249    fn iter_with_optional() {
250        // Initialize two total entities, both with B, one with A.
251        let mut entities = Entities::default();
252        let e1 = entities.create();
253        let e2 = entities.create();
254        let mut components_a = ComponentStore::<A>::default();
255        components_a.insert(e1, A);
256
257        let mut components_b = ComponentStore::<B>::default();
258        components_b.insert(e1, B);
259        components_b.insert(e2, B);
260
261        // Iterate over all entities, optionally retrieve A
262        {
263            let comp_a = Ref::new(&components_a);
264            let mut count_a = 0;
265            let mut count = 0;
266            for (_, a) in entities.iter_with(&Optional(&comp_a)) {
267                count += 1;
268                if a.is_some() {
269                    count_a += 1;
270                }
271            }
272            assert_eq!(count_a, 1);
273            assert_eq!(count, 2);
274        }
275        // Mutably Iterate over all entities, optionally retrieve A
276        {
277            let mut comp_a_mut = RefMut::new(&mut components_a);
278            let mut count_a = 0;
279            let mut count = 0;
280            for (_, a) in entities.iter_with(&mut OptionalMut(&mut comp_a_mut)) {
281                count += 1;
282                if a.is_some() {
283                    count_a += 1;
284                }
285            }
286            assert_eq!(count_a, 1);
287            assert_eq!(count, 2);
288        }
289
290        // Iterate over entities with B and optionaly retrieve A
291        {
292            let comp_a = Ref::new(&components_a);
293            let comp_b = Ref::new(&components_b);
294            let mut count_a = 0;
295            let mut count = 0;
296            for (_, (a, _b)) in entities.iter_with((&Optional(&comp_a), &comp_b)) {
297                count += 1;
298                if a.is_some() {
299                    count_a += 1;
300                }
301            }
302            assert_eq!(count_a, 1);
303            assert_eq!(count, 2);
304        }
305
306        // Iterate over entities with A, and optionally retrieve B
307        {
308            let comp_a = Ref::new(&components_a);
309            let comp_b = Ref::new(&components_b);
310            let mut count = 0;
311            for (_, (_a, b)) in entities.iter_with((&comp_a, &Optional(&comp_b))) {
312                count += 1;
313                assert!(b.is_some());
314            }
315            assert_eq!(count, 1);
316        }
317
318        // Make sure that entities with only optional components are still filtered by others,
319        // and not included in query.
320        //
321        // Case: 4 entities, we query over A and Optionally C, where entities have comps: 0:[AB],1:[B],2:[C],3:[A]
322        // Filtered by A, should iterate over entities 0 and 3. Verify that entitiy 2 with C is not included.
323        {
324            let e3 = entities.create();
325            let e4 = entities.create();
326            let mut components_c = ComponentStore::<A>::default();
327            components_c.insert(e3, A);
328            components_a.insert(e4, A);
329            let comp_a = Ref::new(&components_a);
330            let comp_c = Ref::new(&components_c);
331
332            let mut count = 0;
333            for (_, (_, c)) in entities.iter_with((&comp_a, &Optional(&comp_c))) {
334                count += 1;
335                // Should not iterate over entity with C, as it does not have A.
336                assert!(c.is_none());
337            }
338            // Expected two entities with A
339            assert_eq!(count, 2);
340        }
341    }
342
343    #[derive(Debug, Clone, Copy, PartialEq, Eq, HasSchema, Default)]
344    struct X(u32);
345
346    #[derive(Debug, Clone, Copy, PartialEq, Eq, HasSchema, Default)]
347    struct Y(u32);
348
349    fn entity(index: u32) -> Entity {
350        Entity::new(index, 0)
351    }
352
353    fn store<C: HasSchema>(entities: &[u32], ctor: fn(u32) -> C) -> ComponentStore<C> {
354        let mut store = ComponentStore::default();
355        for &i in entities {
356            store.insert(entity(i), ctor(i));
357        }
358        store
359    }
360
361    fn bitset(enabled: &[usize]) -> BitSetVec {
362        let mut bitset = BitSetVec::default();
363        for &i in enabled {
364            bitset.bit_set(i);
365        }
366        bitset
367    }
368
369    #[test]
370    fn get_single_with_bitset__multiple_required() {
371        {
372            let bitset = bitset(&[]);
373            let (store_x, store_y) = (store(&[], X), store(&[], Y));
374            let query = (&Ref::new(&store_x), &Ref::new(&store_y));
375
376            let result = query.get_single_with_bitset(Rc::new(bitset));
377
378            assert_eq!(result, Err(QuerySingleError::NoEntities));
379        }
380
381        {
382            let bitset = bitset(&[1]);
383            let store_x = store(&[1], X);
384            let store_y = store(&[1], Y);
385            let query = (&Ref::new(&store_x), &Ref::new(&store_y));
386
387            let result = query.get_single_with_bitset(Rc::new(bitset));
388
389            assert_eq!(result, Ok((&X(1), &Y(1))));
390        }
391
392        {
393            let bitset = bitset(&[1, 2]);
394            let store_x = store(&[1, 2], X);
395            let store_y = store(&[1, 2], Y);
396            let query = (&Ref::new(&store_x), &Ref::new(&store_y));
397
398            let result = query.get_single_with_bitset(Rc::new(bitset));
399
400            assert_eq!(result, Err(QuerySingleError::MultipleEntities));
401        }
402    }
403}
bones_ecs/components/iterator.rs

bones_ecs/components/
iterator.rs
