pico_ecs/
query.rs

1// ============================================================================
2// QUERY SYSTEM
3// ============================================================================
4// Provides ECS query functionality to iterate over entities with specific components
5
6use super::entity::Entity;
7use super::component_storage::ComponentStorage;
8use super::component::Component;
9
10/// Query for entities with a single component type
11pub struct Query1<'a, T: Component> {
12    storage: &'a ComponentStorage<T>,
13    entities: &'a [Entity],
14    current: usize,
15}
16
17impl<'a, T: Component> Query1<'a, T> {
18    pub fn new(storage: &'a ComponentStorage<T>, entities: &'a [Entity]) -> Self {
19        Self {
20            storage,
21            entities,
22            current: 0,
23        }
24    }
25}
26
27impl<'a, T: Component> Iterator for Query1<'a, T> {
28    type Item = (Entity, &'a T);
29    
30    fn next(&mut self) -> Option<Self::Item> {
31        while self.current < self.entities.len() {
32            let entity = self.entities[self.current];
33            self.current += 1;
34            
35            if let Some(component) = self.storage.get(entity) {
36                return Some((entity, component));
37            }
38        }
39        None
40    }
41}
42
43/// Query for entities with two component types
44pub struct Query2<'a, T1: Component, T2: Component> {
45    storage1: &'a ComponentStorage<T1>,
46    storage2: &'a ComponentStorage<T2>,
47    entities: &'a [Entity],
48    current: usize,
49}
50
51impl<'a, T1: Component, T2: Component> Query2<'a, T1, T2> {
52    pub fn new(
53        storage1: &'a ComponentStorage<T1>,
54        storage2: &'a ComponentStorage<T2>,
55        entities: &'a [Entity]
56    ) -> Self {
57        Self {
58            storage1,
59            storage2,
60            entities,
61            current: 0,
62        }
63    }
64}
65
66impl<'a, T1: Component, T2: Component> Iterator for Query2<'a, T1, T2> {
67    type Item = (Entity, &'a T1, &'a T2);
68    
69    fn next(&mut self) -> Option<Self::Item> {
70        while self.current < self.entities.len() {
71            let entity = self.entities[self.current];
72            self.current += 1;
73            
74            if let (Some(c1), Some(c2)) = (
75                self.storage1.get(entity),
76                self.storage2.get(entity)
77            ) {
78                return Some((entity, c1, c2));
79            }
80        }
81        None
82    }
83}
84
85/// Query for entities with three component types
86pub struct Query3<'a, T1: Component, T2: Component, T3: Component> {
87    storage1: &'a ComponentStorage<T1>,
88    storage2: &'a ComponentStorage<T2>,
89    storage3: &'a ComponentStorage<T3>,
90    entities: &'a [Entity],
91    current: usize,
92}
93
94impl<'a, T1: Component, T2: Component, T3: Component> Query3<'a, T1, T2, T3> {
95    pub fn new(
96        storage1: &'a ComponentStorage<T1>,
97        storage2: &'a ComponentStorage<T2>,
98        storage3: &'a ComponentStorage<T3>,
99        entities: &'a [Entity]
100    ) -> Self {
101        Self {
102            storage1,
103            storage2,
104            storage3,
105            entities,
106            current: 0,
107        }
108    }
109}
110
111impl<'a, T1: Component, T2: Component, T3: Component> Iterator for Query3<'a, T1, T2, T3> {
112    type Item = (Entity, &'a T1, &'a T2, &'a T3);
113    
114    fn next(&mut self) -> Option<Self::Item> {
115        while self.current < self.entities.len() {
116            let entity = self.entities[self.current];
117            self.current += 1;
118            
119            if let (Some(c1), Some(c2), Some(c3)) = (
120                self.storage1.get(entity),
121                self.storage2.get(entity),
122                self.storage3.get(entity)
123            ) {
124                return Some((entity, c1, c2, c3));
125            }
126        }
127        None
128    }
129}
130
131/// Mutable query for entities with a single component type
132pub struct Query1Mut<'a, T: Component> {
133    storage: &'a mut ComponentStorage<T>,
134    entities: &'a [Entity],
135    current: usize,
136}
137
138impl<'a, T: Component> Query1Mut<'a, T> {
139    pub fn new(storage: &'a mut ComponentStorage<T>, entities: &'a [Entity]) -> Self {
140        Self {
141            storage,
142            entities,
143            current: 0,
144        }
145    }
146}
147
148impl<'a, T: Component> Iterator for Query1Mut<'a, T> {
149    type Item = (Entity, &'a mut T);
150    
151    fn next(&mut self) -> Option<Self::Item> {
152        while self.current < self.entities.len() {
153            let entity = self.entities[self.current];
154            self.current += 1;
155            
156            // SAFETY: We only yield one mutable reference at a time
157            // and never yield the same entity twice in one iteration
158            if let Some(component) = unsafe {
159                let ptr = self.storage.get_mut(entity)? as *mut T;
160                Some(&mut *ptr)
161            } {
162                return Some((entity, component));
163            }
164        }
165        None
166    }
167}
168
169/// Query: 1 mutable, 1 immutable component
170pub struct Query1Mut1<'a, T1: Component, T2: Component> {
171    storage1: &'a mut ComponentStorage<T1>,
172    storage2: &'a ComponentStorage<T2>,
173    entities: &'a [Entity],
174    current: usize,
175}
176
177impl<'a, T1: Component, T2: Component> Query1Mut1<'a, T1, T2> {
178    pub fn new(
179        storage1: &'a mut ComponentStorage<T1>,
180        storage2: &'a ComponentStorage<T2>,
181        entities: &'a [Entity]
182    ) -> Self {
183        Self { storage1, storage2, entities, current: 0 }
184    }
185}
186
187impl<'a, T1: Component, T2: Component> Iterator for Query1Mut1<'a, T1, T2> {
188    type Item = (Entity, &'a mut T1, &'a T2);
189    
190    fn next(&mut self) -> Option<Self::Item> {
191        while self.current < self.entities.len() {
192            let entity = self.entities[self.current];
193            self.current += 1;
194            
195            if let Some(c2) = self.storage2.get(entity) {
196                if let Some(c1) = unsafe {
197                    let ptr = self.storage1.get_mut(entity)? as *mut T1;
198                    Some(&mut *ptr)
199                } {
200                    return Some((entity, c1, c2));
201                }
202            }
203        }
204        None
205    }
206}
207
208/// Query: 1 mutable, 2 immutable components
209pub struct Query1Mut2<'a, T1: Component, T2: Component, T3: Component> {
210    storage1: &'a mut ComponentStorage<T1>,
211    storage2: &'a ComponentStorage<T2>,
212    storage3: &'a ComponentStorage<T3>,
213    entities: &'a [Entity],
214    current: usize,
215}
216
217impl<'a, T1: Component, T2: Component, T3: Component> Query1Mut2<'a, T1, T2, T3> {
218    pub fn new(
219        storage1: &'a mut ComponentStorage<T1>,
220        storage2: &'a ComponentStorage<T2>,
221        storage3: &'a ComponentStorage<T3>,
222        entities: &'a [Entity]
223    ) -> Self {
224        Self { storage1, storage2, storage3, entities, current: 0 }
225    }
226}
227
228impl<'a, T1: Component, T2: Component, T3: Component> Iterator for Query1Mut2<'a, T1, T2, T3> {
229    type Item = (Entity, &'a mut T1, &'a T2, &'a T3);
230    
231    fn next(&mut self) -> Option<Self::Item> {
232        while self.current < self.entities.len() {
233            let entity = self.entities[self.current];
234            self.current += 1;
235            
236            if let (Some(c2), Some(c3)) = (
237                self.storage2.get(entity),
238                self.storage3.get(entity)
239            ) {
240                if let Some(c1) = unsafe {
241                    let ptr = self.storage1.get_mut(entity)? as *mut T1;
242                    Some(&mut *ptr)
243                } {
244                    return Some((entity, c1, c2, c3));
245                }
246            }
247        }
248        None
249    }
250}
251
252/// Query: 1 mutable, 3 immutable components
253pub struct Query1Mut3<'a, T1: Component, T2: Component, T3: Component, T4: Component> {
254    storage1: &'a mut ComponentStorage<T1>,
255    storage2: &'a ComponentStorage<T2>,
256    storage3: &'a ComponentStorage<T3>,
257    storage4: &'a ComponentStorage<T4>,
258    entities: &'a [Entity],
259    current: usize,
260}
261
262impl<'a, T1: Component, T2: Component, T3: Component, T4: Component> Query1Mut3<'a, T1, T2, T3, T4> {
263    pub fn new(
264        storage1: &'a mut ComponentStorage<T1>,
265        storage2: &'a ComponentStorage<T2>,
266        storage3: &'a ComponentStorage<T3>,
267        storage4: &'a ComponentStorage<T4>,
268        entities: &'a [Entity]
269    ) -> Self {
270        Self { storage1, storage2, storage3, storage4, entities, current: 0 }
271    }
272}
273
274impl<'a, T1: Component, T2: Component, T3: Component, T4: Component> Iterator for Query1Mut3<'a, T1, T2, T3, T4> {
275    type Item = (Entity, &'a mut T1, &'a T2, &'a T3, &'a T4);
276    
277    fn next(&mut self) -> Option<Self::Item> {
278        while self.current < self.entities.len() {
279            let entity = self.entities[self.current];
280            self.current += 1;
281            
282            if let (Some(c2), Some(c3), Some(c4)) = (
283                self.storage2.get(entity),
284                self.storage3.get(entity),
285                self.storage4.get(entity)
286            ) {
287                if let Some(c1) = unsafe {
288                    let ptr = self.storage1.get_mut(entity)? as *mut T1;
289                    Some(&mut *ptr)
290                } {
291                    return Some((entity, c1, c2, c3, c4));
292                }
293            }
294        }
295        None
296    }
297}
298
299/// Query: 2 mutable, 1 immutable component
300pub struct Query2Mut1<'a, T1: Component, T2: Component, T3: Component> {
301    storage1: &'a mut ComponentStorage<T1>,
302    storage2: &'a mut ComponentStorage<T2>,
303    storage3: &'a ComponentStorage<T3>,
304    entities: &'a [Entity],
305    current: usize,
306}
307
308impl<'a, T1: Component, T2: Component, T3: Component> Query2Mut1<'a, T1, T2, T3> {
309    pub fn new(
310        storage1: &'a mut ComponentStorage<T1>,
311        storage2: &'a mut ComponentStorage<T2>,
312        storage3: &'a ComponentStorage<T3>,
313        entities: &'a [Entity]
314    ) -> Self {
315        Self { storage1, storage2, storage3, entities, current: 0 }
316    }
317}
318
319impl<'a, T1: Component, T2: Component, T3: Component> Iterator for Query2Mut1<'a, T1, T2, T3> {
320    type Item = (Entity, &'a mut T1, &'a mut T2, &'a T3);
321    
322    fn next(&mut self) -> Option<Self::Item> {
323        while self.current < self.entities.len() {
324            let entity = self.entities[self.current];
325            self.current += 1;
326            
327            if let Some(c3) = self.storage3.get(entity) {
328                if let (Some(c1), Some(c2)) = unsafe {
329                    (
330                        self.storage1.get_mut(entity).map(|p| &mut *(p as *mut T1)),
331                        self.storage2.get_mut(entity).map(|p| &mut *(p as *mut T2))
332                    )
333                } {
334                    return Some((entity, c1, c2, c3));
335                }
336            }
337        }
338        None
339    }
340}
341
342/// Query: 2 mutable components
343pub struct Query2Mut<'a, T1: Component, T2: Component> {
344    storage1: &'a mut ComponentStorage<T1>,
345    storage2: &'a mut ComponentStorage<T2>,
346    entities: &'a [Entity],
347    current: usize,
348}
349
350impl<'a, T1: Component, T2: Component> Query2Mut<'a, T1, T2> {
351    pub fn new(
352        storage1: &'a mut ComponentStorage<T1>,
353        storage2: &'a mut ComponentStorage<T2>,
354        entities: &'a [Entity]
355    ) -> Self {
356        Self { storage1, storage2, entities, current: 0 }
357    }
358}
359
360impl<'a, T1: Component, T2: Component> Iterator for Query2Mut<'a, T1, T2> {
361    type Item = (Entity, &'a mut T1, &'a mut T2);
362    
363    fn next(&mut self) -> Option<Self::Item> {
364        while self.current < self.entities.len() {
365            let entity = self.entities[self.current];
366            self.current += 1;
367            
368            if let (Some(c1), Some(c2)) = unsafe {
369                (
370                    self.storage1.get_mut(entity).map(|p| &mut *(p as *mut T1)),
371                    self.storage2.get_mut(entity).map(|p| &mut *(p as *mut T2))
372                )
373            } {
374                return Some((entity, c1, c2));
375            }
376        }
377        None
378    }
379}
380
381/// Query: 3 mutable components
382pub struct Query3Mut<'a, T1: Component, T2: Component, T3: Component> {
383    storage1: &'a mut ComponentStorage<T1>,
384    storage2: &'a mut ComponentStorage<T2>,
385    storage3: &'a mut ComponentStorage<T3>,
386    entities: &'a [Entity],
387    current: usize,
388}
389
390impl<'a, T1: Component, T2: Component, T3: Component> Query3Mut<'a, T1, T2, T3> {
391    pub fn new(
392        storage1: &'a mut ComponentStorage<T1>,
393        storage2: &'a mut ComponentStorage<T2>,
394        storage3: &'a mut ComponentStorage<T3>,
395        entities: &'a [Entity]
396    ) -> Self {
397        Self { storage1, storage2, storage3, entities, current: 0 }
398    }
399}
400
401impl<'a, T1: Component, T2: Component, T3: Component> Iterator for Query3Mut<'a, T1, T2, T3> {
402    type Item = (Entity, &'a mut T1, &'a mut T2, &'a mut T3);
403    
404    fn next(&mut self) -> Option<Self::Item> {
405        while self.current < self.entities.len() {
406            let entity = self.entities[self.current];
407            self.current += 1;
408            
409            if let (Some(c1), Some(c2), Some(c3)) = unsafe {
410                (
411                    self.storage1.get_mut(entity).map(|p| &mut *(p as *mut T1)),
412                    self.storage2.get_mut(entity).map(|p| &mut *(p as *mut T2)),
413                    self.storage3.get_mut(entity).map(|p| &mut *(p as *mut T3))
414                )
415            } {
416                return Some((entity, c1, c2, c3));
417            }
418        }
419        None
420    }
421}
422
423/// Query for entities with four component types
424pub struct Query4<'a, T1: Component, T2: Component, T3: Component, T4: Component> {
425    storage1: &'a ComponentStorage<T1>,
426    storage2: &'a ComponentStorage<T2>,
427    storage3: &'a ComponentStorage<T3>,
428    storage4: &'a ComponentStorage<T4>,
429    entities: &'a [Entity],
430    current: usize,
431}
432
433impl<'a, T1: Component, T2: Component, T3: Component, T4: Component> Query4<'a, T1, T2, T3, T4> {
434    pub fn new(
435        storage1: &'a ComponentStorage<T1>,
436        storage2: &'a ComponentStorage<T2>,
437        storage3: &'a ComponentStorage<T3>,
438        storage4: &'a ComponentStorage<T4>,
439        entities: &'a [Entity]
440    ) -> Self {
441        Self {
442            storage1,
443            storage2,
444            storage3,
445            storage4,
446            entities,
447            current: 0,
448        }
449    }
450}
451
452impl<'a, T1: Component, T2: Component, T3: Component, T4: Component> Iterator for Query4<'a, T1, T2, T3, T4> {
453    type Item = (Entity, &'a T1, &'a T2, &'a T3, &'a T4);
454    
455    fn next(&mut self) -> Option<Self::Item> {
456        while self.current < self.entities.len() {
457            let entity = self.entities[self.current];
458            self.current += 1;
459            
460            if let (Some(c1), Some(c2), Some(c3), Some(c4)) = (
461                self.storage1.get(entity),
462                self.storage2.get(entity),
463                self.storage3.get(entity),
464                self.storage4.get(entity)
465            ) {
466                return Some((entity, c1, c2, c3, c4));
467            }
468        }
469        None
470    }
471}
472
473/// Query: 4 mutable components
474pub struct Query4Mut<'a, T1: Component, T2: Component, T3: Component, T4: Component> {
475    storage1: &'a mut ComponentStorage<T1>,
476    storage2: &'a mut ComponentStorage<T2>,
477    storage3: &'a mut ComponentStorage<T3>,
478    storage4: &'a mut ComponentStorage<T4>,
479    entities: &'a [Entity],
480    current: usize,
481}
482
483impl<'a, T1: Component, T2: Component, T3: Component, T4: Component> Query4Mut<'a, T1, T2, T3, T4> {
484    pub fn new(
485        storage1: &'a mut ComponentStorage<T1>,
486        storage2: &'a mut ComponentStorage<T2>,
487        storage3: &'a mut ComponentStorage<T3>,
488        storage4: &'a mut ComponentStorage<T4>,
489        entities: &'a [Entity]
490    ) -> Self {
491        Self { storage1, storage2, storage3, storage4, entities, current: 0 }
492    }
493}
494
495impl<'a, T1: Component, T2: Component, T3: Component, T4: Component> Iterator for Query4Mut<'a, T1, T2, T3, T4> {
496    type Item = (Entity, &'a mut T1, &'a mut T2, &'a mut T3, &'a mut T4);
497    
498    fn next(&mut self) -> Option<Self::Item> {
499        while self.current < self.entities.len() {
500            let entity = self.entities[self.current];
501            self.current += 1;
502            
503            if let (Some(c1), Some(c2), Some(c3), Some(c4)) = unsafe {
504                (
505                    self.storage1.get_mut(entity).map(|p| &mut *(p as *mut T1)),
506                    self.storage2.get_mut(entity).map(|p| &mut *(p as *mut T2)),
507                    self.storage3.get_mut(entity).map(|p| &mut *(p as *mut T3)),
508                    self.storage4.get_mut(entity).map(|p| &mut *(p as *mut T4))
509                )
510            } {
511                return Some((entity, c1, c2, c3, c4));
512            }
513        }
514        None
515    }
516}
pico_ecs/query.rs

pico_ecs/
query.rs
