avian3d/dynamics/joints/joint_graph/
mod.rs

1//! A [`JointGraph`] for tracking how [rigid bodies] are connected by [joints].
2//!
3//! [rigid bodies]: crate::dynamics::RigidBody
4//! [joints]: crate::dynamics::joints
5
6mod plugin;
7pub use plugin::{JointComponentId, JointGraphChange, JointGraphPlugin};
8
9use crate::data_structures::{
10    graph::{EdgeIndex, NodeIndex},
11    sparse_secondary_map::SparseSecondaryEntityMap,
12    stable_graph::StableUnGraph,
13};
14use bevy::prelude::*;
15
16// TODO: Once we have many-to-many relationships, we could potentially represent the joint graph in the ECS.
17
18/// A resource for the joint graph, tracking how [rigid bodies] are connected by [joints].
19///
20/// [rigid bodies]: crate::dynamics::RigidBody
21/// [joints]: crate::dynamics::joints
22#[derive(Resource, Clone, Debug, Default)]
23pub struct JointGraph {
24    graph: StableUnGraph<Entity, JointGraphEdge>,
25    entity_to_body: SparseSecondaryEntityMap<NodeIndex>,
26    entity_to_joint: SparseSecondaryEntityMap<EdgeIndex>,
27}
28
29/// A stable identifier for a [`JointGraphEdge`].
30#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Reflect)]
31#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
32#[cfg_attr(feature = "serialize", reflect(Serialize, Deserialize))]
33#[reflect(Debug, PartialEq)]
34pub struct JointId(pub u32);
35
36impl JointId {
37    /// A placeholder identifier for a [`JointGraphEdge`].
38    pub const PLACEHOLDER: Self = Self(u32::MAX);
39}
40
41impl From<JointId> for EdgeIndex {
42    fn from(id: JointId) -> Self {
43        Self(id.0)
44    }
45}
46
47impl From<EdgeIndex> for JointId {
48    fn from(id: EdgeIndex) -> Self {
49        Self(id.0)
50    }
51}
52
53impl core::fmt::Display for JointId {
54    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
55        write!(f, "JointId({})", self.0)
56    }
57}
58
59/// An edge in the [`JointGraph`].
60#[derive(Clone, Debug, Reflect)]
61#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
62#[cfg_attr(feature = "serialize", reflect(Serialize, Deserialize))]
63#[reflect(Debug)]
64pub struct JointGraphEdge {
65    /// The stable identifier of this joint edge.
66    pub id: JointId,
67
68    /// The entity of the joint.
69    pub entity: Entity,
70
71    /// The first [rigid body] connected by this joint.
72    ///
73    /// [rigid body]: crate::dynamics::RigidBody
74    pub body1: Entity,
75
76    /// The second [rigid body] connected by this joint.
77    ///
78    /// [rigid body]: crate::dynamics::RigidBody
79    pub body2: Entity,
80
81    /// If `true`, collisions are disabled between the bodies connected by this joint.
82    ///
83    /// This is controlled by the [`JointCollisionDisabled`] component.
84    ///
85    /// [`JointCollisionDisabled`]: crate::dynamics::joints::JointCollisionDisabled
86    pub collision_disabled: bool,
87}
88
89impl JointGraphEdge {
90    /// Creates a new [`JointGraphEdge`].
91    #[inline]
92    pub fn new(entity: Entity, body1: Entity, body2: Entity, collision_disabled: bool) -> Self {
93        Self {
94            // This gets set to a valid ID when the joint is added to the `JointGraph`.
95            id: JointId::PLACEHOLDER,
96            entity,
97            body1,
98            body2,
99            collision_disabled,
100        }
101    }
102}
103
104impl JointGraph {
105    /// Returns a reference to the underlying [`StableUnGraph`].
106    #[inline]
107    pub fn graph(&self) -> &StableUnGraph<Entity, JointGraphEdge> {
108        &self.graph
109    }
110
111    /// Returns the [`NodeIndex`] of the given entity in the joint graph.
112    ///
113    /// If the entity is not in the graph, `None` is returned.
114    #[inline]
115    pub fn entity_to_body(&self, entity: Entity) -> Option<NodeIndex> {
116        self.entity_to_body.get(entity).copied()
117    }
118
119    /// Returns the [`JointId`] of the joint edge for the given entity.
120    ///
121    /// If the entity is not in the graph, `None` is returned.
122    #[inline]
123    pub fn entity_to_joint(&self, entity: Entity) -> Option<JointId> {
124        self.entity_to_joint.get(entity).copied().map(JointId::from)
125    }
126
127    /// Returns the [`JointGraphEdge`] for the given joint entity.
128    /// If the joint is not in the graph, `None` is returned.
129    #[inline]
130    pub fn get(&self, joint: Entity) -> Option<&JointGraphEdge> {
131        let joint_index = self.entity_to_joint(joint)?;
132        self.get_by_id(joint_index)
133    }
134
135    /// Returns a reference to the [`JointGraphEdge`] for the given [`JointId`].
136    /// If the joint is not in the graph, `None` is returned.
137    #[inline]
138    pub fn get_by_id(&self, joint_id: JointId) -> Option<&JointGraphEdge> {
139        self.graph.edge_weight(joint_id.into())
140    }
141
142    /// Returns a mutable reference to the [`JointGraphEdge`] for the given joint entity.
143    /// If the joint is not in the graph, `None` is returned.
144    #[inline]
145    pub fn get_mut(&mut self, joint: Entity) -> Option<&mut JointGraphEdge> {
146        let joint_index = self.entity_to_joint(joint)?;
147        self.get_mut_by_id(joint_index)
148    }
149
150    /// Returns a mutable reference to the [`JointGraphEdge`] for the given [`JointId`].
151    /// If the joint is not in the graph, `None` is returned.
152    #[inline]
153    pub fn get_mut_by_id(&mut self, joint_id: JointId) -> Option<&mut JointGraphEdge> {
154        self.graph.edge_weight_mut(joint_id.into())
155    }
156
157    /// Returns the [`JointGraphEdge`] between two entities.
158    /// If the edge does not exist, `None` is returned.
159    #[inline]
160    pub fn joints_between(
161        &self,
162        body1: Entity,
163        body2: Entity,
164    ) -> impl Iterator<Item = &JointGraphEdge> {
165        let (Some(index1), Some(index2)) = (self.entity_to_body(body1), self.entity_to_body(body2))
166        else {
167            return itertools::Either::Left(core::iter::empty());
168        };
169
170        let joints = self.graph.edges_between(index1, index2).map(|e| e.weight());
171        itertools::Either::Right(joints)
172    }
173
174    /// Returns an iterator yielding immutable access to all joint edges involving the given entity.
175    #[inline]
176    pub fn joints_of(&self, body: Entity) -> impl Iterator<Item = &JointGraphEdge> {
177        let index = self.entity_to_body(body);
178        if let Some(index) = index {
179            itertools::Either::Left(self.graph.edge_weights(index))
180        } else {
181            itertools::Either::Right(core::iter::empty())
182        }
183    }
184
185    /// Returns an iterator yielding mutable access to all joint edges involving the given entity.
186    #[inline]
187    pub fn joints_of_mut(&mut self, body: Entity) -> impl Iterator<Item = &mut JointGraphEdge> {
188        let index = self.entity_to_body(body);
189        if let Some(index) = index {
190            itertools::Either::Left(self.graph.edge_weights_mut(index))
191        } else {
192            itertools::Either::Right(core::iter::empty())
193        }
194    }
195
196    /// Returns the bodies that are connected by the given joint entity.
197    /// If the joint is not in the graph, `None` is returned.
198    #[inline]
199    pub fn bodies_of(&self, joint: Entity) -> Option<[Entity; 2]> {
200        let joint_index = self.entity_to_joint(joint)?;
201        let (body1_index, body2_index) = self.graph.edge_endpoints(joint_index.into())?;
202
203        Some([
204            *self.graph.node_weight(body1_index)?,
205            *self.graph.node_weight(body2_index)?,
206        ])
207    }
208
209    /// Returns an iterator yielding immutable access to all bodies that are attached
210    /// to the given entity with a joint.
211    #[inline]
212    pub fn bodies_attached_to(&self, body: Entity) -> impl Iterator<Item = Entity> + '_ {
213        self.entity_to_body
214            .get(body)
215            .into_iter()
216            .flat_map(move |&index| {
217                self.graph
218                    .neighbors(index)
219                    .map(|index| *self.graph.node_weight(index).unwrap())
220            })
221    }
222
223    /// Creates a [`JointGraphEdge`] between two entities if it does not already exist,
224    /// and returns the [`EdgeIndex`] of the created joint edge.
225    ///
226    /// # Warning
227    ///
228    /// Creating a joint edge with this method will *not* wake up the entities involved
229    /// or do any other clean-up. Only use this method if you know what you are doing.
230    #[inline]
231    pub fn add_joint(&mut self, joint_edge: JointGraphEdge) -> JointId {
232        let body1 = joint_edge.body1;
233        let body2 = joint_edge.body2;
234
235        // Get the indices of the entities in the graph.
236        let body1_index = self
237            .entity_to_body
238            .get_or_insert_with(body1, || self.graph.add_node(body1));
239        let body2_index = self
240            .entity_to_body
241            .get_or_insert_with(body2, || self.graph.add_node(body2));
242
243        // Add the edge to the graph.
244        let joint_entity = joint_edge.entity;
245        let edge_id = JointId(self.graph.add_edge(body1_index, body2_index, joint_edge).0);
246
247        // Insert the joint edge into the entity-to-joint mapping.
248        self.entity_to_joint
249            .get_or_insert_with(joint_entity, || edge_id.into());
250
251        // Set the joint ID in the joint edge.
252        let edge = self.graph.edge_weight_mut(edge_id.into()).unwrap();
253        edge.id = edge_id;
254
255        edge_id
256    }
257
258    /// Removes a [`JointGraphEdge`] between two entites and returns its value.
259    ///
260    /// # Warning
261    ///
262    /// Removing a joint edge with this method will *not* wake up the entities involved
263    /// or do any other clean-up. Only use this method if you know what you are doing.
264    #[inline]
265    pub fn remove_joint(&mut self, joint_entity: Entity) -> Option<JointGraphEdge> {
266        let joint_index = self.entity_to_joint.remove(joint_entity)?;
267        self.graph.remove_edge(joint_index)
268    }
269
270    /// Removes the body of the given entity from the joint graph, calling the given callback
271    /// for each [`JointGraphEdge`] right before it is removed.
272    ///
273    /// # Warning
274    ///
275    /// Removing a body with this method will *not* wake up the entities involved
276    /// or do any other clean-up. Only use this method if you know what you are doing.
277    #[inline]
278    pub fn remove_body_with<F>(&mut self, entity: Entity, edge_callback: F)
279    where
280        F: FnMut(&mut StableUnGraph<Entity, JointGraphEdge>, EdgeIndex),
281    {
282        // Remove the entity from the entity-to-node mapping,
283        // and get the index of the node in the graph.
284        let Some(index) = self.entity_to_body.remove(entity) else {
285            return;
286        };
287
288        // Remove the entity from the graph.
289        // TODO: Should we remove the joint from the entity-to-joint mapping as well?
290        self.graph.remove_node_with(index, edge_callback);
291
292        // Removing the node swapped the last node to its place,
293        // so we need to remap the entity-to-node mapping of the swapped node.
294        if let Some(swapped) = self.graph.node_weight(index).copied() {
295            let swapped_index = self
296                .entity_to_body
297                .get_mut(swapped)
298                // This should never panic.
299                .expect("swapped entity has no entity-to-node mapping");
300            *swapped_index = index;
301        }
302    }
303}
avian3d/dynamics/joints/joint_graph/mod.rs

avian3d/dynamics/joints/joint_graph/
mod.rs
