ixa/people/
query.rs

1use crate::people::index::IndexValue;
2use crate::{Context, ContextPeopleExt, PersonProperty};
3use seq_macro::seq;
4use std::any::TypeId;
5
6/// Encapsulates a person query.
7///
8/// [`Context::query_people`] actually takes an instance of [`Query`], but because
9/// we implement Query for tuples of up to size 20, that's invisible
10/// to the caller. Do not use this trait directly.
11pub trait Query: Copy {
12    fn setup(&self, context: &Context);
13    fn get_query(&self) -> Vec<(TypeId, IndexValue)>;
14}
15
16impl Query for () {
17    fn setup(&self, _: &Context) {}
18
19    fn get_query(&self) -> Vec<(TypeId, IndexValue)> {
20        vec![]
21    }
22}
23
24// Implement the query version with one parameter.
25impl<T1: PersonProperty> Query for (T1, T1::Value) {
26    fn setup(&self, context: &Context) {
27        context.register_property::<T1>();
28    }
29
30    fn get_query(&self) -> Vec<(TypeId, IndexValue)> {
31        vec![(std::any::TypeId::of::<T1>(), IndexValue::compute(&self.1))]
32    }
33}
34
35// Implement the versions with 1..20 parameters.
36macro_rules! impl_query {
37    ($ct:expr) => {
38        seq!(N in 0..$ct {
39            impl<
40                #(
41                    T~N : PersonProperty,
42                )*
43            > Query for (
44                #(
45                    (T~N, T~N::Value),
46                )*
47            )
48            {
49                fn setup(&self, context: &Context) {
50                    #(
51                        context.register_property::<T~N>();
52                    )*
53                }
54
55                fn get_query(&self) -> Vec<(TypeId, IndexValue)> {
56                    let mut ordered_items = vec![
57                    #(
58                        (std::any::TypeId::of::<T~N>(), IndexValue::compute(&self.N.1)),
59                    )*
60                    ];
61                    ordered_items.sort_by(|a, b| a.0.cmp(&b.0));
62                    ordered_items
63                }
64            }
65        });
66    }
67}
68
69seq!(Z in 1..20 {
70    impl_query!(Z);
71});
72
73/// Helper utility for combining two queries, useful if you want
74/// to iteratively construct a query in multiple parts.
75///
76/// Example:
77/// ```
78/// use ixa::{define_person_property, people::QueryAnd, Context, ContextPeopleExt};
79/// define_person_property!(Age, u8);
80/// define_person_property!(Alive, bool);
81/// let context = Context::new();
82/// let q1 = (Age, 42);
83/// let q2 = (Alive, true);
84/// context.query_people(QueryAnd::new(q1, q2));
85/// ```
86
87#[derive(Copy, Clone)]
88pub struct QueryAnd<Q1, Q2>
89where
90    Q1: Query,
91    Q2: Query,
92{
93    queries: (Q1, Q2),
94}
95
96impl<Q1, Q2> QueryAnd<Q1, Q2>
97where
98    Q1: Query,
99    Q2: Query,
100{
101    pub fn new(q1: Q1, q2: Q2) -> Self {
102        Self { queries: (q1, q2) }
103    }
104}
105
106impl<Q1, Q2> Query for QueryAnd<Q1, Q2>
107where
108    Q1: Query,
109    Q2: Query,
110{
111    fn setup(&self, context: &Context) {
112        Q1::setup(&self.queries.0, context);
113        Q2::setup(&self.queries.1, context);
114    }
115
116    fn get_query(&self) -> Vec<(TypeId, IndexValue)> {
117        let mut query = Vec::new();
118        query.extend_from_slice(&self.queries.0.get_query());
119        query.extend_from_slice(&self.queries.1.get_query());
120        query.sort_by(|a, b| a.0.cmp(&b.0));
121        query
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use crate::people::PeoplePlugin;
128    use crate::people::{Query, QueryAnd};
129    use crate::{
130        define_derived_property, define_multi_property_index, define_person_property, Context,
131        ContextPeopleExt,
132    };
133    use serde_derive::Serialize;
134    use std::any::TypeId;
135
136    define_person_property!(Age, u8);
137    define_person_property!(County, u32);
138    define_person_property!(Height, u32);
139    define_derived_property!(AgeGroup, u8, [Age], |age| (age / 5));
140
141    #[derive(Serialize, Copy, Clone, PartialEq, Eq, Debug)]
142    pub enum RiskCategoryValue {
143        High,
144        Low,
145    }
146
147    define_person_property!(RiskCategory, RiskCategoryValue);
148
149    #[test]
150    fn query_people() {
151        let mut context = Context::new();
152        let _ = context
153            .add_person((RiskCategory, RiskCategoryValue::High))
154            .unwrap();
155
156        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
157        assert_eq!(people.len(), 1);
158    }
159
160    #[test]
161    fn query_people_empty() {
162        let context = Context::new();
163
164        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
165        assert_eq!(people.len(), 0);
166    }
167
168    #[test]
169    fn query_people_count() {
170        let mut context = Context::new();
171        let _ = context
172            .add_person((RiskCategory, RiskCategoryValue::High))
173            .unwrap();
174
175        assert_eq!(
176            context.query_people_count((RiskCategory, RiskCategoryValue::High)),
177            1
178        );
179    }
180
181    #[test]
182    fn query_people_count_empty() {
183        let context = Context::new();
184
185        assert_eq!(
186            context.query_people_count((RiskCategory, RiskCategoryValue::High)),
187            0
188        );
189    }
190
191    #[test]
192    fn query_people_macro_index_first() {
193        let mut context = Context::new();
194        let _ = context
195            .add_person((RiskCategory, RiskCategoryValue::High))
196            .unwrap();
197        context.index_property(RiskCategory);
198        assert!(property_is_indexed::<RiskCategory>(&context));
199        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
200        assert_eq!(people.len(), 1);
201    }
202
203    fn property_is_indexed<T: 'static>(context: &Context) -> bool {
204        context
205            .get_data_container(PeoplePlugin)
206            .unwrap()
207            .get_index_ref(TypeId::of::<T>())
208            .unwrap()
209            .lookup
210            .is_some()
211    }
212
213    #[test]
214    fn query_people_macro_index_second() {
215        let mut context = Context::new();
216        let _ = context.add_person((RiskCategory, RiskCategoryValue::High));
217        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
218        assert!(!property_is_indexed::<RiskCategory>(&context));
219        assert_eq!(people.len(), 1);
220        context.index_property(RiskCategory);
221        assert!(property_is_indexed::<RiskCategory>(&context));
222        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
223        assert_eq!(people.len(), 1);
224    }
225
226    #[test]
227    fn query_people_macro_change() {
228        let mut context = Context::new();
229        let person1 = context
230            .add_person((RiskCategory, RiskCategoryValue::High))
231            .unwrap();
232
233        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
234        assert_eq!(people.len(), 1);
235        let people = context.query_people((RiskCategory, RiskCategoryValue::Low));
236        assert_eq!(people.len(), 0);
237
238        context.set_person_property(person1, RiskCategory, RiskCategoryValue::Low);
239        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
240        assert_eq!(people.len(), 0);
241        let people = context.query_people((RiskCategory, RiskCategoryValue::Low));
242        assert_eq!(people.len(), 1);
243    }
244
245    #[test]
246    fn query_people_index_after_add() {
247        let mut context = Context::new();
248        let _ = context
249            .add_person((RiskCategory, RiskCategoryValue::High))
250            .unwrap();
251        context.index_property(RiskCategory);
252        assert!(property_is_indexed::<RiskCategory>(&context));
253        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
254        assert_eq!(people.len(), 1);
255    }
256
257    #[test]
258    fn query_people_add_after_index() {
259        let mut context = Context::new();
260        let _ = context
261            .add_person((RiskCategory, RiskCategoryValue::High))
262            .unwrap();
263        context.index_property(RiskCategory);
264        assert!(property_is_indexed::<RiskCategory>(&context));
265        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
266        assert_eq!(people.len(), 1);
267
268        let _ = context
269            .add_person((RiskCategory, RiskCategoryValue::High))
270            .unwrap();
271        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
272        assert_eq!(people.len(), 2);
273    }
274
275    #[test]
276    // This is safe because we reindex only when someone queries.
277    fn query_people_add_after_index_without_query() {
278        let mut context = Context::new();
279        let _ = context.add_person(()).unwrap();
280        context.index_property(RiskCategory);
281    }
282
283    #[test]
284    #[should_panic(expected = "Property not initialized")]
285    // This will panic when we query.
286    fn query_people_add_after_index_panic() {
287        let mut context = Context::new();
288        context.add_person(()).unwrap();
289        context.index_property(RiskCategory);
290        context.query_people((RiskCategory, RiskCategoryValue::High));
291    }
292
293    #[test]
294    fn query_people_cast_value() {
295        let mut context = Context::new();
296        let _ = context.add_person((Age, 42)).unwrap();
297
298        // Age is a u8, by default integer literals are i32; the macro should cast it.
299        let people = context.query_people((Age, 42));
300        assert_eq!(people.len(), 1);
301    }
302
303    #[test]
304    fn query_people_intersection() {
305        let mut context = Context::new();
306        let _ = context
307            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
308            .unwrap();
309        let _ = context
310            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::Low)))
311            .unwrap();
312        let _ = context
313            .add_person(((Age, 40), (RiskCategory, RiskCategoryValue::Low)))
314            .unwrap();
315
316        let people = context.query_people(((Age, 42), (RiskCategory, RiskCategoryValue::High)));
317        assert_eq!(people.len(), 1);
318    }
319
320    #[test]
321    fn query_people_intersection_non_macro() {
322        let mut context = Context::new();
323        let _ = context
324            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
325            .unwrap();
326        let _ = context
327            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::Low)))
328            .unwrap();
329        let _ = context
330            .add_person(((Age, 40), (RiskCategory, RiskCategoryValue::Low)))
331            .unwrap();
332
333        let people = context.query_people(((Age, 42), (RiskCategory, RiskCategoryValue::High)));
334        assert_eq!(people.len(), 1);
335    }
336
337    #[test]
338    fn query_people_intersection_one_indexed() {
339        let mut context = Context::new();
340        let _ = context
341            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
342            .unwrap();
343        let _ = context
344            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::Low)))
345            .unwrap();
346        let _ = context
347            .add_person(((Age, 40), (RiskCategory, RiskCategoryValue::Low)))
348            .unwrap();
349
350        context.index_property(Age);
351        let people = context.query_people(((Age, 42), (RiskCategory, RiskCategoryValue::High)));
352        assert_eq!(people.len(), 1);
353    }
354
355    #[test]
356    fn query_derived_prop() {
357        let mut context = Context::new();
358        define_derived_property!(Senior, bool, [Age], |age| age >= 65);
359
360        let person = context.add_person((Age, 64)).unwrap();
361        let _ = context.add_person((Age, 88)).unwrap();
362
363        // Age is a u8, by default integer literals are i32; the macro should cast it.
364        let not_seniors = context.query_people((Senior, false));
365        let seniors = context.query_people((Senior, true));
366        assert_eq!(seniors.len(), 1, "One senior");
367        assert_eq!(not_seniors.len(), 1, "One non-senior");
368
369        context.set_person_property(person, Age, 65);
370
371        let not_seniors = context.query_people((Senior, false));
372        let seniors = context.query_people((Senior, true));
373
374        assert_eq!(seniors.len(), 2, "Two seniors");
375        assert_eq!(not_seniors.len(), 0, "No non-seniors");
376    }
377
378    #[test]
379    fn query_derived_prop_with_index() {
380        let mut context = Context::new();
381        define_derived_property!(Senior, bool, [Age], |age| age >= 65);
382
383        context.index_property(Senior);
384        let person = context.add_person((Age, 64)).unwrap();
385        let _ = context.add_person((Age, 88)).unwrap();
386
387        // Age is a u8, by default integer literals are i32; the macro should cast it.
388        let not_seniors = context.query_people((Senior, false));
389        let seniors = context.query_people((Senior, true));
390        assert_eq!(seniors.len(), 1, "One senior");
391        assert_eq!(not_seniors.len(), 1, "One non-senior");
392
393        context.set_person_property(person, Age, 65);
394
395        let not_seniors = context.query_people((Senior, false));
396        let seniors = context.query_people((Senior, true));
397
398        assert_eq!(seniors.len(), 2, "Two seniors");
399        assert_eq!(not_seniors.len(), 0, "No non-seniors");
400    }
401
402    #[test]
403    fn query_derived_prop_with_optimized_index() {
404        let mut context = Context::new();
405        // create a 'regular' derived property
406        define_derived_property!(
407            Ach,
408            (u8, u32, u32),
409            [Age, County, Height],
410            |age, county, height| { (age, county, height) }
411        );
412
413        // create a multi-property index
414        define_multi_property_index!(Age, County, Height);
415        define_multi_property_index!(County, Height);
416
417        // add some people
418        let _person = context
419            .add_person(((Age, 64), (County, 2), (Height, 120)))
420            .unwrap();
421        let _ = context
422            .add_person(((Age, 88), (County, 2), (Height, 130)))
423            .unwrap();
424        let p2 = context
425            .add_person(((Age, 8), (County, 1), (Height, 140)))
426            .unwrap();
427        let p3 = context
428            .add_person(((Age, 28), (County, 1), (Height, 140)))
429            .unwrap();
430        let p4 = context
431            .add_person(((Age, 28), (County, 2), (Height, 160)))
432            .unwrap();
433        let p5 = context
434            .add_person(((Age, 28), (County, 2), (Height, 160)))
435            .unwrap();
436
437        // 'regular' derived property
438        let ach_people = context.query_people((Ach, (28, 2, 160)));
439        assert_eq!(ach_people.len(), 2, "Should have 2 matches");
440        assert!(ach_people.contains(&p4));
441        assert!(ach_people.contains(&p5));
442
443        // multi-property index
444        let age_county_height2 = context.query_people(((Age, 28), (County, 2), (Height, 160)));
445        assert_eq!(age_county_height2.len(), 2, "Should have 2 matches");
446        assert!(age_county_height2.contains(&p4));
447        assert!(age_county_height2.contains(&p5));
448
449        // multi-property index with different order
450        let age_county_height3 = context.query_people(((County, 2), (Height, 160), (Age, 28)));
451        assert_eq!(age_county_height3.len(), 2, "Should have 2 matches");
452        assert!(age_county_height3.contains(&p4));
453        assert!(age_county_height3.contains(&p5));
454
455        // multi-property index with different order
456        let age_county_height4 = context.query_people(((Height, 160), (County, 2), (Age, 28)));
457        assert_eq!(age_county_height4.len(), 2, "Should have 2 matches");
458        assert!(age_county_height4.contains(&p4));
459        assert!(age_county_height4.contains(&p5));
460
461        // multi-property index with different order and different value
462        let age_county_height5 = context.query_people(((Height, 140), (County, 1), (Age, 28)));
463        assert_eq!(age_county_height5.len(), 1, "Should have 1 matches");
464        assert!(age_county_height5.contains(&p3));
465
466        context.set_person_property(p2, Age, 28);
467        // multi-property index again after changing the value
468        let age_county_height5 = context.query_people(((Height, 140), (County, 1), (Age, 28)));
469        assert_eq!(age_county_height5.len(), 2, "Should have 2 matches");
470        assert!(age_county_height5.contains(&p2));
471        assert!(age_county_height5.contains(&p3));
472
473        let age_county_height5 = context.query_people(((Height, 140), (County, 1)));
474        assert_eq!(age_county_height5.len(), 2, "Should have 2 matches");
475        assert!(age_county_height5.contains(&p2));
476        assert!(age_county_height5.contains(&p3));
477    }
478
479    #[test]
480    fn query_and_returns_people() {
481        let mut context = Context::new();
482        context
483            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
484            .unwrap();
485
486        let q1 = (Age, 42);
487        let q2 = (RiskCategory, RiskCategoryValue::High);
488
489        let people = context.query_people(QueryAnd::new(q1, q2));
490        assert_eq!(people.len(), 1);
491    }
492
493    #[test]
494    fn query_and_conflicting() {
495        let mut context = Context::new();
496        context
497            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
498            .unwrap();
499
500        let q1 = (Age, 42);
501        let q2 = (Age, 64);
502
503        let people = context.query_people(QueryAnd::new(q1, q2));
504        assert_eq!(people.len(), 0);
505    }
506
507    fn query_and_copy_impl<Q: Query>(context: &Context, q: Q) {
508        for _ in 0..2 {
509            context.query_people(q);
510        }
511    }
512    #[test]
513    fn test_query_and_copy() {
514        let mut context = Context::new();
515        context
516            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
517            .unwrap();
518        query_and_copy_impl(
519            &context,
520            QueryAnd::new((Age, 42), (RiskCategory, RiskCategoryValue::High)),
521        );
522    }
523}