rapier2d/dynamics/joint/multibody_joint/
multibody_joint_set.rs1use 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)]
13pub struct MultibodyLinkId {
31 pub(crate) graph_id: RigidBodyGraphIndex,
32 pub multibody: MultibodyIndex,
35 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#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
55#[derive(Clone, Debug)]
56pub struct MultibodyJointSet {
57 pub(crate) multibodies: Arena<Multibody>, pub(crate) rb2mb: Coarena<MultibodyLinkId>,
59 pub(crate) connectivity_graph: InteractionGraph<RigidBodyHandle, ()>,
62 pub(crate) to_wake_up: HashSet<RigidBodyHandle>,
63 pub(crate) to_join: HashSet<(RigidBodyHandle, RigidBodyHandle)>,
65 #[cfg_attr(feature = "serde-serialize", serde(skip))]
70 pub(crate) island_chain_events: Vec<MultibodyIndex>,
71 pub(crate) topology_epoch: u32,
76}
77
78impl MultibodyJointSet {
79 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 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) .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 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 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 #[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 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 self.island_chain_events.push(link2.multibody);
193 self.island_chain_events.push(link1.multibody);
194
195 Some(MultibodyJointHandle(body2.0))
199 }
200
201 #[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 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 let multibodies = multibody.remove_link(removed.id, true);
225
226 for multibody in multibodies {
228 if multibody.num_links() == 1 {
229 let isolated_link = multibody.link(0).unwrap();
231
232 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 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 #[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 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 let removed = self.rb2mb.remove(rb_handle.0, Default::default()).unwrap();
274 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 #[profiling::function]
284 pub fn remove_joints_attached_to_rigid_body(&mut self, rb_to_remove: RigidBodyHandle) {
285 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 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 pub fn rigid_body_link(&self, rb: RigidBodyHandle) -> Option<&MultibodyLinkId> {
309 self.rb2mb.get(rb.0)
310 }
311
312 pub fn get_multibody(&self, index: MultibodyIndex) -> Option<&Multibody> {
314 self.multibodies.get(index.0)
315 }
316
317 pub fn get_multibody_mut(&mut self, index: MultibodyIndex) -> Option<&mut Multibody> {
320 self.multibodies.get_mut(index.0)
322 }
323
324 pub fn get_multibody_mut_internal(&mut self, index: MultibodyIndex) -> Option<&mut Multibody> {
329 self.multibodies.get_mut(index.0)
330 }
331
332 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 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 pub fn get_mut_internal(
350 &mut self,
351 handle: MultibodyJointHandle,
352 ) -> Option<(&mut Multibody, usize)> {
353 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 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 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 if id1.multibody != id2.multibody {
390 return None;
391 }
392
393 let mb = self.multibodies.get(id1.multibody.0)?;
394
395 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 #[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 (inter.0, inter.1, MultibodyJointHandle(inter.1.0))
427 })
428 }
429
430 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 #[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 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