Merge pull request #120 from dimforge/contact_modification

Add the ability to modify the contact points seen by the constraints solver

5 files changed, 100 insertions, 53 deletions

diff --git a/src/dynamics/rigid_body_set.rs b/src/dynamics/rigid_body_set.rs
index 79d5827..f459d4f 100644
--- a/src/dynamics/rigid_body_set.rs
+++ b/src/dynamics/rigid_body_set.rs
@@ -234,13 +234,27 @@ impl RigidBodySet {
         self.bodies.get(handle.0)
     }
 
+    fn mark_as_modified(
+        handle: RigidBodyHandle,
+        rb: &mut RigidBody,
+        modified_bodies: &mut Vec<RigidBodyHandle>,
+        modified_all_bodies: bool,
+    ) {
+        if !modified_all_bodies && !rb.changes.contains(RigidBodyChanges::MODIFIED) {
+            rb.changes = RigidBodyChanges::MODIFIED;
+            modified_bodies.push(handle);
+        }
+    }
+
     /// Gets a mutable reference to the rigid-body with the given handle.
     pub fn get_mut(&mut self, handle: RigidBodyHandle) -> Option<&mut RigidBody> {
         let result = self.bodies.get_mut(handle.0)?;
-        if !self.modified_all_bodies && !result.changes.contains(RigidBodyChanges::MODIFIED) {
-            result.changes = RigidBodyChanges::MODIFIED;
-            self.modified_bodies.push(handle);
-        }
+        Self::mark_as_modified(
+            handle,
+            result,
+            &mut self.modified_bodies,
+            self.modified_all_bodies,
+        );
         Some(result)
     }
 
@@ -300,6 +314,26 @@ impl RigidBodySet {
             .filter_map(move |h| Some((*h, bodies.get(h.0)?)))
     }
 
