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
140    #[derive(Serialize, Copy, Clone, PartialEq, Eq, Debug)]
141    pub enum RiskCategoryValue {
142        High,
143        Low,
144    }
145
146    define_person_property!(RiskCategory, RiskCategoryValue);
147
148    #[test]
149    fn query_people() {
150        let mut context = Context::new();
151        let _ = context
152            .add_person((RiskCategory, RiskCategoryValue::High))
153            .unwrap();
154
155        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
156        assert_eq!(people.len(), 1);
157    }
158
159    #[test]
160    fn query_people_empty() {
161        let context = Context::new();
162
163        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
164        assert_eq!(people.len(), 0);
165    }
166
167    #[test]
168    fn query_people_count() {
169        let mut context = Context::new();
170        let _ = context
171            .add_person((RiskCategory, RiskCategoryValue::High))
172            .unwrap();
173
174        assert_eq!(
175            context.query_people_count((RiskCategory, RiskCategoryValue::High)),
176            1
177        );
178    }
179
180    #[test]
181    fn query_people_count_empty() {
182        let context = Context::new();
183
184        assert_eq!(
185            context.query_people_count((RiskCategory, RiskCategoryValue::High)),
186            0
187        );
188    }
189
190    #[test]
191    fn query_people_macro_index_first() {
192        let mut context = Context::new();
193        let _ = context
194            .add_person((RiskCategory, RiskCategoryValue::High))
195            .unwrap();
196        context.index_property(RiskCategory);
197        assert!(property_is_indexed::<RiskCategory>(&context));
198        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
199        assert_eq!(people.len(), 1);
200    }
201
202    fn property_is_indexed<T: 'static>(context: &Context) -> bool {
203        context
204            .get_data(PeoplePlugin)
205            .get_index_ref(TypeId::of::<T>())
206            .unwrap()
207            .lookup
208            .is_some()
209    }
210
211    #[test]
212    fn query_people_macro_index_second() {
213        let mut context = Context::new();
214        let _ = context.add_person((RiskCategory, RiskCategoryValue::High));
215        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
216        assert!(!property_is_indexed::<RiskCategory>(&context));
217        assert_eq!(people.len(), 1);
218        context.index_property(RiskCategory);
219        assert!(property_is_indexed::<RiskCategory>(&context));
220        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
221        assert_eq!(people.len(), 1);
222    }
223
224    #[test]
225    fn query_people_macro_change() {
226        let mut context = Context::new();
227        let person1 = context
228            .add_person((RiskCategory, RiskCategoryValue::High))
229            .unwrap();
230
231        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
232        assert_eq!(people.len(), 1);
233        let people = context.query_people((RiskCategory, RiskCategoryValue::Low));
234        assert_eq!(people.len(), 0);
235
236        context.set_person_property(person1, RiskCategory, RiskCategoryValue::Low);
237        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
238        assert_eq!(people.len(), 0);
239        let people = context.query_people((RiskCategory, RiskCategoryValue::Low));
240        assert_eq!(people.len(), 1);
241    }
242
243    #[test]
244    fn query_people_index_after_add() {
245        let mut context = Context::new();
246        let _ = context
247            .add_person((RiskCategory, RiskCategoryValue::High))
248            .unwrap();
249        context.index_property(RiskCategory);
250        assert!(property_is_indexed::<RiskCategory>(&context));
251        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
252        assert_eq!(people.len(), 1);
253    }
254
255    #[test]
256    fn query_people_add_after_index() {
257        let mut context = Context::new();
258        let _ = context
259            .add_person((RiskCategory, RiskCategoryValue::High))
260            .unwrap();
261        context.index_property(RiskCategory);
262        assert!(property_is_indexed::<RiskCategory>(&context));
263        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
264        assert_eq!(people.len(), 1);
265
266        let _ = context
267            .add_person((RiskCategory, RiskCategoryValue::High))
268            .unwrap();
269        let people = context.query_people((RiskCategory, RiskCategoryValue::High));
270        assert_eq!(people.len(), 2);
271    }
272
273    #[test]
274    // This is safe because we reindex only when someone queries.
275    fn query_people_add_after_index_without_query() {
276        let mut context = Context::new();
277        let _ = context.add_person(()).unwrap();
278        context.index_property(RiskCategory);
279    }
280
281    #[test]
282    #[should_panic(expected = "Property not initialized")]
283    // This will panic when we query.
284    fn query_people_add_after_index_panic() {
285        let mut context = Context::new();
286        context.add_person(()).unwrap();
287        context.index_property(RiskCategory);
288        context.query_people((RiskCategory, RiskCategoryValue::High));
289    }
290
291    #[test]
292    fn query_people_cast_value() {
293        let mut context = Context::new();
294        let _ = context.add_person((Age, 42)).unwrap();
295
296        // Age is a u8, by default integer literals are i32; the macro should cast it.
297        let people = context.query_people((Age, 42));
298        assert_eq!(people.len(), 1);
299    }
300
301    #[test]
302    fn query_people_intersection() {
303        let mut context = Context::new();
304        let _ = context
305            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
306            .unwrap();
307        let _ = context
308            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::Low)))
309            .unwrap();
310        let _ = context
311            .add_person(((Age, 40), (RiskCategory, RiskCategoryValue::Low)))
312            .unwrap();
313
314        let people = context.query_people(((Age, 42), (RiskCategory, RiskCategoryValue::High)));
315        assert_eq!(people.len(), 1);
316    }
317
318    #[test]
319    fn query_people_intersection_non_macro() {
320        let mut context = Context::new();
321        let _ = context
322            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
323            .unwrap();
324        let _ = context
325            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::Low)))
326            .unwrap();
327        let _ = context
328            .add_person(((Age, 40), (RiskCategory, RiskCategoryValue::Low)))
329            .unwrap();
330
331        let people = context.query_people(((Age, 42), (RiskCategory, RiskCategoryValue::High)));
332        assert_eq!(people.len(), 1);
333    }
334
335    #[test]
336    fn query_people_intersection_one_indexed() {
337        let mut context = Context::new();
338        let _ = context
339            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
340            .unwrap();
341        let _ = context
342            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::Low)))
343            .unwrap();
344        let _ = context
345            .add_person(((Age, 40), (RiskCategory, RiskCategoryValue::Low)))
346            .unwrap();
347
348        context.index_property(Age);
349        let people = context.query_people(((Age, 42), (RiskCategory, RiskCategoryValue::High)));
350        assert_eq!(people.len(), 1);
351    }
352
353    #[test]
354    fn query_derived_prop() {
355        let mut context = Context::new();
356        define_derived_property!(Senior, bool, [Age], |age| age >= 65);
357
358        let person = context.add_person((Age, 64)).unwrap();
359        let _ = context.add_person((Age, 88)).unwrap();
360
361        // Age is a u8, by default integer literals are i32; the macro should cast it.
362        let not_seniors = context.query_people((Senior, false));
363        let seniors = context.query_people((Senior, true));
364        assert_eq!(seniors.len(), 1, "One senior");
365        assert_eq!(not_seniors.len(), 1, "One non-senior");
366
367        context.set_person_property(person, Age, 65);
368
369        let not_seniors = context.query_people((Senior, false));
370        let seniors = context.query_people((Senior, true));
371
372        assert_eq!(seniors.len(), 2, "Two seniors");
373        assert_eq!(not_seniors.len(), 0, "No non-seniors");
374    }
375
376    #[test]
377    fn query_derived_prop_with_index() {
378        let mut context = Context::new();
379        define_derived_property!(Senior, bool, [Age], |age| age >= 65);
380
381        context.index_property(Senior);
382        let person = context.add_person((Age, 64)).unwrap();
383        let _ = context.add_person((Age, 88)).unwrap();
384
385        // Age is a u8, by default integer literals are i32; the macro should cast it.
386        let not_seniors = context.query_people((Senior, false));
387        let seniors = context.query_people((Senior, true));
388        assert_eq!(seniors.len(), 1, "One senior");
389        assert_eq!(not_seniors.len(), 1, "One non-senior");
390
391        context.set_person_property(person, Age, 65);
392
393        let not_seniors = context.query_people((Senior, false));
394        let seniors = context.query_people((Senior, true));
395
396        assert_eq!(seniors.len(), 2, "Two seniors");
397        assert_eq!(not_seniors.len(), 0, "No non-seniors");
398    }
399
400    #[test]
401    fn query_derived_prop_with_optimized_index() {
402        let mut context = Context::new();
403        // create a 'regular' derived property
404        define_derived_property!(
405            Ach,
406            (u8, u32, u32),
407            [Age, County, Height],
408            |age, county, height| { (age, county, height) }
409        );
410
411        // create a multi-property index
412        define_multi_property_index!(Age, County, Height);
413        define_multi_property_index!(County, Height);
414
415        // add some people
416        let _person = context
417            .add_person(((Age, 64), (County, 2), (Height, 120)))
418            .unwrap();
419        let _ = context
420            .add_person(((Age, 88), (County, 2), (Height, 130)))
421            .unwrap();
422        let p2 = context
423            .add_person(((Age, 8), (County, 1), (Height, 140)))
424            .unwrap();
425        let p3 = context
426            .add_person(((Age, 28), (County, 1), (Height, 140)))
427            .unwrap();
428        let p4 = context
429            .add_person(((Age, 28), (County, 2), (Height, 160)))
430            .unwrap();
431        let p5 = context
432            .add_person(((Age, 28), (County, 2), (Height, 160)))
433            .unwrap();
434
435        // 'regular' derived property
436        let ach_people = context.query_people((Ach, (28, 2, 160)));
437        assert_eq!(ach_people.len(), 2, "Should have 2 matches");
438        assert!(ach_people.contains(&p4));
439        assert!(ach_people.contains(&p5));
440
441        // multi-property index
442        let age_county_height2 = context.query_people(((Age, 28), (County, 2), (Height, 160)));
443        assert_eq!(age_county_height2.len(), 2, "Should have 2 matches");
444        assert!(age_county_height2.contains(&p4));
445        assert!(age_county_height2.contains(&p5));
446
447        // multi-property index with different order
448        let age_county_height3 = context.query_people(((County, 2), (Height, 160), (Age, 28)));
449        assert_eq!(age_county_height3.len(), 2, "Should have 2 matches");
450        assert!(age_county_height3.contains(&p4));
451        assert!(age_county_height3.contains(&p5));
452
453        // multi-property index with different order
454        let age_county_height4 = context.query_people(((Height, 160), (County, 2), (Age, 28)));
455        assert_eq!(age_county_height4.len(), 2, "Should have 2 matches");
456        assert!(age_county_height4.contains(&p4));
457        assert!(age_county_height4.contains(&p5));
458
459        // multi-property index with different order and different value
460        let age_county_height5 = context.query_people(((Height, 140), (County, 1), (Age, 28)));
461        assert_eq!(age_county_height5.len(), 1, "Should have 1 matches");
462        assert!(age_county_height5.contains(&p3));
463
464        context.set_person_property(p2, Age, 28);
465        // multi-property index again after changing the value
466        let age_county_height5 = context.query_people(((Height, 140), (County, 1), (Age, 28)));
467        assert_eq!(age_county_height5.len(), 2, "Should have 2 matches");
468        assert!(age_county_height5.contains(&p2));
469        assert!(age_county_height5.contains(&p3));
470
471        let age_county_height5 = context.query_people(((Height, 140), (County, 1)));
472        assert_eq!(age_county_height5.len(), 2, "Should have 2 matches");
473        assert!(age_county_height5.contains(&p2));
474        assert!(age_county_height5.contains(&p3));
475    }
476
477    #[test]
478    fn query_and_returns_people() {
479        let mut context = Context::new();
480        context
481            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
482            .unwrap();
483
484        let q1 = (Age, 42);
485        let q2 = (RiskCategory, RiskCategoryValue::High);
486
487        let people = context.query_people(QueryAnd::new(q1, q2));
488        assert_eq!(people.len(), 1);
489    }
490
491    #[test]
492    fn query_and_conflicting() {
493        let mut context = Context::new();
494        context
495            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
496            .unwrap();
497
498        let q1 = (Age, 42);
499        let q2 = (Age, 64);
500
501        let people = context.query_people(QueryAnd::new(q1, q2));
502        assert_eq!(people.len(), 0);
503    }
504
505    fn query_and_copy_impl<Q: Query>(context: &Context, q: Q) {
506        for _ in 0..2 {
507            context.query_people(q);
508        }
509    }
510    #[test]
511    fn test_query_and_copy() {
512        let mut context = Context::new();
513        context
514            .add_person(((Age, 42), (RiskCategory, RiskCategoryValue::High)))
515            .unwrap();
516        query_and_copy_impl(
517            &context,
518            QueryAnd::new((Age, 42), (RiskCategory, RiskCategoryValue::High)),
519        );
520    }
521}