rapier2d/dynamics/joint/multibody_joint/
multibody_joint_set.rs1use 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#[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#[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 pub fn into_raw_parts(self) -> (u32, u32) {
23 self.0.into_raw_parts()
24 }
25
26 pub fn from_raw_parts(id: u32, generation: u32) -> Self {
28 Self(Index::from_raw_parts(id, generation))
29 }
30
31 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)]
48pub struct MultibodyLinkId {
66 pub(crate) graph_id: RigidBodyGraphIndex,
67 pub multibody: MultibodyIndex,
70 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#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
90#[derive(Clone, Debug)]
91pub struct MultibodyJointSet {
92 pub(crate) multibodies: Arena<Multibody>, pub(crate) rb2mb: Coarena<MultibodyLinkId>,
94 pub(crate) connectivity_graph: InteractionGraph<RigidBodyHandle, ()>,
97 pub(crate) to_wake_up: HashSet<RigidBodyHandle>,
98 pub(crate) to_join: HashSet<(RigidBodyHandle, RigidBodyHandle)>,
100}
101
102impl MultibodyJointSet {
103 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 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) .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 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 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 #[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 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 Some(MultibodyJointHandle(body2.0))
217 }
218
219 #[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 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 let multibodies = multibody.remove_link(removed.id, true);
241
242 for multibody in multibodies {
244 if multibody.num_links() == 1 {
245 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 #[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 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 let removed = self.rb2mb.remove(rb_handle.0, Default::default()).unwrap();
281 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 #[profiling::function]
291 pub fn remove_joints_attached_to_rigid_body(&mut self, rb_to_remove: RigidBodyHandle) {
292 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 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 pub fn rigid_body_link(&self, rb: RigidBodyHandle) -> Option<&MultibodyLinkId> {
316 self.rb2mb.get(rb.0)
317 }
318
319 pub fn get_multibody(&self, index: MultibodyIndex) -> Option<&Multibody> {
321 self.multibodies.get(index.0)
322 }
323
324 pub fn get_multibody_mut(&mut self, index: MultibodyIndex) -> Option<&mut Multibody> {
327 self.multibodies.get_mut(index.0)
329 }
330
331 pub fn get_multibody_mut_internal(&mut self, index: MultibodyIndex) -> Option<&mut Multibody> {
336 self.multibodies.get_mut(index.0)
337 }
338
339 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 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 pub fn get_mut_internal(
357 &mut self,
358 handle: MultibodyJointHandle,
359 ) -> Option<(&mut Multibody, usize)> {
360 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 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 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 if id1.multibody != id2.multibody {
397 return None;
398 }
399
400 let mb = self.multibodies.get(id1.multibody.0)?;
401
402 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 #[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 (inter.0, inter.1, MultibodyJointHandle(inter.1.0))
434 })
435 }
436
437 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 #[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 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