Fix body status modification.

2 files changed, 119 insertions, 54 deletions

diff --git a/src/dynamics/rigid_body.rs b/src/dynamics/rigid_body.rs
index ebf71de..47eea8a 100644
--- a/src/dynamics/rigid_body.rs
+++ b/src/dynamics/rigid_body.rs
@@ -182,10 +182,13 @@ impl RigidBody {
         self.body_status
     }
 
-    // pub fn set_body_status(&mut self, status: BodyStatus) {
-    //     self.changes.insert(RigidBodyChanges::BODY_STATUS);
-    //     self.body_status = status;
-    // }
+    /// Sets the status of this rigid-body.
+    pub fn set_body_status(&mut self, status: BodyStatus) {
+        if status != self.body_status {
+            self.changes.insert(RigidBodyChanges::BODY_STATUS);
+            self.body_status = status;
+        }
+    }
 
     /// The mass properties of this rigid-body.
     #[inline]
@@ -948,7 +951,7 @@ impl RigidBodyBuilder {
     /// equal to the sum of this additional mass and the mass computed from the colliders
     /// (with non-zero densities) attached to this rigid-body.
     #[deprecated(note = "renamed to `additional_mass`.")]
-    pub fn mass(mut self, mass: Real) -> Self {
+    pub fn mass(self, mass: Real) -> Self {
         self.additional_mass(mass)
     }
 
@@ -993,7 +996,7 @@ impl RigidBodyBuilder {
     /// Sets the principal angular inertia of this rigid-body.
     #[cfg(feature = "dim3")]
     #[deprecated(note = "renamed to `additional_principal_angular_inertia`.")]
-    pub fn principal_angular_inertia(mut self, inertia: AngVector<Real>) -> Self {
+    pub fn principal_angular_inertia(self, inertia: AngVector<Real>) -> Self {
         self.additional_principal_angular_inertia(inertia)
     }
 
diff --git a/src/dynamics/rigid_body_set.rs b/src/dynamics/rigid_body_set.rs
index 2ae298f..5a85ada 100644
--- a/src/dynamics/rigid_body_set.rs
+++ b/src/dynamics/rigid_body_set.rs
@@ -2,7 +2,7 @@
 use rayon::prelude::*;
 
 use crate::data::arena::Arena;
-use crate::dynamics::{Joint, JointSet, RigidBody, RigidBodyChanges};
+use crate::dynamics::{BodyStatus, Joint, JointSet, RigidBody, RigidBodyChanges};
 use crate::geometry::{ColliderSet, InteractionGraph, NarrowPhase};
 use parry::partitioning::IndexedData;
 use std::ops::{Index, IndexMut};
@@ -453,71 +453,133 @@ impl RigidBodySet {
 
     // Utility function to avoid some borrowing issue in the `maintain` method.
     fn maintain_one(
+        bodies: &mut Arena<RigidBody>,
         colliders: &mut ColliderSet,
         handle: RigidBodyHandle,
-        rb: &mut RigidBody,
         modified_inactive_set: &mut Vec<RigidBodyHandle>,
         active_kinematic_set: &mut Vec<RigidBodyHandle>,
         active_dynamic_set: &mut Vec<RigidBodyHandle>,
     ) {
-        // Update the positions of the colliders.
-        if rb.changes.contains(RigidBodyChanges::POSITION)
-            || rb.changes.contains(RigidBodyChanges::COLLIDERS)
-        {
-            rb.update_colliders_positions(colliders);
-
-            if rb.is_static() {
-                modified_inactive_set.push(handle);
+        enum FinalAction {
+            UpdateActiveKinematicSetId,
+            UpdateActiveDynamicSetId,
+        }
+
+        if let Some(rb) = bodies.get_mut(handle.0) {
+            let mut final_action = None;
+
+            // The body's status changed. We need to make sure
+            // it is on the correct active set.
+            if rb.changes.contains(RigidBodyChanges::BODY_STATUS) {
+                match rb.body_status() {
+                    BodyStatus::Dynamic => {
+                        // Remove from the active kinematic set if it was there.
+                        if active_kinematic_set.get(rb.active_set_id) == Some(&handle) {
+                            active_kinematic_set.swap_remove(rb.active_set_id);
+                            final_action =
+                                Some((FinalAction::UpdateActiveKinematicSetId, rb.active_set_id));
+                        }
+
+                        // Add to the active dynamic set.
+                        rb.wake_up(true);
+                        // Make sure the sleep change flag is set (even if for some
+                        // reasons the rigid-body was already awake) to make
+                        // sure the code handling sleeping change adds the body to
+                        // the active_dynamic_set.
+                        rb.changes.set(RigidBodyChanges::SLEEP, true);
+                    }
+                    BodyStatus::Kinematic => {
+                        // Remove from the active dynamic set if it was there.
+                        if active_dynamic_set.get(rb.active_set_id) == Some(&handle) {
+                            active_dynamic_set.swap_remove(rb.active_set_id);
+                            final_action =
+                                Some((FinalAction::UpdateActiveDynamicSetId, rb.active_set_id));
+                        }
+
+                        // Add to the active kinematic set.
+                        if active_kinematic_set.get(rb.active_set_id) != Some(&handle) {
+                            rb.active_set_id = active_kinematic_set.len();
+                            active_kinematic_set.push(handle);
+                        }
+                    }
+                    BodyStatus::Static => {}
+                }
+            }
+
+            // Update the positions of the colliders.
+            if rb.changes.contains(RigidBodyChanges::POSITION)
+                || rb.changes.contains(RigidBodyChanges::COLLIDERS)
+            {
+                rb.update_colliders_positions(colliders);
+
+                if rb.is_static() {
+                    modified_inactive_set.push(handle);
+                }
+
+                if rb.is_kinematic() && active_kinematic_set.get(rb.active_set_id) != Some(&handle)
+                {
+                    rb.active_set_id = active_kinematic_set.len();
+                    active_kinematic_set.push(handle);
+                }
             }
 
-            if rb.is_kinematic() && active_kinematic_set.get(rb.active_set_id) != Some(&handle) {
-                rb.active_set_id = active_kinematic_set.len();
-                active_kinematic_set.push(handle);
+            // Push the body to the active set if it is not
+            // sleeping and if it is not already inside of the active set.
+            if rb.changes.contains(RigidBodyChanges::SLEEP)
+                && !rb.is_sleeping() // May happen if the body was put to sleep manually.
+                && rb.is_dynamic() // Only dynamic bodies are in the active dynamic set.
+                && active_dynamic_set.get(rb.active_set_id) != Some(&handle)
+            {
+                rb.active_set_id = active_dynamic_set.len(); // This will handle the case where the activation_channel contains duplicates.
+                active_dynamic_set.push(handle);
             }
-        }
 
-        // Push the body to the active set if it is not
-        // sleeping and if it is not already inside of the active set.
-        if rb.changes.contains(RigidBodyChanges::SLEEP)
-            && !rb.is_sleeping() // May happen if the body was put to sleep manually.
-            && rb.is_dynamic() // Only dynamic bodies are in the active dynamic set.
-            && active_dynamic_set.get(rb.active_set_id) != Some(&handle)
-        {
-            rb.active_set_id = active_dynamic_set.len(); // This will handle the case where the activation_channel contains duplicates.
-            active_dynamic_set.push(handle);
-        }
+            rb.changes = RigidBodyChanges::empty();
 
-        rb.changes = RigidBodyChanges::empty();
+            // Adjust some ids, if needed.
+            if let Some((action, id)) = final_action {
+                let active_set = match action {
+                    FinalAction::UpdateActiveKinematicSetId => active_kinematic_set,
+                    FinalAction::UpdateActiveDynamicSetId => active_dynamic_set,
+                };
+
+                if id < active_set.len() {
+                    if let Some(rb2) = bodies.get_mut(active_set[id].0) {
+                        rb2.active_set_id = id;
+                    }
+                }
+            }
+        }
     }
 
     pub(crate) fn handle_user_changes(&mut self, colliders: &mut ColliderSet) {
         if self.modified_all_bodies {
-            for (handle, rb) in self.bodies.iter_mut() {
-                Self::maintain_one(
-                    colliders,
-                    RigidBodyHandle(handle),
-                    rb,
-                    &mut self.modified_inactive_set,
-                    &mut self.active_kinematic_set,
-                    &mut self.active_dynamic_set,
-                )
+            // Unfortunately, we have to push all the bodies to `modified_bodies`
+            // instead of just calling `maintain_one` on each element i
+            // `self.bodies.iter_mut()` because otherwise it would be difficult to
+            // handle the final  change of active_set_id in Self::maintain_one
+            // (because it has  to modify another rigid-body because of the swap-remove.
+            // So this causes borrowing problems if we do this while iterating
+            // through self.bodies.iter_mut()).
+            for (handle, _) in self.bodies.iter_mut() {
+                self.modified_bodies.push(RigidBodyHandle(handle));
             }
+        }
+
+        for handle in self.modified_bodies.drain(..) {
+            Self::maintain_one(
+                &mut self.bodies,
+                colliders,
+                handle,
+                &mut self.modified_inactive_set,
+                &mut self.active_kinematic_set,
+                &mut self.active_dynamic_set,
+            )
+        }
 
-            self.modified_bodies.clear();
+        if self.modified_all_bodies {
+            self.modified_bodies.shrink_to_fit(); // save some memory.
             self.modified_all_bodies = false;
-        } else {
-            for handle in self.modified_bodies.drain(..) {
-                if let Some(rb) = self.bodies.get_mut(handle.0) {
-                    Self::maintain_one(
-                        colliders,
-                        handle,
-                        rb,
-                        &mut self.modified_inactive_set,
-                        &mut self.active_kinematic_set,
-                        &mut self.active_dynamic_set,
-                    )
-                }
-            }
         }
     }