bevy_either/
either_both.rs

1use super::*;
2
3/// A type that contains either the [first](EitherBoth::Left) type, [second](EitherBoth::Right)
4/// type, or [both](EitherBoth::Both).
5#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
6pub enum EitherBoth<T, U> {
7    Left(T),
8    Right(U),
9    Both(T, U),
10}
11
12enum Matches {
13    Left,
14    Right,
15    Both,
16}
17
18pub struct EitherBothState<T: FetchState, U: FetchState> {
19    pub(super) left_state: T,
20    pub(super) right_state: U,
21}
22
23unsafe impl<T: FetchState, U: FetchState> FetchState for EitherBothState<T, U> {
24    fn init(world: &mut World) -> Self {
25        EitherBothState {
26            left_state: T::init(world),
27            right_state: U::init(world),
28        }
29    }
30
31    fn update_component_access(&self, access: &mut FilteredAccess<ComponentId>) {
32        self.left_state.update_component_access(access);
33        self.right_state.update_component_access(access);
34    }
35
36    fn update_archetype_component_access(
37        &self,
38        archetype: &Archetype,
39        access: &mut Access<ArchetypeComponentId>,
40    ) {
41        self.left_state.update_archetype_component_access(archetype, access);
42        self.right_state.update_archetype_component_access(archetype, access);
43    }
44
45    fn matches_archetype(&self, archetype: &Archetype) -> bool {
46        self.left_state.matches_archetype(archetype) || self.right_state.matches_archetype(archetype)
47    }
48
49    fn matches_table(&self, table: &Table) -> bool {
50        self.left_state.matches_table(table) || self.right_state.matches_table(table)
51    }
52}
53
54pub struct EitherBothFetch<T, U> {
55    left: T,
56    right: U,
57    matches: Matches,
58}
59
60unsafe impl<'w, T: Fetch<'w>, U: Fetch<'w>> ReadOnlyFetch for EitherBothFetch<T, U> {}
61
62impl<'w, T: Fetch<'w>, U: Fetch<'w>> Fetch<'w> for EitherBothFetch<T, U> {
63    type Item = EitherBoth<T::Item, U::Item>;
64    type State = EitherBothState<T::State, U::State>;
65
66    fn is_dense(&self) -> bool {
67        self.left.is_dense() && self.right.is_dense()
68    }
69
70    unsafe fn init(
71        world: &World,
72        state: &Self::State,
73        last_change_tick: u32,
74        change_tick: u32,
75    ) -> Self {
76        EitherBothFetch {
77            left: T::init(world, &state.left_state, last_change_tick, change_tick),
78            right: U::init(world, &state.right_state, last_change_tick, change_tick),
79            matches: Matches::Left,
80        }
81    }
82
83    unsafe fn set_archetype(
84        &mut self,
85        state: &Self::State,
86        archetype: &Archetype,
87        tables: &Tables,
88    ) {
89        let left_match = state.left_state.matches_archetype(archetype);
90        let right_match = state.right_state.matches_archetype(archetype);
91        if left_match {
92            self.left.set_archetype(&state.left_state, archetype, tables);
93            if right_match {
94                self.matches = Matches::Both;
95                self.right.set_archetype(&state.right_state, archetype, tables);
96            } else {
97                self.matches = Matches::Left;
98            }
99        } else if right_match {
100            self.matches = Matches::Right;
101            self.right.set_archetype(&state.right_state, archetype, tables);
102        } else if cfg!(not(all(not(debug_assertions), unchecked))) {
103            unreachable!("neither left nor right side matched. what?");
104        }
105    }
106
107    unsafe fn set_table(&mut self, state: &Self::State, table: &Table) {
108        let left_match = state.left_state.matches_table(table);
109        let right_match = state.right_state.matches_table(table);
110        if left_match {
111            self.left.set_table(&state.left_state, table);
112            if right_match {
113                self.matches = Matches::Both;
114                self.right.set_table(&state.right_state, table);
115            } else {
116                self.matches = Matches::Left;
117            }
118        } else if right_match {
119            self.matches = Matches::Right;
120            self.right.set_table(&state.right_state, table);
121        } else if cfg!(not(all(not(debug_assertions), unchecked))) {
122            unreachable!("neither left nor right side matched. what?");
123        }
124    }
125
126    unsafe fn archetype_fetch(&mut self, archetype_index: usize) -> Self::Item {
127        match self.matches {
128            Matches::Both => EitherBoth::Both(
129                self.left.archetype_fetch(archetype_index),
130                self.right.archetype_fetch(archetype_index),
131            ),
132            Matches::Left => EitherBoth::Left(self.left.archetype_fetch(archetype_index)),
133            Matches::Right => EitherBoth::Right(self.right.archetype_fetch(archetype_index)),
134        }
135    }
136
137    unsafe fn table_fetch(&mut self, table_row: usize) -> Self::Item {
138        match self.matches {
139            Matches::Both => EitherBoth::Both(
140                self.left.table_fetch(table_row),
141                self.right.table_fetch(table_row),
142            ),
143            Matches::Left => EitherBoth::Left(self.left.table_fetch(table_row)),
144            Matches::Right => EitherBoth::Right(self.right.table_fetch(table_row)),
145        }
146    }
147}
148
149impl<T: WorldQuery, U: WorldQuery> WorldQuery for EitherBoth<T, U> {
150    type Fetch = EitherBothFetch<T::Fetch, U::Fetch>;
151    type State = EitherBothState<T::State, U::State>;
152}
153
154#[cfg(test)]
155mod tests {
156    use bevy::prelude::*;
157    use super::*;
158
159    #[derive(Clone, Copy)]
160    struct LeftElem;
161
162    #[derive(Clone, Copy)]
163    struct RightElem;
164
165    fn push_entities(world: &mut World) -> (u32, u32, u32) {
166        // 0 = None, 1 = Some(Left), 2 = Some(Right), 3 = Some(Both)
167        static SUPERPERM: &[Option<EitherBoth<LeftElem, RightElem>>] = &[
168            None,
169            Some(EitherBoth::Left(LeftElem)),
170            Some(EitherBoth::Right(RightElem)),
171            Some(EitherBoth::Both(LeftElem, RightElem)),
172            None,
173            Some(EitherBoth::Left(LeftElem)),
174            Some(EitherBoth::Right(RightElem)),
175            None,
176            Some(EitherBoth::Both(LeftElem, RightElem)),
177            Some(EitherBoth::Left(LeftElem)),
178            Some(EitherBoth::Right(RightElem)),
179            None,
180            Some(EitherBoth::Left(LeftElem)),
181            Some(EitherBoth::Both(LeftElem, RightElem)),
182            Some(EitherBoth::Right(RightElem)),
183            None,
184            Some(EitherBoth::Left(LeftElem)),
185            None,
186            Some(EitherBoth::Right(RightElem)),
187            Some(EitherBoth::Both(LeftElem, RightElem)),
188            Some(EitherBoth::Left(LeftElem)),
189            None,
190            Some(EitherBoth::Right(RightElem)),
191            Some(EitherBoth::Left(LeftElem)),
192            Some(EitherBoth::Both(LeftElem, RightElem)),
193            None,
194            Some(EitherBoth::Right(RightElem)),
195            Some(EitherBoth::Left(LeftElem)),
196            None,
197            Some(EitherBoth::Both(LeftElem, RightElem)),
198            Some(EitherBoth::Right(RightElem)),
199            Some(EitherBoth::Left(LeftElem)),
200            None,
201        ];
202        let mut left_count = 0;
203        let mut right_count = 0;
204        let mut both_count = 0;
205
206        for &p in SUPERPERM {
207            match p {
208                Some(EitherBoth::Both(l, r)) => {
209                    world.spawn().insert(l).insert(r);
210                    both_count += 1;
211                },
212                Some(EitherBoth::Left(l)) => {
213                    world.spawn().insert(l);
214                    left_count += 1;
215                },
216                Some(EitherBoth::Right(r)) => {
217                    world.spawn().insert(r);
218                    right_count += 1;
219                },
220                None => {
221                    world.spawn();
222                },
223            }
224        }
225
226        (left_count, right_count, both_count)
227    }
228
229    #[derive(Debug, PartialEq, Eq)]
230    struct LeftCount(u32);
231
232    #[derive(Debug, PartialEq, Eq)]
233    struct RightCount(u32);
234
235    #[derive(Debug, PartialEq, Eq)]
236    struct BothCount(u32);
237
238    #[test]
239    fn test_eitherboth() {
240        let mut world = World::default();
241        world.insert_resource(LeftCount(0));
242        world.insert_resource(RightCount(0));
243        world.insert_resource(BothCount(0));
244        let (real_left_count, real_right_count, real_both_count) = push_entities(&mut world);
245        let mut update_stage = SystemStage::single((|
246                q: Query<EitherBoth<&LeftElem, &RightElem>>,
247                mut l: ResMut<LeftCount>,
248                mut r: ResMut<RightCount>,
249                mut b: ResMut<BothCount>,
250            | {
251                for eb in q.iter() {
252                    match eb {
253                        EitherBoth::Left(_) => {
254                            l.0 += 1;
255                        },
256                        EitherBoth::Right(_) => {
257                            r.0 += 1;
258                        },
259                        EitherBoth::Both(_, _) => {
260                            b.0 += 1;
261                        },
262                    }
263                }
264            }
265        ).system());
266        update_stage.run(&mut world);
267        assert_eq!(world.get_resource::<LeftCount>().unwrap().0, real_left_count);
268        assert_eq!(world.get_resource::<RightCount>().unwrap().0, real_right_count);
269        assert_eq!(world.get_resource::<BothCount>().unwrap().0, real_both_count);
270    }
271}
bevy_either/either_both.rs

bevy_either/
either_both.rs
