Fix broken multibody joint removal.

2 files changed, 80 insertions, 9 deletions

diff --git a/src/dynamics/joint/multibody_joint/multibody.rs b/src/dynamics/joint/multibody_joint/multibody.rs
index 5546f9d..6c8344b 100644
--- a/src/dynamics/joint/multibody_joint/multibody.rs
+++ b/src/dynamics/joint/multibody_joint/multibody.rs
@@ -145,6 +145,7 @@ impl Multibody {
         let mut link2mb = vec![usize::MAX; self.links.len()];
         let mut link_id2new_id = vec![usize::MAX; self.links.len()];
 
+        // Split multibody and update the set of links and ndofs.
         for (i, mut link) in self.links.0.into_iter().enumerate() {
             let is_new_root = i == 0
                 || !joint_only && link.parent_internal_id == to_remove
@@ -192,7 +193,14 @@ impl Multibody {
 
                 link.internal_id = i;
                 link.assembly_id = assembly_id;
-                link.parent_internal_id = link_id2new_id[link.parent_internal_id];
+
+                // NOTE: for the root, the current`link.parent_internal_id` is invalid since that
+                //       parent lies in a different multibody now.
+                link.parent_internal_id = if i != 0 {
+                    link_id2new_id[link.parent_internal_id]
+                } else {
+                    0
+                };
                 assembly_id += link_ndofs;
             }
         }
@@ -1363,8 +1371,11 @@ impl IndexSequence {
 #[cfg(test)]
 mod test {
     use super::IndexSequence;
+    use crate::dynamics::{ImpulseJointSet, IslandManager};
     use crate::math::{Real, SPATIAL_DIM};
-    use crate::prelude::{MultibodyJointSet, RevoluteJoint, RigidBodyBuilder, RigidBodySet};
+    use crate::prelude::{
+        ColliderSet, MultibodyJointSet, RevoluteJoint, RigidBodyBuilder, RigidBodySet,
+    };
     use na::{DVector, RowDVector};
 
     #[test]
@@ -1389,6 +1400,63 @@ mod test {
         assert_eq!(joints.get(mb_handle).unwrap().0.ndofs, SPATIAL_DIM + 3);
     }
 
+    #[test]
+    fn test_multibody_remove() {
+        let mut rnd = oorandom::Rand32::new(1234);
+
+        for k in 0..10 {
+            let mut bodies = RigidBodySet::new();
+            let mut multibody_joints = MultibodyJointSet::new();
+            let mut colliders = ColliderSet::new();
+            let mut impulse_joints = ImpulseJointSet::new();
+            let mut islands = IslandManager::new();
+
+            let num_links = 100;
+            let mut handles = vec![];
+
+            for _ in 0..num_links {
+                handles.push(bodies.insert(RigidBodyBuilder::dynamic()));
+            }
+
+            #[cfg(feature = "dim2")]
+            let joint = RevoluteJoint::new();
+            #[cfg(feature = "dim3")]
+            let joint = RevoluteJoint::new(na::Vector::x_axis());
+
+            for i in 0..num_links - 1 {
+                multibody_joints
+                    .insert(handles[i], handles[i + 1], joint, true)
+                    .unwrap();
+            }
+
+            match k {
+                0 => {} // Remove in insertion order.
+                1 => {
+                    // Remove from leaf to root.
+                    handles.reverse();
+                }
+                _ => {
+                    // Shuffle the vector a bit.
+                    // (This test checks multiple shuffle arrangements due to k > 2).
+                    for l in 0..num_links {
+                        handles.swap(l, rnd.rand_range(0..num_links as u32) as usize);
+                    }
+                }
+            }
+
+            for handle in handles {
+                bodies.remove(
+                    handle,
+                    &mut islands,
+                    &mut colliders,
+                    &mut impulse_joints,
+                    &mut multibody_joints,
+                    true,
+                );
+            }
+        }
+    }
+
     fn test_sequence() -> IndexSequence {
         let mut seq = IndexSequence::new();
         seq.remove(2);
diff --git a/src/dynamics/joint/multibody_joint/multibody_joint_set.rs b/src/dynamics/joint/multibody_joint/multibody_joint_set.rs
index 6ffdece..67df159 100644
--- a/src/dynamics/joint/multibody_joint/multibody_joint_set.rs
+++ b/src/dynamics/joint/multibody_joint/multibody_joint_set.rs
@@ -209,7 +209,7 @@ impl MultibodyJointSet {
         Some(MultibodyJointHandle(body2.0))
     }
 
-    /// Removes an multibody_joint from this set.
+    /// Removes a multibody_joint from this set.
     pub fn remove(&mut self, handle: MultibodyJointHandle, wake_up: bool) {
         if let Some(removed) = self.rb2mb.get(handle.0).copied() {
             let multibody = self.multibodies.remove(removed.multibody.0).unwrap();
@@ -217,10 +217,9 @@ impl MultibodyJointSet {
             // Remove the edge from the connectivity graph.
             if let Some(parent_link) = multibody.link(removed.id).unwrap().parent_id() {
                 let parent_rb = multibody.link(parent_link).unwrap().rigid_body;
-                self.connectivity_graph.remove_edge(
-                    self.rb2mb.get(parent_rb.0).unwrap().graph_id,
-                    removed.graph_id,
-                );
+                let parent_graph_id = self.rb2mb.get(parent_rb.0).unwrap().graph_id;
+                self.connectivity_graph
+                    .remove_edge(parent_graph_id, removed.graph_id);
 
                 if wake_up {
                     self.to_wake_up.push(RigidBodyHandle(handle.0));
@@ -236,8 +235,12 @@ impl MultibodyJointSet {
                 for multibody in multibodies {
                     if multibody.num_links() == 1 {
                         // We don’t have any multibody_joint attached to this body, remove it.
-                        if let Some(other) = self.connectivity_graph.remove_node(removed.graph_id) {
-                            self.rb2mb.get_mut(other.0).unwrap().graph_id = removed.graph_id;
+                        let isolated_link = multibody.link(0).unwrap();
+                        let isolated_graph_id =
+                            self.rb2mb.get(isolated_link.rigid_body.0).unwrap().graph_id;
+                        if let Some(other) = self.connectivity_graph.remove_node(isolated_graph_id)
+                        {
+                            self.rb2mb.get_mut(other.0).unwrap().graph_id = isolated_graph_id;
                         }
                     } else {
                         let mb_id = self.multibodies.insert(multibody);