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}