+    /// Applies the given function on all the active dynamic rigid-bodies
+    /// contained by this set.
+    #[inline(always)]
+    pub fn foreach_active_dynamic_body_mut(
+        &mut self,
+        mut f: impl FnMut(RigidBodyHandle, &mut RigidBody),
+    ) {
+        for handle in &self.active_dynamic_set {
+            if let Some(rb) = self.bodies.get_mut(handle.0) {
+                Self::mark_as_modified(
+                    *handle,
+                    rb,
+                    &mut self.modified_bodies,
+                    self.modified_all_bodies,
+                );
+                f(*handle, rb)
+            }
+        }
+    }
+
     #[inline(always)]
     pub(crate) fn foreach_active_body_mut_internal(
         &mut self,
diff --git a/src/dynamics/solver/velocity_constraint.rs b/src/dynamics/solver/velocity_constraint.rs
index 2cac93d..2de9807 100644
--- a/src/dynamics/solver/velocity_constraint.rs
+++ b/src/dynamics/solver/velocity_constraint.rs
@@ -124,7 +124,7 @@ pub(crate) struct VelocityConstraint {
     pub mj_lambda1: usize,
     pub mj_lambda2: usize,
     pub manifold_id: ContactManifoldIndex,
-    pub manifold_contact_id: usize,
+    pub manifold_contact_id: [u8; MAX_MANIFOLD_POINTS],
     pub num_contacts: u8,
     pub elements: [VelocityConstraintElement; MAX_MANIFOLD_POINTS],
 }
@@ -152,7 +152,7 @@ impl VelocityConstraint {
         let force_dir1 = -manifold.data.normal;
         let warmstart_coeff = manifold.data.warmstart_multiplier * params.warmstart_coeff;
 
-        for (l, manifold_points) in manifold
+        for (_l, manifold_points) in manifold
             .data
             .solver_contacts
             .chunks(MAX_MANIFOLD_POINTS)
@@ -168,7 +168,7 @@ impl VelocityConstraint {
                 mj_lambda1,
                 mj_lambda2,
                 manifold_id,
-                manifold_contact_id: l * MAX_MANIFOLD_POINTS,
+                manifold_contact_id: [0; MAX_MANIFOLD_POINTS],
                 num_contacts: manifold_points.len() as u8,
             };
 
@@ -211,7 +211,7 @@ impl VelocityConstraint {
                 constraint.mj_lambda1 = mj_lambda1;
                 constraint.mj_lambda2 = mj_lambda2;
                 constraint.manifold_id = manifold_id;
-                constraint.manifold_contact_id = l * MAX_MANIFOLD_POINTS;
+                constraint.manifold_contact_id = [0; MAX_MANIFOLD_POINTS];
                 constraint.num_contacts = manifold_points.len() as u8;
             }
 
@@ -224,6 +224,8 @@ impl VelocityConstraint {
                 let vel2 = rb2.linvel + rb2.angvel.gcross(dp2);
 
                 constraint.limit = manifold_point.friction;
+                constraint.manifold_contact_id[k] = manifold_point.contact_id;
+
                 // Normal part.
                 {
                     let gcross1 = rb1
@@ -271,7 +273,8 @@ impl VelocityConstraint {
                                 + rb2.effective_inv_mass
                                 + gcross1.gdot(gcross1)
                                 + gcross2.gdot(gcross2));
-                        let rhs = (vel1 - vel2).dot(&tangents1[j]);
+                        let rhs =
+                            (vel1 - vel2 + manifold_point.tangent_velocity).dot(&tangents1[j]);
                         #[cfg(feature = "dim2")]
                         let impulse = manifold_point.data.tangent_impulse * warmstart_coeff;
                         #[cfg(feature = "dim3")]
@@ -292,7 +295,7 @@ impl VelocityConstraint {
             if push {
                 out_constraints.push(AnyVelocityConstraint::Nongrouped(constraint));
             } else {
-                out_constraints[manifold.data.constraint_index + l] =
+                out_constraints[manifold.data.constraint_index + _l] =
                     AnyVelocityConstraint::Nongrouped(constraint);
             }
         }
@@ -382,19 +385,18 @@ impl VelocityConstraint {
 
     pub fn writeback_impulses(&self, manifolds_all: &mut [&mut ContactManifold]) {
         let manifold = &mut manifolds_all[self.manifold_id];
-        let k_base = self.manifold_contact_id;
 
         for k in 0..self.num_contacts as usize {
-            let active_contacts = &mut manifold.points[..manifold.data.num_active_contacts()];
-            active_contacts[k_base + k].data.impulse = self.elements[k].normal_part.impulse;
+            let contact_id = self.manifold_contact_id[k];
+            let active_contact = &mut manifold.points[contact_id as usize];
+            active_contact.data.impulse = self.elements[k].normal_part.impulse;
             #[cfg(feature = "dim2")]
             {
-                active_contacts[k_base + k].data.tangent_impulse =
-                    self.elements[k].tangent_part[0].impulse;
+                active_contact.data.tangent_impulse = self.elements[k].tangent_part[0].impulse;
             }
             #[cfg(feature = "dim3")]
             {
-                active_contacts[k_base + k].data.tangent_impulse = [
+                active_contact.data.tangent_impulse = [
                     self.elements[k].tangent_part[0].impulse,
                     self.elements[k].tangent_part[1].impulse,
                 ];
diff --git a/src/dynamics/solver/velocity_constraint_wide.rs b/src/dynamics/solver/velocity_constraint_wide.rs
index 3e068e7..97fa261 100644
--- a/src/dynamics/solver/velocity_constraint_wide.rs
+++ b/src/dynamics/solver/velocity_constraint_wide.rs
@@ -55,7 +55,7 @@ pub(crate) struct WVelocityConstraint {
     pub mj_lambda1: [usize; SIMD_WIDTH],
     pub mj_lambda2: [usize; SIMD_WIDTH],
     pub manifold_id: [ContactManifoldIndex; SIMD_WIDTH],
-    pub manifold_contact_id: usize,
+    pub manifold_contact_id: [[u8; SIMD_WIDTH]; MAX_MANIFOLD_POINTS],
 }
 
 impl WVelocityConstraint {
@@ -116,7 +116,7 @@ impl WVelocityConstraint {
                 mj_lambda1,
                 mj_lambda2,
                 manifold_id,
-                manifold_contact_id: l,
+                manifold_contact_id: [[0; SIMD_WIDTH]; MAX_MANIFOLD_POINTS],
                 num_contacts: num_points as u8,
             };
 
@@ -130,6 +130,8 @@ impl WVelocityConstraint {
                 );
                 let point = Point::from(array![|ii| manifold_points[ii][k].point; SIMD_WIDTH]);
                 let dist = SimdReal::from(array![|ii| manifold_points[ii][k].dist; SIMD_WIDTH]);
+                let tangent_velocity =
+                    Vector::from(array![|ii| manifold_points[ii][k].tangent_velocity; SIMD_WIDTH]);
 
                 let impulse =
                     SimdReal::from(array![|ii| manifold_points[ii][k].data.impulse; SIMD_WIDTH]);
@@ -141,6 +143,8 @@ impl WVelocityConstraint {
                 let vel2 = linvel2 + angvel2.gcross(dp2);
 
                 constraint.limit = friction;
+                constraint.manifold_contact_id[k] =
+                    array![|ii| manifold_points[ii][k].contact_id; SIMD_WIDTH];
 
                 // Normal part.
                 {
@@ -179,7 +183,7 @@ impl WVelocityConstraint {
                     let gcross2 = ii2.transform_vector(dp2.gcross(-tangents1[j]));
                     let r = SimdReal::splat(1.0)
                         / (im1 + im2 + gcross1.gdot(gcross1) + gcross2.gdot(gcross2));
-                    let rhs = (vel1 - vel2).dot(&tangents1[j]);
+                    let rhs = (vel1 - vel2 + tangent_velocity).dot(&tangents1[j]);
 
                     constraint.elements[k].tangent_parts[j] = WVelocityConstraintElementPart {
                         gcross1,
@@ -332,17 +336,17 @@ impl WVelocityConstraint {
 
             for ii in 0..SIMD_WIDTH {
                 let manifold = &mut manifolds_all[self.manifold_id[ii]];
-                let k_base = self.manifold_contact_id;
-                let active_contacts = &mut manifold.points[..manifold.data.num_active_contacts()];
-                active_contacts[k_base + k].data.impulse = impulses[ii];
+                let contact_id = self.manifold_contact_id[k][ii];
+                let active_contact = &mut manifold.points[contact_id as usize];
+                active_contact.data.impulse = impulses[ii];
 
                 #[cfg(feature = "dim2")]
                 {
-                    active_contacts[k_base + k].data.tangent_impulse = tangent_impulses[ii];
+                    active_contact.data.tangent_impulse = tangent_impulses[ii];
                 }
                 #[cfg(feature = "dim3")]
                 {
-                    active_contacts[k_base + k].data.tangent_impulse =
+                    active_contact.data.tangent_impulse =
                         [tangent_impulses[ii], bitangent_impulses[ii]];
                 }
             }
diff --git a/src/dynamics/solver/velocity_ground_constraint.rs b/src/dynamics/solver/velocity_ground_constraint.rs
index 88f864a..beb07ed 100644
--- a/src/dynamics/solver/velocity_ground_constraint.rs
+++ b/src/dynamics/solver/velocity_ground_constraint.rs
@@ -49,7 +49,7 @@ pub(crate) struct VelocityGroundConstraint {
     pub limit: Real,
     pub mj_lambda2: usize,
     pub manifold_id: ContactManifoldIndex,
-    pub manifold_contact_id: usize,
+    pub manifold_contact_id: [u8; MAX_MANIFOLD_POINTS],
     pub num_contacts: u8,
     pub elements: [VelocityGroundConstraintElement; MAX_MANIFOLD_POINTS],
 }
@@ -68,17 +68,17 @@ impl VelocityGroundConstraint {
         let mut rb2 = &bodies[manifold.data.body_pair.body2];
         let flipped = !rb2.is_dynamic();
 
-        let force_dir1 = if flipped {
+        let (force_dir1, flipped_multiplier) = if flipped {
             std::mem::swap(&mut rb1, &mut rb2);
-            manifold.data.normal
+            (manifold.data.normal, -1.0)
         } else {
-            -manifold.data.normal
+            (-manifold.data.normal, 1.0)
         };
 
         let mj_lambda2 = rb2.active_set_offset;
         let warmstart_coeff = manifold.data.warmstart_multiplier * params.warmstart_coeff;
 
-        for (l, manifold_points) in manifold
+        for (_l, manifold_points) in manifold
             .data
             .solver_contacts
             .chunks(MAX_MANIFOLD_POINTS)
@@ -92,7 +92,7 @@ impl VelocityGroundConstraint {
                 limit: 0.0,
                 mj_lambda2,
                 manifold_id,
-                manifold_contact_id: l * MAX_MANIFOLD_POINTS,
+                manifold_contact_id: [0; MAX_MANIFOLD_POINTS],
                 num_contacts: manifold_points.len() as u8,
             };
 
@@ -123,7 +123,7 @@ impl VelocityGroundConstraint {
                     .as_nongrouped_ground_mut()
                     .unwrap()
             } else {
-                unreachable!(); // We don't have parallelization on WASM yet, so this is unreachable.
+                unreachable!(); // We don't have parallelization on WASM yet, so this is unreachable.
             };
 
             #[cfg(target_arch = "wasm32")]
@@ -133,7 +133,7 @@ impl VelocityGroundConstraint {
                 constraint.limit = 0.0;
                 constraint.mj_lambda2 = mj_lambda2;
                 constraint.manifold_id = manifold_id;
-                constraint.manifold_contact_id = l * MAX_MANIFOLD_POINTS;
+                constraint.manifold_contact_id = [0; MAX_MANIFOLD_POINTS];
                 constraint.num_contacts = manifold_points.len() as u8;
             }
 
@@ -145,6 +145,7 @@ impl VelocityGroundConstraint {
                 let vel2 = rb2.linvel + rb2.angvel.gcross(dp2);
 
                 constraint.limit = manifold_point.friction;
+                constraint.manifold_contact_id[k] = manifold_point.contact_id;
 
                 // Normal part.
                 {
@@ -178,7 +179,9 @@ impl VelocityGroundConstraint {
                             .effective_world_inv_inertia_sqrt
                             .transform_vector(dp2.gcross(-tangents1[j]));
                         let r = 1.0 / (rb2.effective_inv_mass + gcross2.gdot(gcross2));
-                        let rhs = -vel2.dot(&tangents1[j]) + vel1.dot(&tangents1[j]);
+                        let rhs = (vel1 - vel2
+                            + flipped_multiplier * manifold_point.tangent_velocity)
+                            .dot(&tangents1[j]);
                         #[cfg(feature = "dim2")]
                         let impulse = manifold_points[k].data.tangent_impulse * warmstart_coeff;
                         #[cfg(feature = "dim3")]
@@ -199,7 +202,7 @@ impl VelocityGroundConstraint {
             if push {
                 out_constraints.push(AnyVelocityConstraint::NongroupedGround(constraint));
             } else {
-                out_constraints[manifold.data.constraint_index + l] =
+                out_constraints[manifold.data.constraint_index + _l] =
                     AnyVelocityConstraint::NongroupedGround(constraint);
             }
         }
@@ -267,19 +270,18 @@ impl VelocityGroundConstraint {
     // FIXME: duplicated code. This is exactly the same as in the non-ground velocity constraint.
     pub fn writeback_impulses(&self, manifolds_all: &mut [&mut ContactManifold]) {
         let manifold = &mut manifolds_all[self.manifold_id];
-        let k_base = self.manifold_contact_id;
 
         for k in 0..self.num_contacts as usize {
-            let active_contacts = &mut manifold.points[..manifold.data.num_active_contacts()];
-            active_contacts[k_base + k].data.impulse = self.elements[k].normal_part.impulse;
+            let contact_id = self.manifold_contact_id[k];
+            let active_contact = &mut manifold.points[contact_id as usize];
+            active_contact.data.impulse = self.elements[k].normal_part.impulse;
             #[cfg(feature = "dim2")]
             {
-                active_contacts[k_base + k].data.tangent_impulse =
-                    self.elements[k].tangent_part[0].impulse;
+                active_contact.data.tangent_impulse = self.elements[k].tangent_part[0].impulse;
             }
             #[cfg(feature = "dim3")]
             {
-                active_contacts[k_base + k].data.tangent_impulse = [
+                active_contact.data.tangent_impulse = [
                     self.elements[k].tangent_part[0].impulse,
                     self.elements[k].tangent_part[1].impulse,
                 ];
diff --git a/src/dynamics/solver/velocity_ground_constraint_wide.rs b/src/dynamics/solver/velocity_ground_constraint_wide.rs
index 29ba5e9..6e7216a 100644
--- a/src/dynamics/solver/velocity_ground_constraint_wide.rs
+++ b/src/dynamics/solver/velocity_ground_constraint_wide.rs
@@ -51,7 +51,7 @@ pub(crate) struct WVelocityGroundConstraint {
     pub limit: SimdReal,
     pub mj_lambda2: [usize; SIMD_WIDTH],
     pub manifold_id: [ContactManifoldIndex; SIMD_WIDTH],
-    pub manifold_contact_id: usize,
+    pub manifold_contact_id: [[u8; SIMD_WIDTH]; MAX_MANIFOLD_POINTS],
 }
 
 impl WVelocityGroundConstraint {
@@ -66,15 +66,17 @@ impl WVelocityGroundConstraint {
         let inv_dt = SimdReal::splat(params.inv_dt());
         let mut rbs1 = array![|ii| &bodies[manifolds[ii].data.body_pair.body1]; SIMD_WIDTH];
         let mut rbs2 = array![|ii| &bodies[manifolds[ii].data.body_pair.body2]; SIMD_WIDTH];
-        let mut flipped = [false; SIMD_WIDTH];
+        let mut flipped = [1.0; SIMD_WIDTH];
 
         for ii in 0..SIMD_WIDTH {
             if !rbs2[ii].is_dynamic() {
                 std::mem::swap(&mut rbs1[ii], &mut rbs2[ii]);
-                flipped[ii] = true;
+                flipped[ii] = -1.0;
             }
         }
 
+        let flipped_sign = SimdReal::from(flipped);
+
         let im2 = SimdReal::from(array![|ii| rbs2[ii].effective_inv_mass; SIMD_WIDTH]);
         let ii2: AngularInertia<SimdReal> = AngularInertia::from(
             array![|ii| rbs2[ii].effective_world_inv_inertia_sqrt; SIMD_WIDTH],
@@ -89,9 +91,8 @@ impl WVelocityGroundConstraint {
         let world_com1 = Point::from(array![|ii| rbs1[ii].world_com; SIMD_WIDTH]);
         let world_com2 = Point::from(array![|ii| rbs2[ii].world_com; SIMD_WIDTH]);
 
-        let force_dir1 = Vector::from(
-            array![|ii| if flipped[ii] { manifolds[ii].data.normal } else { -manifolds[ii].data.normal }; SIMD_WIDTH],
-        );
+        let normal1 = Vector::from(array![|ii| manifolds[ii].data.normal; SIMD_WIDTH]);
+        let force_dir1 = normal1 * -flipped_sign;
 
         let mj_lambda2 = array![|ii| rbs2[ii].active_set_offset; SIMD_WIDTH];
 
@@ -111,7 +112,7 @@ impl WVelocityGroundConstraint {
                 limit: SimdReal::splat(0.0),
                 mj_lambda2,
                 manifold_id,
-                manifold_contact_id: l,
+                manifold_contact_id: [[0; SIMD_WIDTH]; MAX_MANIFOLD_POINTS],
                 num_contacts: num_points as u8,
             };
 
@@ -125,6 +126,8 @@ impl WVelocityGroundConstraint {
                 );
                 let point = Point::from(array![|ii| manifold_points[ii][k].point; SIMD_WIDTH]);
                 let dist = SimdReal::from(array![|ii| manifold_points[ii][k].dist; SIMD_WIDTH]);
+                let tangent_velocity =
+                    Vector::from(array![|ii| manifold_points[ii][k].tangent_velocity; SIMD_WIDTH]);
 
                 let impulse =
                     SimdReal::from(array![|ii| manifold_points[ii][k].data.impulse; SIMD_WIDTH]);
@@ -135,6 +138,8 @@ impl WVelocityGroundConstraint {
                 let vel2 = linvel2 + angvel2.gcross(dp2);
 
                 constraint.limit = friction;
+                constraint.manifold_contact_id[k] =
+                    array![|ii| manifold_points[ii][k].contact_id; SIMD_WIDTH];
 
                 // Normal part.
                 {
@@ -168,7 +173,7 @@ impl WVelocityGroundConstraint {
 
                     let gcross2 = ii2.transform_vector(dp2.gcross(-tangents1[j]));
                     let r = SimdReal::splat(1.0) / (im2 + gcross2.gdot(gcross2));
-                    let rhs = -vel2.dot(&tangents1[j]) + vel1.dot(&tangents1[j]);
+                    let rhs = (vel1 - vel2 + tangent_velocity * flipped_sign).dot(&tangents1[j]);
 
                     constraint.elements[k].tangent_parts[j] =
                         WVelocityGroundConstraintElementPart {
@@ -281,17 +286,17 @@ impl WVelocityGroundConstraint {
 
             for ii in 0..SIMD_WIDTH {
                 let manifold = &mut manifolds_all[self.manifold_id[ii]];
-                let k_base = self.manifold_contact_id;
-                let active_contacts = &mut manifold.points[..manifold.data.num_active_contacts()];
-                active_contacts[k_base + k].data.impulse = impulses[ii];
+                let contact_id = self.manifold_contact_id[k][ii];
+                let active_contact = &mut manifold.points[contact_id as usize];
+                active_contact.data.impulse = impulses[ii];
 
                 #[cfg(feature = "dim2")]
                 {
-                    active_contacts[k_base + k].data.tangent_impulse = tangent_impulses[ii];
+                    active_contact.data.tangent_impulse = tangent_impulses[ii];
                 }
                 #[cfg(feature = "dim3")]
                 {
-                    active_contacts[k_base + k].data.tangent_impulse =
+                    active_contact.data.tangent_impulse =
                         [tangent_impulses[ii], bitangent_impulses[ii]];
                 }
             }