avian3d/sync/
ancestor_marker.rs

1//! Functionality for marking ancestors of entities with marker components.
2
3use core::marker::PhantomData;
4
5use bevy::prelude::*;
6
7/// A plugin that marks the ancestors of entities that have the given component `C`
8/// with the [`AncestorMarker`] component.
9///
10/// One use case is speeding up transform propagation: we only need to propagate
11/// down trees that have a certain type of entity, like a collider or a rigid body.
12pub struct AncestorMarkerPlugin<C: Component>(PhantomData<C>);
13
14impl<C: Component> Default for AncestorMarkerPlugin<C> {
15    fn default() -> Self {
16        Self(PhantomData)
17    }
18}
19
20impl<C: Component> Plugin for AncestorMarkerPlugin<C> {
21    fn build(&self, app: &mut App) {
22        // Add `AncestorMarker<C>` for the ancestors of colliders that are inserted as children,
23        // until an ancestor that has other `AncestorMarker<C>` entities as children is encountered.
24        app.add_observer(
25            |trigger: Trigger<OnInsert, (ChildOf, C)>,
26             mut commands: Commands,
27             collider_query: Query<&C>,
28             parent_query: Query<&ChildOf>,
29             ancestor_query: Query<(), With<AncestorMarker<C>>>| {
30                let entity = trigger.target();
31                if collider_query.contains(entity) {
32                    add_ancestor_markers(
33                        entity,
34                        &mut commands,
35                        &parent_query,
36                        &ancestor_query,
37                        false,
38                    );
39                }
40            },
41        );
42
43        // Remove `AncestorMarker<C>` from the ancestors of colliders that have been removed or moved to another parent,
44        // until an ancestor that has other `AncestorMarker<C>` entities as children is encountered.
45        #[allow(clippy::type_complexity)]
46        app.add_observer(
47            |trigger: Trigger<OnReplace, (ChildOf, C)>,
48            mut commands: Commands,
49            collider_query: Query<&C>,
50            child_query: Query<&Children>,
51            parent_query: Query<&ChildOf>,
52            ancestor_query: Query<
53                (Entity, Has<C>),
54                Or<(With<AncestorMarker<C>>, With<C>)>
55            >| {
56                let entity = trigger.target();
57                if collider_query.contains(entity) {
58                    remove_ancestor_markers(entity, &mut commands, &parent_query, &child_query, &ancestor_query, false);
59                }
60            },
61        );
62    }
63}
64
65/// A marker component that marks an entity as an ancestor of an entity with the given component `C`.
66///
67/// This is added and removed automatically by the [`AncestorMarkerPlugin`] if it is enabled.
68#[derive(Component, Copy, Reflect)]
69#[reflect(Component, Default)]
70pub struct AncestorMarker<C: Component> {
71    #[reflect(ignore)]
72    _phantom: PhantomData<C>,
73}
74
75impl<C: Component> Clone for AncestorMarker<C> {
76    fn clone(&self) -> Self {
77        Self::default()
78    }
79}
80
81impl<C: Component> Default for AncestorMarker<C> {
82    fn default() -> Self {
83        Self {
84            _phantom: PhantomData,
85        }
86    }
87}
88
89fn add_ancestor_markers<C: Component>(
90    entity: Entity,
91    commands: &mut Commands,
92    parent_query: &Query<&ChildOf>,
93    ancestor_query: &Query<(), With<AncestorMarker<C>>>,
94    include_self: bool,
95) {
96    if include_self {
97        commands
98            .entity(entity)
99            .insert(AncestorMarker::<C>::default());
100    }
101
102    // Traverse up the tree, marking entities with `AncestorMarker<C>`
103    // until an entity that already has it is encountered.
104    for parent_entity in parent_query.iter_ancestors(entity) {
105        if ancestor_query.contains(parent_entity) {
106            break;
107        } else {
108            commands
109                .entity(parent_entity)
110                .insert(AncestorMarker::<C>::default());
111        }
112    }
113}
114
115/// Remove the component from entity, unless it is already despawned.
116fn remove_component<T: Bundle>(commands: &mut Commands, entity: Entity) {
117    if let Ok(mut entity_commands) = commands.get_entity(entity) {
118        entity_commands.try_remove::<T>();
119    }
120}
121
122#[allow(clippy::type_complexity)]
123fn remove_ancestor_markers<C: Component>(
124    entity: Entity,
125    commands: &mut Commands,
126    parent_query: &Query<&ChildOf>,
127    child_query: &Query<&Children>,
128    ancestor_query: &Query<(Entity, Has<C>), Or<(With<AncestorMarker<C>>, With<C>)>>,
129    include_self: bool,
130) {
131    if include_self {
132        // Remove the marker from the `parent` unless a sibling of the `child`
133        // is also a marked ancestor or has `C`.
134        if let Ok(children) = child_query.get(entity) {
135            let keep_marker = ancestor_query
136                .iter_many(children)
137                .any(|(parent_child, _has_c)| parent_child != entity);
138            if keep_marker {
139                return;
140            } else {
141                remove_component::<AncestorMarker<C>>(commands, entity);
142            }
143        } else {
144            // The parent has no children, so it cannot be an ancestor.
145            remove_component::<AncestorMarker<C>>(commands, entity);
146        }
147    }
148
149    // Iterate over ancestors, removing `AncestorMarker<C>` markers until
150    // an entity that has other `AncestorMarker<C>` children is encountered.
151    let mut previous_parent = entity;
152    for parent_entity in parent_query.iter_ancestors(entity) {
153        if let Ok(children) = child_query.get(parent_entity) {
154            // Keep the marker if `parent_entity` has a child that is a marked ancestor
155            // or an entity that has `C`, but not the one that was removed.
156            let keep_marker = ancestor_query
157                .iter_many(children)
158                .any(|(child, has_c)| child != previous_parent || (has_c && child != entity));
159
160            if keep_marker {
161                return;
162            } else {
163                remove_component::<AncestorMarker<C>>(commands, parent_entity);
164            }
165        } else {
166            // The parent has no children, so it cannot be an ancestor.
167            remove_component::<AncestorMarker<C>>(commands, entity);
168        }
169
170        previous_parent = parent_entity;
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    #[derive(Component)]
179    struct C;
180
181    #[test]
182    fn add_and_remove_component() {
183        let mut app = App::new();
184
185        app.add_plugins(AncestorMarkerPlugin::<C>::default());
186
187        // Set up an entity tree like the following:
188        //
189        //     AN
190        //    /  \
191        //  BN    CY
192        //       /  \
193        //     DN    EN
194        //    /  \
195        //  FY    GY
196        //
197        // where Y means that the entity has `C`,
198        // and N means that the entity does not have `C`.
199
200        let an = app.world_mut().spawn_empty().id();
201
202        let bn = app.world_mut().spawn(ChildOf(an)).id();
203        let cy = app.world_mut().spawn((C, ChildOf(an))).id();
204
205        let dn = app.world_mut().spawn(ChildOf(cy)).id();
206        let en = app.world_mut().spawn(ChildOf(cy)).id();
207
208        let fy = app.world_mut().spawn((C, ChildOf(dn))).id();
209        let gy = app.world_mut().spawn((C, ChildOf(dn))).id();
210
211        // Check that the correct entities have the `AncestorMarker<C>` component.
212        assert!(app.world().entity(an).contains::<AncestorMarker<C>>());
213        assert!(!app.world().entity(bn).contains::<AncestorMarker<C>>());
214        assert!(app.world().entity(cy).contains::<AncestorMarker<C>>());
215        assert!(app.world().entity(dn).contains::<AncestorMarker<C>>());
216        assert!(!app.world().entity(en).contains::<AncestorMarker<C>>());
217        assert!(!app.world().entity(fy).contains::<AncestorMarker<C>>());
218        assert!(!app.world().entity(gy).contains::<AncestorMarker<C>>());
219
220        // Remove `C` from FY. DN, CY, and AN should all keep the `AncestorMarker<C>` marker.
221        let mut entity_mut = app.world_mut().entity_mut(fy);
222        entity_mut.remove::<C>();
223
224        assert!(app.world().entity(dn).contains::<AncestorMarker<C>>());
225        assert!(app.world().entity(cy).contains::<AncestorMarker<C>>());
226        assert!(app.world().entity(an).contains::<AncestorMarker<C>>());
227
228        // Remove `C` from GY. The `AncestorMarker<C>` marker should
229        // now be removed from DN and CY, but it should remain on AN.
230        let mut entity_mut = app.world_mut().entity_mut(gy);
231        entity_mut.remove::<C>();
232
233        assert!(!app.world().entity(dn).contains::<AncestorMarker<C>>());
234        assert!(!app.world().entity(cy).contains::<AncestorMarker<C>>());
235        assert!(app.world().entity(an).contains::<AncestorMarker<C>>());
236
237        // Remove `C` from CY. The `AncestorMarker<C>` marker should
238        // now be removed from AN.
239        let mut entity_mut = app.world_mut().entity_mut(cy);
240        entity_mut.remove::<C>();
241
242        assert!(!app.world().entity(an).contains::<AncestorMarker<C>>());
243    }
244
245    #[test]
246    fn remove_children() {
247        let mut app = App::new();
248
249        app.add_plugins(AncestorMarkerPlugin::<C>::default());
250
251        // Set up an entity tree like the following:
252        //
253        //     AN
254        //    /  \
255        //  BN    CY
256        //       /  \
257        //     DN    EN
258        //    /  \
259        //  FY    GY
260        //
261        // where Y means that the entity has `C`,
262        // and N means that the entity does not have `C`.
263
264        let an = app.world_mut().spawn_empty().id();
265
266        let bn = app.world_mut().spawn(ChildOf(an)).id();
267        let cy = app.world_mut().spawn((C, ChildOf(an))).id();
268
269        let dn = app.world_mut().spawn(ChildOf(cy)).id();
270        let en = app.world_mut().spawn(ChildOf(cy)).id();
271
272        let fy = app.world_mut().spawn((C, ChildOf(dn))).id();
273        let gy = app.world_mut().spawn((C, ChildOf(dn))).id();
274
275        // Check that the correct entities have the `AncestorMarker<C>` component.
276        assert!(app.world().entity(an).contains::<AncestorMarker<C>>());
277        assert!(!app.world().entity(bn).contains::<AncestorMarker<C>>());
278        assert!(app.world().entity(cy).contains::<AncestorMarker<C>>());
279        assert!(app.world().entity(dn).contains::<AncestorMarker<C>>());
280        assert!(!app.world().entity(en).contains::<AncestorMarker<C>>());
281        assert!(!app.world().entity(fy).contains::<AncestorMarker<C>>());
282        assert!(!app.world().entity(gy).contains::<AncestorMarker<C>>());
283
284        // Make FY an orphan.
285        let mut entity_mut = app.world_mut().entity_mut(fy);
286        entity_mut.remove::<ChildOf>();
287
288        assert!(app.world().entity(dn).contains::<AncestorMarker<C>>());
289        assert!(app.world().entity(cy).contains::<AncestorMarker<C>>());
290        assert!(app.world().entity(an).contains::<AncestorMarker<C>>());
291
292        // Make GY an orphan. The `AncestorMarker<C>` marker should
293        // now be removed from DN and CY, but it should remain on AN.
294        let mut entity_mut = app.world_mut().entity_mut(gy);
295        entity_mut.remove::<ChildOf>();
296
297        assert!(!app.world().entity(dn).contains::<AncestorMarker<C>>());
298        assert!(!app.world().entity(cy).contains::<AncestorMarker<C>>());
299        assert!(app.world().entity(an).contains::<AncestorMarker<C>>());
300
301        // Make CY an orphan. The `AncestorMarker<C>` marker should
302        // now be removed from AN.
303        let mut entity_mut = app.world_mut().entity_mut(cy);
304        entity_mut.remove::<ChildOf>();
305
306        assert!(!app.world().entity(an).contains::<AncestorMarker<C>>());
307
308        // Make CY a child of AN again. The `AncestorMarker<C>` marker should
309        // now be added to AN.
310        let mut entity_mut = app.world_mut().entity_mut(an);
311        entity_mut.add_child(cy);
312        assert!(app.world().entity(an).contains::<AncestorMarker<C>>());
313
314        // Make CY an orphan and delete AN. This must not crash.
315        let mut entity_mut = app.world_mut().entity_mut(cy);
316        entity_mut.remove::<ChildOf>();
317        entity_mut.despawn();
318    }
319
320    #[test]
321    fn move_children() {
322        let mut app = App::new();
323
324        app.add_plugins(AncestorMarkerPlugin::<C>::default());
325
326        // Set up an entity tree like the following:
327        //
328        //     AN
329        //    /  \
330        //  BN    CY
331        //       /  \
332        //     DN    EN
333        //    /  \
334        //  FY    GY
335        //
336        // where Y means that the entity has `C`,
337        // and N means that the entity does not have `C`.
338
339        let an = app.world_mut().spawn_empty().id();
340
341        let bn = app.world_mut().spawn(ChildOf(an)).id();
342        let cy = app.world_mut().spawn((C, ChildOf(an))).id();
343
344        let dn = app.world_mut().spawn(ChildOf(cy)).id();
345        let en = app.world_mut().spawn(ChildOf(cy)).id();
346
347        let fy = app.world_mut().spawn((C, ChildOf(dn))).id();
348        let gy = app.world_mut().spawn((C, ChildOf(dn))).id();
349
350        // Check that the correct entities have the `AncestorMarker<C>` component.
351        assert!(app.world().entity(an).contains::<AncestorMarker<C>>());
352        assert!(!app.world().entity(bn).contains::<AncestorMarker<C>>());
353        assert!(app.world().entity(cy).contains::<AncestorMarker<C>>());
354        assert!(app.world().entity(dn).contains::<AncestorMarker<C>>());
355        assert!(!app.world().entity(en).contains::<AncestorMarker<C>>());
356        assert!(!app.world().entity(fy).contains::<AncestorMarker<C>>());
357        assert!(!app.world().entity(gy).contains::<AncestorMarker<C>>());
358
359        // Move FY to be a child of BN. BN should get the `AncestorMarker<C>` component.
360        let mut entity_mut = app.world_mut().entity_mut(bn);
361        entity_mut.add_child(fy);
362
363        assert!(app.world().entity(bn).contains::<AncestorMarker<C>>());
364        assert!(app.world().entity(dn).contains::<AncestorMarker<C>>());
365        assert!(app.world().entity(cy).contains::<AncestorMarker<C>>());
366        assert!(app.world().entity(an).contains::<AncestorMarker<C>>());
367
368        // Move GY to be a child of BN. The `AncestorMarker<C>` marker should
369        // now be removed from DN and CY.
370        let mut entity_mut = app.world_mut().entity_mut(bn);
371        entity_mut.add_child(gy);
372
373        assert!(!app.world().entity(dn).contains::<AncestorMarker<C>>());
374        assert!(!app.world().entity(cy).contains::<AncestorMarker<C>>());
375        assert!(app.world().entity(an).contains::<AncestorMarker<C>>());
376
377        // Move all children from BN to CY. The `AncestorMarker<C>` marker should
378        // now be removed from BN, and CY should get it back.
379        let mut entity_mut = app.world_mut().entity_mut(cy);
380        entity_mut.add_children(&[fy, gy]);
381
382        assert!(!app.world().entity(bn).contains::<AncestorMarker<C>>());
383        assert!(app.world().entity(cy).contains::<AncestorMarker<C>>());
384        assert!(app.world().entity(an).contains::<AncestorMarker<C>>());
385
386        // Move all children from CY to BN and remove CY. This must not crash.
387        let mut entity_mut = app.world_mut().entity_mut(bn);
388        entity_mut.add_children(&[dn, en, fy, gy]);
389    }
390}
avian3d/sync/ancestor_marker.rs

avian3d/sync/
ancestor_marker.rs
