Skip to main content

cougr_core/query/
mod.rs

1mod helpers;
2mod state;
3#[cfg(test)]
4mod tests;
5
6pub use state::{SimpleQueryCache, SimpleQueryState};
7
8use crate::simple_world::EntityId as SimpleEntityId;
9use crate::simple_world::SimpleWorld;
10use soroban_sdk::{Env, Symbol, Vec};
11
12/// Which backing storage a `SimpleQuery` should consider.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum QueryStorage {
15    /// Only scan table-backed components, optimized for gameplay loops.
16    Table,
17    /// Include both table-backed and sparse components.
18    Any,
19}
20
21/// Query builder and executable query for `SimpleWorld`.
22///
23/// `SimpleQuery` is the preferred query surface for Soroban gameplay loops:
24/// it supports required, excluded, and optional-match components while
25/// selecting the narrowest indexed candidate set available.
26#[derive(Debug, Clone)]
27pub struct SimpleQuery {
28    required_components: Vec<Symbol>,
29    excluded_components: Vec<Symbol>,
30    any_components: Vec<Symbol>,
31    storage: QueryStorage,
32}
33
34impl SimpleQuery {
35    /// Create an empty query over table-backed components.
36    pub fn new(env: &Env) -> Self {
37        Self {
38            required_components: Vec::new(env),
39            excluded_components: Vec::new(env),
40            any_components: Vec::new(env),
41            storage: QueryStorage::Table,
42        }
43    }
44
45    /// Require the entity to have this component.
46    pub fn with_component(mut self, component_type: Symbol) -> Self {
47        self.required_components.push_back(component_type);
48        self
49    }
50
51    /// Require the entity to have all components from this slice.
52    pub fn with_components(mut self, component_types: &[Symbol]) -> Self {
53        for component_type in component_types {
54            self.required_components.push_back(component_type.clone());
55        }
56        self
57    }
58
59    /// Exclude entities with this component.
60    pub fn without_component(mut self, component_type: Symbol) -> Self {
61        self.excluded_components.push_back(component_type);
62        self
63    }
64
65    /// Exclude entities that have any component from this slice.
66    pub fn without_components(mut self, component_types: &[Symbol]) -> Self {
67        for component_type in component_types {
68            self.excluded_components.push_back(component_type.clone());
69        }
70        self
71    }
72
73    /// Require the entity to match at least one component from this set.
74    pub fn with_any_component(mut self, component_type: Symbol) -> Self {
75        self.any_components.push_back(component_type);
76        self
77    }
78
79    /// Require the entity to match at least one component from this slice.
80    pub fn with_any_components(mut self, component_types: &[Symbol]) -> Self {
81        for component_type in component_types {
82            self.any_components.push_back(component_type.clone());
83        }
84        self
85    }
86
87    /// Include both table and sparse storage during execution.
88    pub fn include_sparse(mut self) -> Self {
89        self.storage = QueryStorage::Any;
90        self
91    }
92
93    /// Returns whether the query has no constraints.
94    pub fn is_empty(&self) -> bool {
95        self.required_components.is_empty()
96            && self.excluded_components.is_empty()
97            && self.any_components.is_empty()
98    }
99
100    /// Returns the current storage mode for the query.
101    pub fn storage(&self) -> QueryStorage {
102        self.storage
103    }
104
105    /// Execute the query against a `SimpleWorld`.
106    pub fn execute(&self, world: &SimpleWorld, env: &Env) -> Vec<SimpleEntityId> {
107        let candidates = self.candidate_entities(world, env);
108        let mut results = Vec::new(env);
109
110        for i in 0..candidates.len() {
111            if let Some(entity_id) = candidates.get(i) {
112                if self.matches(world, entity_id) {
113                    results.push_back(entity_id);
114                }
115            }
116        }
117
118        results
119    }
120
121    /// Returns the number of entities that must be scanned before final filtering.
122    pub fn candidate_count(&self, world: &SimpleWorld, env: &Env) -> usize {
123        self.candidate_entities(world, env)
124            .len()
125            .try_into()
126            .unwrap()
127    }
128
129    fn candidate_entities(&self, world: &SimpleWorld, env: &Env) -> Vec<SimpleEntityId> {
130        let mut best: Option<Vec<SimpleEntityId>> = None;
131
132        for i in 0..self.required_components.len() {
133            if let Some(component_type) = self.required_components.get(i) {
134                let entities = self.entities_for_component(world, &component_type, env);
135                if best
136                    .as_ref()
137                    .map(|current| entities.len() < current.len())
138                    .unwrap_or(true)
139                {
140                    best = Some(entities);
141                }
142            }
143        }
144
145        if let Some(candidates) = best {
146            return candidates;
147        }
148
149        if !self.any_components.is_empty() {
150            let mut union = Vec::new(env);
151            for i in 0..self.any_components.len() {
152                if let Some(component_type) = self.any_components.get(i) {
153                    let entities = self.entities_for_component(world, &component_type, env);
154                    for j in 0..entities.len() {
155                        if let Some(entity_id) = entities.get(j) {
156                            if !helpers::contains_entity(&union, entity_id) {
157                                union.push_back(entity_id);
158                            }
159                        }
160                    }
161                }
162            }
163            return union;
164        }
165
166        let mut all_entities = Vec::new(env);
167        for entity_id in world.entity_components.keys().iter() {
168            all_entities.push_back(entity_id);
169        }
170        all_entities
171    }
172
173    fn entities_for_component(
174        &self,
175        world: &SimpleWorld,
176        component_type: &Symbol,
177        env: &Env,
178    ) -> Vec<SimpleEntityId> {
179        match self.storage {
180            QueryStorage::Table => world.get_table_entities_with_component(component_type, env),
181            QueryStorage::Any => world.get_all_entities_with_component(component_type, env),
182        }
183    }
184
185    fn matches(&self, world: &SimpleWorld, entity_id: SimpleEntityId) -> bool {
186        for i in 0..self.required_components.len() {
187            if let Some(component_type) = self.required_components.get(i) {
188                if !world.has_component(entity_id, &component_type) {
189                    return false;
190                }
191            }
192        }
193
194        for i in 0..self.excluded_components.len() {
195            if let Some(component_type) = self.excluded_components.get(i) {
196                if world.has_component(entity_id, &component_type) {
197                    return false;
198                }
199            }
200        }
201
202        if self.any_components.is_empty() {
203            return true;
204        }
205
206        for i in 0..self.any_components.len() {
207            if let Some(component_type) = self.any_components.get(i) {
208                if world.has_component(entity_id, &component_type) {
209                    return true;
210                }
211            }
212        }
213
214        false
215    }
216}
217
218/// Builder for `SimpleQuery`.
219pub struct SimpleQueryBuilder {
220    query: SimpleQuery,
221}
222
223impl SimpleQueryBuilder {
224    pub fn new(env: &Env) -> Self {
225        Self {
226            query: SimpleQuery::new(env),
227        }
228    }
229
230    pub fn with_component(mut self, component_type: Symbol) -> Self {
231        self.query = self.query.with_component(component_type);
232        self
233    }
234
235    pub fn with_components(mut self, component_types: &[Symbol]) -> Self {
236        self.query = self.query.with_components(component_types);
237        self
238    }
239
240    pub fn without_component(mut self, component_type: Symbol) -> Self {
241        self.query = self.query.without_component(component_type);
242        self
243    }
244
245    pub fn without_components(mut self, component_types: &[Symbol]) -> Self {
246        self.query = self.query.without_components(component_types);
247        self
248    }
249
250    pub fn with_any_component(mut self, component_type: Symbol) -> Self {
251        self.query = self.query.with_any_component(component_type);
252        self
253    }
254
255    pub fn with_any_components(mut self, component_types: &[Symbol]) -> Self {
256        self.query = self.query.with_any_components(component_types);
257        self
258    }
259
260    pub fn include_sparse(mut self) -> Self {
261        self.query = self.query.include_sparse();
262        self
263    }
264
265    pub fn build(self) -> SimpleQuery {
266        self.query
267    }
268
269    pub fn build_state(self, env: &Env) -> SimpleQueryState {
270        SimpleQueryState::new(self.query, env)
271    }
272}