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.
66#[derive(Deref, DerefMut)]
67pub struct UntypedComponentOptionalBitsetIterator<'a>(pub UntypedComponentBitsetIterator<'a>);
68
69impl<'a> Iterator for UntypedComponentOptionalBitsetIterator<'a> {
70    type Item = Option<SchemaRef<'a>>;
71    fn next(&mut self) -> Option<Self::Item> {
72        // We stop iterating at bitset length, not component store length, as we want to iterate over
73        // whole bitset and return None for entities that don't have this optional component.
74        let max_id = self.bitset.bit_len();
75        while self.current_id < max_id && !self.bitset.bit_test(self.current_id) {
76            self.current_id += 1;
77        }
78        let ret = if self.current_id < max_id {
79            // SAFE: Here we are just getting a pointer, not doing anything unsafe with it.
80            if self.components.bitset.bit_test(self.current_id) {
81                Some(Some(unsafe {
82                    SchemaRef::from_ptr_schema(
83                        self.components.storage.unchecked_idx(self.current_id),
84                        self.components.schema,
85                    )
86                }))
87            } else {
88                // Component at current_id is not in store, however we are still iterating,
89                // later ids in self.bitset may have components in store.
90                Some(None)
91            }
92        } else {
93            // Iterated through whole bitset
94            None
95        };
96        self.current_id += 1;
97        ret
98    }
99}
100
101/// Iterates over components using a provided bitset. Each time the bitset has a 1 in index i, the
102/// iterator will fetch data from the storage at index i.
103pub struct UntypedComponentBitsetIteratorMut<'a> {
104    pub(crate) current_id: usize,
105    pub(crate) components: &'a mut UntypedComponentStore,
106    pub(crate) bitset: Rc<BitSetVec>,
107}
108
109impl<'a> Iterator for UntypedComponentBitsetIteratorMut<'a> {
110    type Item = SchemaRefMut<'a>;
111    fn next(&mut self) -> Option<Self::Item> {
112        let max_id = self.components.max_id;
113        while self.current_id < max_id
114            && !(self.bitset.bit_test(self.current_id)
115                && self.components.bitset.bit_test(self.current_id))
116        {
117            self.current_id += 1;
118        }
119        let ret = if self.current_id < max_id {
120            // SAFE: We know that the index is within bounds, and we know that the pointer will be
121            // valid for the new lifetime.
122            Some(unsafe {
123                SchemaRefMut::from_ptr_schema(
124                    self.components.storage.unchecked_idx(self.current_id),
125                    self.components.schema,
126                )
127            })
128        } else {
129            None
130        };
131        self.current_id += 1;
132        ret
133    }
134}
135
136/// Iterate mutably over component store returning `Option<SchemaRef<'a>>`,
137/// filtered by bitset of iterator, but not bitset of own ComponentStore. Returns None on
138/// bitset entries that do not have this Component.
139#[derive(Deref, DerefMut)]
140pub struct UntypedComponentOptionalBitsetIteratorMut<'a>(pub UntypedComponentBitsetIteratorMut<'a>);
141
142impl<'a> Iterator for UntypedComponentOptionalBitsetIteratorMut<'a> {
143    type Item = Option<SchemaRefMut<'a>>;
144    fn next(&mut self) -> Option<Self::Item> {
145        // We do not stop iterating at component store length, as we want to iterate over
146        // whole bitset and return None for entities that don't have this optional component.
147        let max_id = self.bitset.bit_len();
148        while self.current_id < max_id && !self.bitset.bit_test(self.current_id) {
149            self.current_id += 1;
150        }
151        let ret = if self.current_id < max_id {
152            // SAFE: Here we are just getting a pointer, not doing anything unsafe with it.
153            if self.components.bitset.bit_test(self.current_id) {
154                Some(Some(unsafe {
155                    SchemaRefMut::from_ptr_schema(
156                        self.components.storage.unchecked_idx(self.current_id),
157                        self.components.schema,
158                    )
159                }))
160            } else {
161                // Component at current_id is not in store, however we are still iterating,
162                // later ids in self.bitset may have components in store.
163                Some(None)
164            }
165        } else {
166            // Iterated through whole bitset
167            None
168        };
169        self.current_id += 1;
170        ret
171    }
172}
173
174#[cfg(test)]
175mod test {
176    use super::*;
177
178    #[derive(Clone, HasSchema, Default)]
179    struct A;
180
181    #[derive(Clone, HasSchema, Default)]
182    struct B;
183
184    #[test]
185    fn iter_with_empty_bitset() {
186        let mut entities = Entities::default();
187        let e = entities.create();
188        let mut components = ComponentStore::<A>::default();
189
190        components.insert(e, A);
191
192        let bitset = Rc::new(BitSetVec::default());
193        assert_eq!(components.iter_with_bitset(bitset.clone()).count(), 0);
194        assert_eq!(components.iter_mut_with_bitset(bitset).count(), 0);
195    }
196
197    #[test]
198    /// Test that iterating with optional components does not filter entities.
199    fn iter_with_optional() {
200        // Initialize two total entities, both with B, one with A.
201        let mut entities = Entities::default();
202        let e1 = entities.create();
203        let e2 = entities.create();
204        let mut components_a = ComponentStore::<A>::default();
205        components_a.insert(e1, A);
206
207        let mut components_b = ComponentStore::<B>::default();
208        components_b.insert(e1, B);
209        components_b.insert(e2, B);
210
211        // Iterate over all entities, optionally retrieve A
212        {
213            let comp_a = Ref::new(&components_a);
214            let mut count_a = 0;
215            let mut count = 0;
216            for (_, a) in entities.iter_with(&Optional(&comp_a)) {
217                count += 1;
218                if a.is_some() {
219                    count_a += 1;
220                }
221            }
222            assert_eq!(count_a, 1);
223            assert_eq!(count, 2);
224        }
225        // Mutably Iterate over all entities, optionally retrieve A
226        {
227            let mut comp_a_mut = RefMut::new(&mut components_a);
228            let mut count_a = 0;
229            let mut count = 0;
230            for (_, a) in entities.iter_with(&mut OptionalMut(&mut comp_a_mut)) {
231                count += 1;
232                if a.is_some() {
233                    count_a += 1;
234                }
235            }
236            assert_eq!(count_a, 1);
237            assert_eq!(count, 2);
238        }
239
240        // Iterate over entities with B and optionaly retrieve A
241        {
242            let comp_a = Ref::new(&components_a);
243            let comp_b = Ref::new(&components_b);
244            let mut count_a = 0;
245            let mut count = 0;
246            for (_, (a, _b)) in entities.iter_with((&Optional(&comp_a), &comp_b)) {
247                count += 1;
248                if a.is_some() {
249                    count_a += 1;
250                }
251            }
252            assert_eq!(count_a, 1);
253            assert_eq!(count, 2);
254        }
255
256        // Iterate over entities with A, and optionally retrieve B
257        {
258            let comp_a = Ref::new(&components_a);
259            let comp_b = Ref::new(&components_b);
260            let mut count = 0;
261            for (_, (_a, b)) in entities.iter_with((&comp_a, &Optional(&comp_b))) {
262                count += 1;
263                assert!(b.is_some());
264            }
265            assert_eq!(count, 1);
266        }
267
268        // Make sure that entities with only optional components are still filtered by others,
269        // and not included in query.
270        //
271        // Case: 4 entities, we query over A and Optionally C, where entities have comps: 0:[AB],1:[B],2:[C],3:[A]
272        // Filtered by A, should iterate over entities 0 and 3. Verify that entitiy 2 with C is not included.
273        {
274            let e3 = entities.create();
275            let e4 = entities.create();
276            let mut components_c = ComponentStore::<A>::default();
277            components_c.insert(e3, A);
278            components_a.insert(e4, A);
279            let comp_a = Ref::new(&components_a);
280            let comp_c = Ref::new(&components_c);
281
282            let mut count = 0;
283            for (_, (_, c)) in entities.iter_with((&comp_a, &Optional(&comp_c))) {
284                count += 1;
285                // Should not iterate over entity with C, as it does not have A.
286                assert!(c.is_none());
287            }
288            // Expected two entities with A
289            assert_eq!(count, 2);
290        }
291    }
292}
bones_ecs/components/iterator.rs

bones_ecs/components/
iterator.rs
