jecs/entities/
query.rs

1use super::Entities;
2use crate::entities::Component;
3use crate::errors::JellyEcsError;
4use eyre::Result;
5use std::any::{Any, TypeId};
6
7pub type QueryIndexes = Vec<usize>;
8pub type QueryComponents = Vec<Vec<Component>>;
9
10#[derive(Debug)]
11pub struct Query<'a> {
12    map: u32,
13    entities: &'a Entities,
14    type_ids: Vec<TypeId>,
15}
16
17impl<'a> Query<'a> {
18    pub fn new(entities: &'a Entities) -> Self {
19        Self {
20            entities,
21            map: 0,
22            type_ids: vec![],
23        }
24    }
25
26    pub fn with_component<T: Any>(&mut self) -> Result<&mut Self> {
27        let type_id = TypeId::of::<T>();
28        if let Some(bit_mask) = self.entities.get_bit_mask(&type_id) {
29            self.map |= bit_mask;
30            self.type_ids.push(type_id);
31        } else {
32            return Err(JellyEcsError::ComponentNotRegistered.into());
33        }
34
35        Ok(self)
36    }
37
38    pub fn run(&self) -> (QueryIndexes, QueryComponents) {
39        let indexes: Vec<usize> = self
40            .entities
41            .map
42            .iter()
43            .enumerate()
44            .filter_map(|(index, entity_map)| {
45                if entity_map & self.map == self.map {
46                    Some(index)
47                } else {
48                    None
49                }
50            })
51            .collect();
52
53        let mut result = vec![];
54
55        for type_id in &self.type_ids {
56            let entity_components = self.entities.components.get(type_id).unwrap();
57            let mut components_to_keep = vec![];
58            for index in &indexes {
59                components_to_keep.push(entity_components[*index].as_ref().unwrap().clone());
60            }
61
62            result.push(components_to_keep);
63        }
64
65        (indexes, result)
66    }
67}
68
69#[cfg(test)]
70mod test {
71    use super::*;
72
73    #[test]
74    fn query_mask_updating_with_component() -> Result<()> {
75        let mut entities = Entities::default();
76        entities.register_component::<u32>();
77        entities.register_component::<f32>();
78
79        let mut query = Query::new(&entities);
80        query.with_component::<u32>()?.with_component::<f32>()?;
81
82        assert_eq!(query.map, 3);
83        assert_eq!(TypeId::of::<u32>(), query.type_ids[0]);
84        assert_eq!(TypeId::of::<f32>(), query.type_ids[1]);
85        Ok(())
86    }
87
88    #[test]
89    #[allow(clippy::float_cmp)]
90    fn run_query() -> Result<()> {
91        let mut entities = Entities::default();
92        entities.register_component::<u32>();
93        entities.register_component::<f32>();
94
95        entities
96            .create_entity()
97            .with_component(10_u32)?
98            .with_component(16.0_f32)?;
99        entities.create_entity().with_component(20_u32)?;
100        entities.create_entity().with_component(32.0_f32)?;
101        entities
102            .create_entity()
103            .with_component(16_u32)?
104            .with_component(64.0_f32)?;
105
106        let mut query = Query::new(&entities);
107        let query_result = query
108            .with_component::<u32>()?
109            .with_component::<f32>()?
110            .run();
111        let u32s = &query_result.1[0];
112        let f32s = &query_result.1[1];
113        let indexes = &query_result.0;
114
115        assert!(u32s.len() == f32s.len() && u32s.len() == indexes.len());
116        assert_eq!(u32s.len(), 2);
117
118        let borrowed_first_u32 = u32s[0].borrow();
119        let first_u32 = borrowed_first_u32.downcast_ref::<u32>().unwrap();
120        assert_eq!(*first_u32, 10);
121
122        let borrowed_first_f32 = f32s[0].borrow();
123        let first_f32 = borrowed_first_f32.downcast_ref::<f32>().unwrap();
124        assert_eq!(*first_f32, 16.0);
125
126        let borrowed_second_u32 = u32s[1].borrow();
127        let second_u32 = borrowed_second_u32.downcast_ref::<u32>().unwrap();
128        assert_eq!(*second_u32, 16);
129
130        let borrowed_second_f32 = f32s[1].borrow();
131        let second_f32 = borrowed_second_f32.downcast_ref::<f32>().unwrap();
132        assert_eq!(*second_f32, 64.0);
133
134        assert_eq!(indexes[0], 0);
135        assert_eq!(indexes[1], 3);
136
137        Ok(())
138    }
139}