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

avian2d/dynamics/solver/joint_graph/
mod.rs
