Skip to main content

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