bevy_either/
either.rs

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

bevy_either/
either.rs
