Skip to main content

rapier2d/dynamics/joint/multibody_joint/
multibody_joint_set.rs

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