rapier2d/dynamics/joint/multibody_joint/
multibody_joint_set.rs

1use parry::utils::hashset::HashSet;
2
3use crate::data::{Arena, Coarena, Index};
4use crate::dynamics::joint::MultibodyLink;
5use crate::dynamics::{GenericJoint, Multibody, MultibodyJoint, RigidBodyHandle};
6use crate::geometry::{InteractionGraph, RigidBodyGraphIndex};
7
8/// The unique handle of an multibody_joint added to a `MultibodyJointSet`.
9#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
10#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
11#[repr(transparent)]
12pub struct MultibodyJointHandle(pub Index);
13
14/// The temporary index of a multibody added to a `MultibodyJointSet`.
15#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
16#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
17#[repr(transparent)]
18pub struct MultibodyIndex(pub Index);
19
20impl MultibodyJointHandle {
21    /// Converts this handle into its (index, generation) components.
22    pub fn into_raw_parts(self) -> (u32, u32) {
23        self.0.into_raw_parts()
24    }
25
26    /// Reconstructs an handle from its (index, generation) components.
27    pub fn from_raw_parts(id: u32, generation: u32) -> Self {
28        Self(Index::from_raw_parts(id, generation))
29    }
30
31    /// An always-invalid rigid-body handle.
32    pub fn invalid() -> Self {
33        Self(Index::from_raw_parts(
34            crate::INVALID_U32,
35            crate::INVALID_U32,
36        ))
37    }
38}
39
40impl Default for MultibodyJointHandle {
41    fn default() -> Self {
42        Self::invalid()
43    }
44}
45
46#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
47#[derive(Copy, Clone, Debug, PartialEq, Eq)]
48/// Indexes usable to get a multibody link from a `MultibodyJointSet`.
49///
50/// ```
51/// # use rapier3d::prelude::*;
52/// # let mut bodies = RigidBodySet::new();
53/// # let mut multibody_joint_set = MultibodyJointSet::new();
54/// # let body1 = bodies.insert(RigidBodyBuilder::dynamic());
55/// # let body2 = bodies.insert(RigidBodyBuilder::dynamic());
56/// # let joint = RevoluteJointBuilder::new(Vector::y_axis());
57/// # multibody_joint_set.insert(body1, body2, joint, true);
58/// # let multibody_link_id = multibody_joint_set.rigid_body_link(body2).unwrap();
59/// // With:
60/// //     multibody_joint_set: MultibodyJointSet
61/// //     multibody_link_id: MultibodyLinkId
62/// let multibody = &multibody_joint_set[multibody_link_id.multibody];
63/// let link = multibody.link(multibody_link_id.id).expect("Link not found.");
64/// ```
65pub struct MultibodyLinkId {
66    pub(crate) graph_id: RigidBodyGraphIndex,
67    /// The multibody index to be used as `&multibody_joint_set[multibody]` to
68    /// retrieve the multibody reference.
69    pub multibody: MultibodyIndex,
70    /// The multibody link index to be given to [`Multibody::link`].
71    pub id: usize,
72}
73
74impl Default for MultibodyLinkId {
75    fn default() -> Self {
76        Self {
77            graph_id: RigidBodyGraphIndex::new(crate::INVALID_U32),
78            multibody: MultibodyIndex(Index::from_raw_parts(
79                crate::INVALID_U32,
80                crate::INVALID_U32,
81            )),
82            id: 0,
83        }
84    }
85}
86
87#[derive(Default)]
88/// A set of rigid bodies that can be handled by a physics pipeline.
89#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
90#[derive(Clone, Debug)]
91pub struct MultibodyJointSet {
92    pub(crate) multibodies: Arena<Multibody>, // NOTE: a Slab would be sufficient.
93    pub(crate) rb2mb: Coarena<MultibodyLinkId>,
94    // NOTE: this is mostly for the island extraction. So perhaps we won’t need
95    //       that any more in the future when we improve our island builder.
96    pub(crate) connectivity_graph: InteractionGraph<RigidBodyHandle, ()>,
97    pub(crate) to_wake_up: HashSet<RigidBodyHandle>,
98}
99
100impl MultibodyJointSet {
101    /// Create a new empty set of multibodies.
102    pub fn new() -> Self {
103        Self {
104            multibodies: Arena::new(),
105            rb2mb: Coarena::new(),
106            connectivity_graph: InteractionGraph::new(),
107            to_wake_up: HashSet::default(),
108        }
109    }
110
111    /// Iterates through all the multibody joints from this set.
112    pub fn iter(
113        &self,
114    ) -> impl Iterator<
115        Item = (
116            MultibodyJointHandle,
117            &MultibodyLinkId,
118            &Multibody,
119            &MultibodyLink,
120        ),
121    > {
122        self.rb2mb
123            .iter()
124            .filter(|(_, link)| link.id > 0) // The first link of a rigid-body hasn’t been added by the user.
125            .map(|(h, link)| {
126                let mb = &self.multibodies[link.multibody.0];
127                (MultibodyJointHandle(h), link, mb, mb.link(link.id).unwrap())
128            })
129    }
130
131    /// Inserts a new kinematic multibody joint into this set.
132    pub fn insert_kinematic(
133        &mut self,
134        body1: RigidBodyHandle,
135        body2: RigidBodyHandle,
136        data: impl Into<GenericJoint>,
137        wake_up: bool,
138    ) -> Option<MultibodyJointHandle> {
139        self.do_insert(body1, body2, data, true, wake_up)
140    }
141
142    /// Inserts a new multibody joint into this set.
143    pub fn insert(
144        &mut self,
145        body1: RigidBodyHandle,
146        body2: RigidBodyHandle,
147        data: impl Into<GenericJoint>,
148        wake_up: bool,
149    ) -> Option<MultibodyJointHandle> {
150        self.do_insert(body1, body2, data, false, wake_up)
151    }
152
153    /// Inserts a new multibody_joint into this set.
154    #[profiling::function]
155    fn do_insert(
156        &mut self,
157        body1: RigidBodyHandle,
158        body2: RigidBodyHandle,
159        data: impl Into<GenericJoint>,
160        kinematic: bool,
161        wake_up: bool,
162    ) -> Option<MultibodyJointHandle> {
163        let link1 = self.rb2mb.get(body1.0).copied().unwrap_or_else(|| {
164            let mb_handle = self.multibodies.insert(Multibody::with_root(body1, true));
165            MultibodyLinkId {
166                graph_id: self.connectivity_graph.graph.add_node(body1),
167                multibody: MultibodyIndex(mb_handle),
168                id: 0,
169            }
170        });
171
172        let link2 = self.rb2mb.get(body2.0).copied().unwrap_or_else(|| {
173            let mb_handle = self.multibodies.insert(Multibody::with_root(body2, true));
174            MultibodyLinkId {
175                graph_id: self.connectivity_graph.graph.add_node(body2),
176                multibody: MultibodyIndex(mb_handle),
177                id: 0,
178            }
179        });
180
181        if link1.multibody == link2.multibody || link2.id != 0 {
182            // This would introduce an invalid configuration.
183            return None;
184        }
185
186        self.connectivity_graph
187            .graph
188            .add_edge(link1.graph_id, link2.graph_id, ());
189        self.rb2mb.insert(body1.0, link1);
190        self.rb2mb.insert(body2.0, link2);
191
192        let mb2 = self.multibodies.remove(link2.multibody.0).unwrap();
193        let multibody1 = &mut self.multibodies[link1.multibody.0];
194
195        for mb_link2 in mb2.links() {
196            let link = self.rb2mb.get_mut(mb_link2.rigid_body.0).unwrap();
197            link.multibody = link1.multibody;
198            link.id += multibody1.num_links();
199        }
200
201        multibody1.append(mb2, link1.id, MultibodyJoint::new(data.into(), kinematic));
202
203        if wake_up {
204            self.to_wake_up.insert(body1);
205            self.to_wake_up.insert(body2);
206        }
207
208        // Because each rigid-body can only have one parent link,
209        // we can use the second rigid-body’s handle as the multibody_joint’s
210        // handle.
211        Some(MultibodyJointHandle(body2.0))
212    }
213
214    /// Removes a multibody_joint from this set.
215    #[profiling::function]
216    pub fn remove(&mut self, handle: MultibodyJointHandle, wake_up: bool) {
217        if let Some(removed) = self.rb2mb.get(handle.0).copied() {
218            let multibody = self.multibodies.remove(removed.multibody.0).unwrap();
219
220            // Remove the edge from the connectivity graph.
221            if let Some(parent_link) = multibody.link(removed.id).unwrap().parent_id() {
222                let parent_rb = multibody.link(parent_link).unwrap().rigid_body;
223                let parent_graph_id = self.rb2mb.get(parent_rb.0).unwrap().graph_id;
224                self.connectivity_graph
225                    .remove_edge(parent_graph_id, removed.graph_id);
226
227                if wake_up {
228                    self.to_wake_up.insert(RigidBodyHandle(handle.0));
229                    self.to_wake_up.insert(parent_rb);
230                }
231
232                // TODO: remove the node if it no longer has any attached edges?
233
234                // Extract the individual sub-trees generated by this removal.
235                let multibodies = multibody.remove_link(removed.id, true);
236
237                // Update the rb2mb mapping.
238                for multibody in multibodies {
239                    if multibody.num_links() == 1 {
240                        // We don’t have any multibody_joint attached to this body, remove it.
241                        let isolated_link = multibody.link(0).unwrap();
242                        let isolated_graph_id =
243                            self.rb2mb.get(isolated_link.rigid_body.0).unwrap().graph_id;
244                        if let Some(other) = self.connectivity_graph.remove_node(isolated_graph_id)
245                        {
246                            self.rb2mb.get_mut(other.0).unwrap().graph_id = isolated_graph_id;
247                        }
248                    } else {
249                        let mb_id = self.multibodies.insert(multibody);
250                        for link in self.multibodies[mb_id].links() {
251                            let ids = self.rb2mb.get_mut(link.rigid_body.0).unwrap();
252                            ids.multibody = MultibodyIndex(mb_id);
253                            ids.id = link.internal_id;
254                        }
255                    }
256                }
257            }
258        }
259    }
260
261    /// Removes all the multibody_joints from the multibody the given rigid-body is part of.
262    #[profiling::function]
263    pub fn remove_multibody_articulations(&mut self, handle: RigidBodyHandle, wake_up: bool) {
264        if let Some(removed) = self.rb2mb.get(handle.0).copied() {
265            // Remove the multibody.
266            let multibody = self.multibodies.remove(removed.multibody.0).unwrap();
267            for link in multibody.links() {
268                let rb_handle = link.rigid_body;
269
270                if wake_up {
271                    self.to_wake_up.insert(rb_handle);
272                }
273
274                // Remove the rigid-body <-> multibody mapping for this link.
275                let removed = self.rb2mb.remove(rb_handle.0, Default::default()).unwrap();
276                // Remove the node (and all it’s edges) from the connectivity graph.
277                if let Some(other) = self.connectivity_graph.remove_node(removed.graph_id) {
278                    self.rb2mb.get_mut(other.0).unwrap().graph_id = removed.graph_id;
279                }
280            }
281        }
282    }
283
284    /// Removes all the multibody joints attached to a rigid-body.
285    #[profiling::function]
286    pub fn remove_joints_attached_to_rigid_body(&mut self, rb_to_remove: RigidBodyHandle) {
287        // TODO: optimize this.
288        if let Some(link_to_remove) = self.rb2mb.get(rb_to_remove.0).copied() {
289            let mut articulations_to_remove = vec![];
290            for (rb1, rb2, _) in self
291                .connectivity_graph
292                .interactions_with(link_to_remove.graph_id)
293            {
294                // There is a multibody_joint handle is equal to the second rigid-body’s handle.
295                articulations_to_remove.push(MultibodyJointHandle(rb2.0));
296
297                self.to_wake_up.insert(rb1);
298                self.to_wake_up.insert(rb2);
299            }
300
301            for articulation_handle in articulations_to_remove {
302                self.remove(articulation_handle, true);
303            }
304        }
305    }
306
307    /// Returns the link of this multibody attached to the given rigid-body.
308    ///
309    /// Returns `None` if `rb` isn’t part of any rigid-body.
310    pub fn rigid_body_link(&self, rb: RigidBodyHandle) -> Option<&MultibodyLinkId> {
311        self.rb2mb.get(rb.0)
312    }
313
314    /// Gets a reference to a multibody, based on its temporary index.
315    pub fn get_multibody(&self, index: MultibodyIndex) -> Option<&Multibody> {
316        self.multibodies.get(index.0)
317    }
318
319    /// Gets a mutable reference to a multibody, based on its temporary index.
320    /// `MultibodyJointSet`.
321    pub fn get_multibody_mut(&mut self, index: MultibodyIndex) -> Option<&mut Multibody> {
322        // TODO: modification tracking.
323        self.multibodies.get_mut(index.0)
324    }
325
326    /// Gets a mutable reference to a multibody, based on its temporary index.
327    ///
328    /// This method will bypass any modification-detection automatically done by the
329    /// `MultibodyJointSet`.
330    pub fn get_multibody_mut_internal(&mut self, index: MultibodyIndex) -> Option<&mut Multibody> {
331        self.multibodies.get_mut(index.0)
332    }
333
334    /// Gets a reference to the multibody identified by its `handle`.
335    pub fn get(&self, handle: MultibodyJointHandle) -> Option<(&Multibody, usize)> {
336        let link = self.rb2mb.get(handle.0)?;
337        let multibody = self.multibodies.get(link.multibody.0)?;
338        Some((multibody, link.id))
339    }
340
341    /// Gets a mutable reference to the multibody identified by its `handle`.
342    pub fn get_mut(&mut self, handle: MultibodyJointHandle) -> Option<(&mut Multibody, usize)> {
343        let link = self.rb2mb.get(handle.0)?;
344        let multibody = self.multibodies.get_mut(link.multibody.0)?;
345        Some((multibody, link.id))
346    }
347
348    /// Gets a mutable reference to the multibody identified by its `handle`.
349    ///
350    /// This method will bypass any modification-detection automatically done by the MultibodyJointSet.
351    pub fn get_mut_internal(
352        &mut self,
353        handle: MultibodyJointHandle,
354    ) -> Option<(&mut Multibody, usize)> {
355        // TODO: modification tracking?
356        let link = self.rb2mb.get(handle.0)?;
357        let multibody = self.multibodies.get_mut(link.multibody.0)?;
358        Some((multibody, link.id))
359    }
360
361    /// Gets the joint with the given handle without a known generation.
362    ///
363    /// This is useful when you know you want the joint at index `i` but
364    /// don't know what is its current generation number. Generation numbers are
365    /// used to protect from the ABA problem because the joint position `i`
366    /// are recycled between two insertion and a removal.
367    ///
368    /// Using this is discouraged in favor of `self.get(handle)` which does not
369    /// suffer form the ABA problem.
370    pub fn get_unknown_gen(&self, i: u32) -> Option<(&Multibody, usize, MultibodyJointHandle)> {
371        let link = self.rb2mb.get_unknown_gen(i)?;
372        let generation = self.rb2mb.get_gen(i)?;
373        let multibody = self.multibodies.get(link.multibody.0)?;
374        Some((
375            multibody,
376            link.id,
377            MultibodyJointHandle(Index::from_raw_parts(i, generation)),
378        ))
379    }
380
381    /// Returns the joint between two rigid-bodies (if it exists).
382    pub fn joint_between(
383        &self,
384        rb1: RigidBodyHandle,
385        rb2: RigidBodyHandle,
386    ) -> Option<(MultibodyJointHandle, &Multibody, &MultibodyLink)> {
387        let id1 = self.rb2mb.get(rb1.0)?;
388        let id2 = self.rb2mb.get(rb2.0)?;
389
390        // Both bodies must be part of the same multibody.
391        if id1.multibody != id2.multibody {
392            return None;
393        }
394
395        let mb = self.multibodies.get(id1.multibody.0)?;
396
397        // NOTE: if there is a joint between these two bodies, then
398        //       one of the bodies must be the parent of the other.
399        let link1 = mb.link(id1.id)?;
400        let parent1 = link1.parent_id();
401
402        if parent1 == Some(id2.id) {
403            Some((MultibodyJointHandle(rb1.0), mb, link1))
404        } else {
405            let link2 = mb.link(id2.id)?;
406            let parent2 = link2.parent_id();
407
408            if parent2 == Some(id1.id) {
409                Some((MultibodyJointHandle(rb2.0), mb, link2))
410            } else {
411                None
412            }
413        }
414    }
415
416    /// Iterates through all the joints attached to the given rigid-body.
417    #[profiling::function]
418    pub fn attached_joints(
419        &self,
420        rb: RigidBodyHandle,
421    ) -> impl Iterator<Item = (RigidBodyHandle, RigidBodyHandle, MultibodyJointHandle)> + '_ {
422        self.rb2mb
423            .get(rb.0)
424            .into_iter()
425            .flat_map(move |link| self.connectivity_graph.interactions_with(link.graph_id))
426            .map(|inter| {
427                // NOTE: the joint handle is always equal to the handle of the second rigid-body.
428                (inter.0, inter.1, MultibodyJointHandle(inter.1.0))
429            })
430    }
431
432    /// Iterate through the handles of all the rigid-bodies attached to this rigid-body
433    /// by a multibody_joint.
434    pub fn attached_bodies(
435        &self,
436        body: RigidBodyHandle,
437    ) -> impl Iterator<Item = RigidBodyHandle> + '_ {
438        self.rb2mb
439            .get(body.0)
440            .into_iter()
441            .flat_map(move |id| self.connectivity_graph.interactions_with(id.graph_id))
442            .map(move |inter| crate::utils::select_other((inter.0, inter.1), body))
443    }
444
445    /// Iterate through the handles of all the rigid-bodies attached to this rigid-body
446    /// by an enabled multibody_joint.
447    #[profiling::function]
448    pub fn bodies_attached_with_enabled_joint(
449        &self,
450        body: RigidBodyHandle,
451    ) -> impl Iterator<Item = RigidBodyHandle> + '_ {
452        self.attached_bodies(body).filter(move |other| {
453            if let Some((_, _, link)) = self.joint_between(body, *other) {
454                link.joint.data.is_enabled()
455            } else {
456                false
457            }
458        })
459    }
460
461    /// Iterates through all the multibodies on this set.
462    pub fn multibodies(&self) -> impl Iterator<Item = &Multibody> {
463        self.multibodies.iter().map(|e| e.1)
464    }
465}
466
467impl std::ops::Index<MultibodyIndex> for MultibodyJointSet {
468    type Output = Multibody;
469
470    fn index(&self, index: MultibodyIndex) -> &Multibody {
471        &self.multibodies[index.0]
472    }
473}
474
475// impl Index<MultibodyJointHandle> for MultibodyJointSet {
476//     type Output = Multibody;
477//
478//     fn index(&self, index: MultibodyJointHandle) -> &Multibody {
479//         &self.multibodies[index.0]
480//     }
481// }