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}
99
100impl MultibodyJointSet {
101 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 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) .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 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 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 #[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 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 Some(MultibodyJointHandle(body2.0))
212 }
213
214 #[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 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 let multibodies = multibody.remove_link(removed.id, true);
236
237 for multibody in multibodies {
239 if multibody.num_links() == 1 {
240 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 #[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 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 let removed = self.rb2mb.remove(rb_handle.0, Default::default()).unwrap();
276 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 #[profiling::function]
286 pub fn remove_joints_attached_to_rigid_body(&mut self, rb_to_remove: RigidBodyHandle) {
287 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 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 pub fn rigid_body_link(&self, rb: RigidBodyHandle) -> Option<&MultibodyLinkId> {
311 self.rb2mb.get(rb.0)
312 }
313
314 pub fn get_multibody(&self, index: MultibodyIndex) -> Option<&Multibody> {
316 self.multibodies.get(index.0)
317 }
318
319 pub fn get_multibody_mut(&mut self, index: MultibodyIndex) -> Option<&mut Multibody> {
322 self.multibodies.get_mut(index.0)
324 }
325
326 pub fn get_multibody_mut_internal(&mut self, index: MultibodyIndex) -> Option<&mut Multibody> {
331 self.multibodies.get_mut(index.0)
332 }
333
334 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 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 pub fn get_mut_internal(
352 &mut self,
353 handle: MultibodyJointHandle,
354 ) -> Option<(&mut Multibody, usize)> {
355 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 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 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 if id1.multibody != id2.multibody {
392 return None;
393 }
394
395 let mb = self.multibodies.get(id1.multibody.0)?;
396
397 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 #[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 (inter.0, inter.1, MultibodyJointHandle(inter.1.0))
429 })
430 }
431
432 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 #[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 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