avian3d/collision/collider/collider_hierarchy/
plugin.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
use crate::{prelude::*, sync::ancestor_marker::AncestorMarker};
use bevy::prelude::*;

/// A plugin for managing [`ColliderOf`] relationships based on the entity hierarchy.
///
/// By default, colliders are automatically attached to the nearest rigid body ancestor.
///
/// This plugin requires that colliders have the [`ColliderMarker`] component,
/// which is added automatically for colliders if the [`ColliderBackendPlugin`] is enabled.
pub struct ColliderHierarchyPlugin;

impl Plugin for ColliderHierarchyPlugin {
    fn build(&self, app: &mut App) {
        // Imsert `ColliderOf` for colliders that are added to a rigid body.
        app.add_observer(
            |trigger: Trigger<OnAdd, (RigidBody, ColliderMarker)>,
             mut commands: Commands,
             query: Query<(), With<ColliderMarker>>,
             rb_query: Query<Entity, With<RigidBody>>,
             parent_query: Query<&ChildOf>| {
                let entity = trigger.target();

                // Make sure the entity is a collider.
                if !query.contains(entity) {
                    return;
                }

                // Find the closest rigid body ancestor.
                if let Some(body) = rb_query.get(entity).ok().map_or_else(
                    || {
                        parent_query
                            .iter_ancestors(trigger.target())
                            .find(|&entity| rb_query.contains(entity))
                    },
                    Some,
                ) {
                    commands.entity(entity).insert(ColliderOf { body });
                }
            },
        );

        // Remove `ColliderOf` when the rigid body or collider is removed.
        app.add_observer(
            |trigger: Trigger<OnRemove, (RigidBody, ColliderMarker)>,
             mut commands: Commands,
             query: Query<(), With<ColliderMarker>>| {
                let entity = trigger.target();

                // Make sure the collider is on the same entity as the rigid body.
                if query.contains(entity) {
                    commands.entity(entity).try_remove::<ColliderOf>();
                }
            },
        );

        // Update `ColliderOf` components for colliders when their ancestors change.
        app.add_observer(on_collider_body_changed);

        // Remove `ColliderOf` from colliders when their rigid bodies are removed.
        app.add_observer(on_body_removed);
    }
}

/// Updates [`ColliderOf`] components for colliders when their ancestors change
/// or when a rigid body is added to the hierarchy.
fn on_collider_body_changed(
    trigger: Trigger<OnInsert, (ChildOf, RigidBody, AncestorMarker<ColliderMarker>)>,
    mut commands: Commands,
    query: Query<
        Has<ColliderMarker>,
        Or<(With<ColliderMarker>, With<AncestorMarker<ColliderMarker>>)>,
    >,
    child_query: Query<&Children>,
    parent_query: Query<&ChildOf>,
    body_query: Query<Entity, With<RigidBody>>,
    collider_query: Query<(), With<ColliderMarker>>,
) {
    let entity = trigger.target();

    // Skip if the entity is not a collider or an ancestor of a collider.
    let Ok(is_collider) = query.get(entity) else {
        return;
    };

    // Find the closest rigid body ancestor.
    let Some(body) = body_query.get(entity).ok().map_or_else(
        || {
            parent_query
                .iter_ancestors(trigger.target())
                .find(|&entity| body_query.contains(entity))
        },
        Some,
    ) else {
        return;
    };

    // Insert `ColliderOf` for the new child entity.
    if is_collider {
        commands.entity(entity).insert(ColliderOf { body });
    }

    // Iterate over all descendants starting from the new child entity,
    // and update the `ColliderOf` component to point to the closest rigid body ancestor.
    for child in child_query.iter_descendants(entity) {
        // Stop traversal if we reach another rigid body.
        // FIXME: This might lead to siblings of the rigid body being skipped.
        if body_query.contains(child) {
            break;
        }

        // Update the `ColliderOf` component.
        if collider_query.contains(child) {
            commands.entity(child).insert(ColliderOf { body });
        }
    }
}

/// Removes [`ColliderOf`] from colliders when their rigid bodies are removed.
fn on_body_removed(
    trigger: Trigger<OnRemove, RigidBody>,
    mut commands: Commands,
    body_collider_query: Query<&RigidBodyColliders>,
) {
    // TODO: Here we assume that rigid bodies are not nested, so `ColliderOf` is simply removed
    //       instead of being updated to point to a new rigid body in the hierarchy.

    let body = trigger.target();

    // Remove `ColliderOf` from all colliders attached to the rigid body.
    if let Ok(colliders) = body_collider_query.get(body) {
        for collider in colliders.iter() {
            commands
                .entity(collider)
                .try_remove::<(ColliderOf, ColliderTransform)>();
        }
    }
}