flax/fetch/
satisfied.rs

1use alloc::vec::Vec;
2
3use crate::{archetype::Slice, Fetch, FetchItem};
4
5use super::{FetchAccessData, FmtQuery, PreparedFetch};
6
7/// Yields true iff `F` would match the query
8pub struct Satisfied<F>(pub F);
9
10impl<'q, F: FetchItem<'q>> FetchItem<'q> for Satisfied<F> {
11    type Item = bool;
12}
13
14impl<'w, F: Fetch<'w>> Fetch<'w> for Satisfied<F> {
15    const MUTABLE: bool = false;
16
17    type Prepared = PreparedSatisfied<F::Prepared>;
18
19    fn prepare(&'w self, data: super::FetchPrepareData<'w>) -> Option<Self::Prepared> {
20        if self.0.filter_arch(data.into()) {
21            Some(PreparedSatisfied(self.0.prepare(data)))
22        } else {
23            Some(PreparedSatisfied(None))
24        }
25    }
26
27    fn filter_arch(&self, _: FetchAccessData) -> bool {
28        true
29    }
30
31    fn describe(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
32        write!(f, "satisfied {:?}", FmtQuery(&self.0))
33    }
34
35    fn access(&self, _: super::FetchAccessData, _: &mut Vec<crate::system::Access>) {}
36}
37
38#[doc(hidden)]
39pub struct PreparedSatisfied<F>(Option<F>);
40
41impl<'q, F: PreparedFetch<'q>> PreparedFetch<'q> for PreparedSatisfied<F> {
42    type Item = bool;
43    type Chunk = bool;
44
45    const HAS_FILTER: bool = true;
46
47    unsafe fn create_chunk(&'q mut self, slots: Slice) -> Self::Chunk {
48        match &mut self.0 {
49            Some(f) => {
50                let res = f.filter_slots(slots);
51                !res.is_empty()
52            }
53            None => false,
54        }
55    }
56
57    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
58        *chunk
59    }
60
61    unsafe fn filter_slots(&mut self, slots: Slice) -> Slice {
62        match &mut self.0 {
63            Some(f) => {
64                let res = f.filter_slots(slots);
65
66                // Something was missed
67                if res.start != slots.start {
68                    // Catch the negative slice
69                    Slice::new(slots.start, res.start)
70                } else {
71                    res
72                }
73            }
74            None => slots,
75        }
76    }
77}
78
79#[cfg(test)]
80mod test {
81    use itertools::Itertools;
82    use pretty_assertions::assert_eq;
83
84    use crate::{component, components::name, Entity, FetchExt, Query, World};
85
86    component! {
87        a: i32,
88    }
89
90    #[test]
91    fn satisfied() {
92        let mut world = World::new();
93
94        ('a'..='c')
95            .map(|v| Entity::builder().set(name(), v.into()).spawn(&mut world))
96            .collect_vec();
97
98        ('d'..='f')
99            .map(|v| {
100                Entity::builder()
101                    .set(name(), v.into())
102                    .set(a(), 5)
103                    .spawn(&mut world)
104            })
105            .collect_vec();
106
107        let mut query = Query::new((name().cloned(), a().satisfied()));
108        assert_eq!(
109            query.collect_vec(&world),
110            [
111                ("a".into(), false),
112                ("b".into(), false),
113                ("c".into(), false),
114                ("d".into(), true),
115                ("e".into(), true),
116                ("f".into(), true),
117            ]
118        );
119    }
120
121    #[test]
122    fn satisfied_modified() {
123        let mut world = World::new();
124
125        ('a'..='c')
126            .map(|v| Entity::builder().set(name(), v.into()).spawn(&mut world))
127            .collect_vec();
128
129        let ids = ('d'..='f')
130            .map(|v| {
131                Entity::builder()
132                    .set(name(), v.into())
133                    .set(a(), 5)
134                    .spawn(&mut world)
135            })
136            .collect_vec();
137
138        let mut query = Query::new((name().cloned(), a().modified().satisfied()));
139
140        assert_eq!(
141            query.collect_vec(&world),
142            [
143                ("a".into(), false),
144                ("b".into(), false),
145                ("c".into(), false),
146                ("d".into(), true),
147                ("e".into(), true),
148                ("f".into(), true),
149            ]
150        );
151
152        assert_eq!(
153            query.collect_vec(&world),
154            [
155                ("a".into(), false),
156                ("b".into(), false),
157                ("c".into(), false),
158                ("d".into(), false),
159                ("e".into(), false),
160                ("f".into(), false),
161            ]
162        );
163
164        *world.get_mut(ids[1], a()).unwrap() = 5;
165
166        assert_eq!(
167            query.collect_vec(&world),
168            [
169                ("a".into(), false),
170                ("b".into(), false),
171                ("c".into(), false),
172                ("d".into(), false),
173                ("e".into(), true),
174                ("f".into(), false),
175            ]
176        );
177    }
178}