Merge pull request #579 from dimforge/joints-improvements

Feat: implement a "small-steps" velocity-based constraints solver + joint improvements

37 files changed, 5572 insertions, 3942 deletions

diff --git a/src/dynamics/solver/categorization.rs b/src/dynamics/solver/categorization.rs
index d110971..cef9c9e 100644
--- a/src/dynamics/solver/categorization.rs
+++ b/src/dynamics/solver/categorization.rs
@@ -6,10 +6,10 @@ pub(crate) fn categorize_contacts(
     multibody_joints: &MultibodyJointSet,
     manifolds: &[&mut ContactManifold],
     manifold_indices: &[ContactManifoldIndex],
-    out_ground: &mut Vec<ContactManifoldIndex>,
-    out_not_ground: &mut Vec<ContactManifoldIndex>,
-    out_generic_ground: &mut Vec<ContactManifoldIndex>,
-    out_generic_not_ground: &mut Vec<ContactManifoldIndex>,
+    out_one_body: &mut Vec<ContactManifoldIndex>,
+    out_two_body: &mut Vec<ContactManifoldIndex>,
+    out_generic_one_body: &mut Vec<ContactManifoldIndex>,
+    out_generic_two_body: &mut Vec<ContactManifoldIndex>,
 ) {
     for manifold_i in manifold_indices {
         let manifold = &manifolds[*manifold_i];
@@ -26,14 +26,14 @@ pub(crate) fn categorize_contacts(
                 .is_some()
         {
             if manifold.data.relative_dominance != 0 {
-                out_generic_ground.push(*manifold_i);
+                out_generic_one_body.push(*manifold_i);
             } else {
-                out_generic_not_ground.push(*manifold_i);
+                out_generic_two_body.push(*manifold_i);
             }
         } else if manifold.data.relative_dominance != 0 {
-            out_ground.push(*manifold_i)
+            out_one_body.push(*manifold_i)
         } else {
-            out_not_ground.push(*manifold_i)
+            out_two_body.push(*manifold_i)
         }
     }
 }
@@ -43,10 +43,10 @@ pub(crate) fn categorize_joints(
     multibody_joints: &MultibodyJointSet,
     impulse_joints: &[JointGraphEdge],
     joint_indices: &[JointIndex],
-    ground_joints: &mut Vec<JointIndex>,
-    nonground_joints: &mut Vec<JointIndex>,
-    generic_ground_joints: &mut Vec<JointIndex>,
-    generic_nonground_joints: &mut Vec<JointIndex>,
+    one_body_joints: &mut Vec<JointIndex>,
+    two_body_joints: &mut Vec<JointIndex>,
+    generic_one_body_joints: &mut Vec<JointIndex>,
+    generic_two_body_joints: &mut Vec<JointIndex>,
 ) {
     for joint_i in joint_indices {
         let joint = &impulse_joints[*joint_i].weight;
@@ -57,14 +57,14 @@ pub(crate) fn categorize_joints(
             || multibody_joints.rigid_body_link(joint.body2).is_some()
         {
             if !rb1.is_dynamic() || !rb2.is_dynamic() {
-                generic_ground_joints.push(*joint_i);
+                generic_one_body_joints.push(*joint_i);
             } else {
-                generic_nonground_joints.push(*joint_i);
+                generic_two_body_joints.push(*joint_i);
             }
         } else if !rb1.is_dynamic() || !rb2.is_dynamic() {
-            ground_joints.push(*joint_i);
+            one_body_joints.push(*joint_i);
         } else {
-            nonground_joints.push(*joint_i);
+            two_body_joints.push(*joint_i);
         }
     }
 }
diff --git a/src/dynamics/solver/contact_constraint/contact_constraints_set.rs b/src/dynamics/solver/contact_constraint/contact_constraints_set.rs
new file mode 100644
index 0000000..4d02270
--- /dev/null
+++ b/src/dynamics/solver/contact_constraint/contact_constraints_set.rs
@@ -0,0 +1,528 @@
+use crate::dynamics::solver::categorization::categorize_contacts;
+use crate::dynamics::solver::contact_constraint::{
+    GenericOneBodyConstraint, GenericOneBodyConstraintBuilder, GenericTwoBodyConstraint,
+    GenericTwoBodyConstraintBuilder, OneBodyConstraint, OneBodyConstraintBuilder,
+    TwoBodyConstraint, TwoBodyConstraintBuilder,
+};
+use crate::dynamics::solver::solver_body::SolverBody;
+use crate::dynamics::solver::solver_vel::SolverVel;
+use crate::dynamics::solver::{reset_buffer, ConstraintTypes, SolverConstraintsSet};
+use crate::dynamics::{
+    ImpulseJoint, IntegrationParameters, IslandManager, JointAxesMask, MultibodyJointSet,
+    RigidBodySet,
+};
+use crate::geometry::{ContactManifold, ContactManifoldIndex};
+use crate::math::{Real, MAX_MANIFOLD_POINTS};
+use na::DVector;
+use parry::math::DIM;
+
+#[cfg(feature = "simd-is-enabled")]
+use {
+    crate::dynamics::solver::contact_constraint::{
+        OneBodyConstraintSimd, SimdOneBodyConstraintBuilder, TwoBodyConstraintBuilderSimd,
+        TwoBodyConstraintSimd,
+    },
+    crate::math::SIMD_WIDTH,
+};
+
+pub struct ConstraintsCounts {
+    pub num_constraints: usize,
+    pub num_jacobian_lines: usize,
+}
+
+impl ConstraintsCounts {
+    pub fn from_contacts(manifold: &ContactManifold) -> Self {
+        let rest = manifold.data.solver_contacts.len() % MAX_MANIFOLD_POINTS != 0;
+        Self {
+            num_constraints: manifold.data.solver_contacts.len() / MAX_MANIFOLD_POINTS
+                + rest as usize,
+            num_jacobian_lines: manifold.data.solver_contacts.len() * DIM,
+        }
+    }
+
+    pub fn from_joint(joint: &ImpulseJoint) -> Self {
+        let joint = &joint.data;
+        let locked_axes = joint.locked_axes.bits();
+        let motor_axes = joint.motor_axes.bits() & !locked_axes;
+        let limit_axes = joint.limit_axes.bits() & !locked_axes;
+        let coupled_axes = joint.coupled_axes.bits();
+
+        let num_constraints = (motor_axes & !coupled_axes).count_ones() as usize
+            + ((motor_axes & coupled_axes) & JointAxesMask::ANG_AXES.bits() != 0) as usize
+            + ((motor_axes & coupled_axes) & JointAxesMask::LIN_AXES.bits() != 0) as usize
+            + locked_axes.count_ones() as usize
+            + (limit_axes & !coupled_axes).count_ones() as usize
+            + ((limit_axes & coupled_axes) & JointAxesMask::ANG_AXES.bits() != 0) as usize
+            + ((limit_axes & coupled_axes) & JointAxesMask::LIN_AXES.bits() != 0) as usize;
+        Self {
+            num_constraints,
+            num_jacobian_lines: num_constraints,
+        }
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub(crate) struct ContactConstraintTypes;
+
+impl ConstraintTypes for ContactConstraintTypes {
+    type OneBody = OneBodyConstraint;
+    type TwoBodies = TwoBodyConstraint;
+    type GenericOneBody = GenericOneBodyConstraint;
+    type GenericTwoBodies = GenericTwoBodyConstraint;
+    #[cfg(feature = "simd-is-enabled")]
+    type SimdOneBody = OneBodyConstraintSimd;
+    #[cfg(feature = "simd-is-enabled")]
+    type SimdTwoBodies = TwoBodyConstraintSimd;
+
+    type BuilderOneBody = OneBodyConstraintBuilder;
+    type BuilderTwoBodies = TwoBodyConstraintBuilder;
+    type GenericBuilderOneBody = GenericOneBodyConstraintBuilder;
+    type GenericBuilderTwoBodies = GenericTwoBodyConstraintBuilder;
+    #[cfg(feature = "simd-is-enabled")]
+    type SimdBuilderOneBody = SimdOneBodyConstraintBuilder;
+    #[cfg(feature = "simd-is-enabled")]
+    type SimdBuilderTwoBodies = TwoBodyConstraintBuilderSimd;
+}
+
+pub type ContactConstraintsSet = SolverConstraintsSet<ContactConstraintTypes>;
+
+impl ContactConstraintsSet {
+    pub fn init_constraint_groups(
+        &mut self,
+        island_id: usize,
+        islands: &IslandManager,
+        bodies: &RigidBodySet,
+        multibody_joints: &MultibodyJointSet,
+        manifolds: &[&mut ContactManifold],
+        manifold_indices: &[ContactManifoldIndex],
+    ) {
+        self.two_body_interactions.clear();
+        self.one_body_interactions.clear();
+        self.generic_two_body_interactions.clear();
+        self.generic_one_body_interactions.clear();
+
+        categorize_contacts(
+            bodies,
+            multibody_joints,
+            manifolds,
+            manifold_indices,
+            &mut self.one_body_interactions,
+            &mut self.two_body_interactions,
+            &mut self.generic_one_body_interactions,
+            &mut self.generic_two_body_interactions,
+        );
+
+        self.interaction_groups.clear_groups();
+        self.interaction_groups.group_manifolds(
+            island_id,
+            islands,
+            bodies,
+            manifolds,
+            &self.two_body_interactions,
+        );
+
+        self.one_body_interaction_groups.clear_groups();
+        self.one_body_interaction_groups.group_manifolds(
+            island_id,
+            islands,
+            bodies,
+            manifolds,
+            &self.one_body_interactions,
+        );
+
+        // NOTE: uncomment this do disable SIMD contact resolution.
+        //        self.interaction_groups
+        //            .nongrouped_interactions
+        //            .append(&mut self.interaction_groups.simd_interactions);
+        //        self.one_body_interaction_groups
+        //            .nongrouped_interactions
+        //            .append(&mut self.one_body_interaction_groups.simd_interactions);
+    }
+
+    pub fn init(
+        &mut self,
+        island_id: usize,
+        islands: &IslandManager,
+        bodies: &RigidBodySet,
+        multibody_joints: &MultibodyJointSet,
+        manifolds: &[&mut ContactManifold],
+        manifold_indices: &[ContactManifoldIndex],
+    ) {
+        self.clear_constraints();
+        self.clear_builders();
+
+        self.init_constraint_groups(
+            island_id,
+            islands,
+            bodies,
+            multibody_joints,
+            manifolds,
+            manifold_indices,
+        );
+
+        let mut jacobian_id = 0;
+
+        #[cfg(feature = "simd-is-enabled")]
+        {
+            self.simd_compute_constraints(bodies, manifolds);
+        }
+        self.compute_constraints(bodies, manifolds);
+        self.compute_generic_constraints(bodies, multibody_joints, manifolds, &mut jacobian_id);
+
+        #[cfg(feature = "simd-is-enabled")]
+        {
+            self.simd_compute_one_body_constraints(bodies, manifolds);
+        }
+        self.compute_one_body_constraints(bodies, manifolds);
+        self.compute_generic_one_body_constraints(
+            bodies,
+            multibody_joints,
+            manifolds,
+            &mut jacobian_id,
+        );
+    }
+
+    #[cfg(feature = "simd-is-enabled")]
+    fn simd_compute_constraints(
+        &mut self,
+        bodies: &RigidBodySet,
+        manifolds_all: &[&mut ContactManifold],
+    ) {
+        let total_num_constraints = self
+            .interaction_groups
+            .simd_interactions
+            .chunks_exact(SIMD_WIDTH)
+            .map(|i| ConstraintsCounts::from_contacts(&manifolds_all[i[0]]).num_constraints)
+            .sum();
+
+        unsafe {
+            reset_buffer(
+                &mut self.simd_velocity_constraints_builder,
+                total_num_constraints,
+            );
+            reset_buffer(&mut self.simd_velocity_constraints, total_num_constraints);
+        }
+
+        let mut curr_start = 0;
+
+        for manifolds_i in self
+            .interaction_groups
+            .simd_interactions
+            .chunks_exact(SIMD_WIDTH)
+        {
+            let num_to_add =
+                ConstraintsCounts::from_contacts(&manifolds_all[manifolds_i[0]]).num_constraints;
+            let manifold_id = gather![|ii| manifolds_i[ii]];
+            let manifolds = gather![|ii| &*manifolds_all[manifolds_i[ii]]];
+
+            TwoBodyConstraintBuilderSimd::generate(
+                manifold_id,
+                manifolds,
+                bodies,
+                &mut self.simd_velocity_constraints_builder[curr_start..],
+                &mut self.simd_velocity_constraints[curr_start..],
+            );
+
+            curr_start += num_to_add;
+        }
+
+        assert_eq!(curr_start, total_num_constraints);
+    }
+
+    fn compute_constraints(
+        &mut self,
+        bodies: &RigidBodySet,
+        manifolds_all: &[&mut ContactManifold],
+    ) {
+        let total_num_constraints = self
+            .interaction_groups
+            .nongrouped_interactions
+            .iter()
+            .map(|i| ConstraintsCounts::from_contacts(&manifolds_all[*i]).num_constraints)
+            .sum();
+
+        unsafe {
+            reset_buffer(
+                &mut self.velocity_constraints_builder,
+                total_num_constraints,
+            );
+            reset_buffer(&mut self.velocity_constraints, total_num_constraints);
+        }
+
+        let mut curr_start = 0;
+
+        for manifold_i in &self.interaction_groups.nongrouped_interactions {
+            let manifold = &manifolds_all[*manifold_i];
+            let num_to_add = ConstraintsCounts::from_contacts(manifold).num_constraints;
+
+            TwoBodyConstraintBuilder::generate(
+                *manifold_i,
+                manifold,
+                bodies,
+                &mut self.velocity_constraints_builder[curr_start..],
+                &mut self.velocity_constraints[curr_start..],
+            );
+
+            curr_start += num_to_add;
+        }
+        assert_eq!(curr_start, total_num_constraints);
+    }
+
+    fn compute_generic_constraints(
+        &mut self,
+        bodies: &RigidBodySet,
+        multibody_joints: &MultibodyJointSet,
+        manifolds_all: &[&mut ContactManifold],
+        jacobian_id: &mut usize,
+    ) {
+        let total_num_constraints = self
+            .generic_two_body_interactions
+            .iter()
+            .map(|i| ConstraintsCounts::from_contacts(&manifolds_all[*i]).num_constraints)
+            .sum();
+
+        self.generic_velocity_constraints_builder.resize(
+            total_num_constraints,
+            GenericTwoBodyConstraintBuilder::invalid(),
+        );
+        self.generic_velocity_constraints
+            .resize(total_num_constraints, GenericTwoBodyConstraint::invalid());
+
+        let mut curr_start = 0;
+
+        for manifold_i in &self.generic_two_body_interactions {
+            let manifold = &manifolds_all[*manifold_i];
+            let num_to_add = ConstraintsCounts::from_contacts(manifold).num_constraints;
+
+            GenericTwoBodyConstraintBuilder::generate(
+                *manifold_i,
+                manifold,
+                bodies,
+                multibody_joints,
+                &mut self.generic_velocity_constraints_builder[curr_start..],
+                &mut self.generic_velocity_constraints[curr_start..],
+                &mut self.generic_jacobians,
+                jacobian_id,
+            );
+
+            curr_start += num_to_add;
+        }
+
+        assert_eq!(curr_start, total_num_constraints);
+    }
+
+    fn compute_generic_one_body_constraints(
+        &mut self,
+        bodies: &RigidBodySet,
+        multibody_joints: &MultibodyJointSet,
+        manifolds_all: &[&mut ContactManifold],
+        jacobian_id: &mut usize,
+    ) {
+        let total_num_constraints = self
+            .generic_one_body_interactions
+            .iter()
+            .map(|i| ConstraintsCounts::from_contacts(&manifolds_all[*i]).num_constraints)
+            .sum();
+        self.generic_velocity_one_body_constraints_builder.resize(
+            total_num_constraints,
+            GenericOneBodyConstraintBuilder::invalid(),
+        );
+        self.generic_velocity_one_body_constraints
+            .resize(total_num_constraints, GenericOneBodyConstraint::invalid());
+
+        let mut curr_start = 0;
+
+        for manifold_i in &self.generic_one_body_interactions {
+            let manifold = &manifolds_all[*manifold_i];
+            let num_to_add = ConstraintsCounts::from_contacts(manifold).num_constraints;
+
+            GenericOneBodyConstraintBuilder::generate(
+                *manifold_i,
+                manifold,
+                bodies,
+                multibody_joints,
+                &mut self.generic_velocity_one_body_constraints_builder[curr_start..],
+                &mut self.generic_velocity_one_body_constraints[curr_start..],
+                &mut self.generic_jacobians,
+                jacobian_id,
+            );
+
+            curr_start += num_to_add;
+        }
+        assert_eq!(curr_start, total_num_constraints);
+    }
+
+    #[cfg(feature = "simd-is-enabled")]
+    fn simd_compute_one_body_constraints(
+        &mut self,
+        bodies: &RigidBodySet,
+        manifolds_all: &[&mut ContactManifold],
+    ) {
+        let total_num_constraints = self
+            .one_body_interaction_groups
+            .simd_interactions
+            .chunks_exact(SIMD_WIDTH)
+            .map(|i| ConstraintsCounts::from_contacts(&manifolds_all[i[0]]).num_constraints)
+            .sum();
+
+        unsafe {
+            reset_buffer(
+                &mut self.simd_velocity_one_body_constraints_builder,
+                total_num_constraints,
+            );
+            reset_buffer(
+                &mut self.simd_velocity_one_body_constraints,
+                total_num_constraints,
+            );
+        }
+
+        let mut curr_start = 0;
+
+        for manifolds_i in self
+            .one_body_interaction_groups
+            .simd_interactions
+            .chunks_exact(SIMD_WIDTH)
+        {
+            let num_to_add =
+                ConstraintsCounts::from_contacts(&manifolds_all[manifolds_i[0]]).num_constraints;
+            let manifold_id = gather![|ii| manifolds_i[ii]];
+            let manifolds = gather![|ii| &*manifolds_all[manifolds_i[ii]]];
+            SimdOneBodyConstraintBuilder::generate(
+                manifold_id,
+                manifolds,
+                bodies,
+                &mut self.simd_velocity_one_body_constraints_builder[curr_start..],
+                &mut self.simd_velocity_one_body_constraints[curr_start..],
+            );
+            curr_start += num_to_add;
+        }
+        assert_eq!(curr_start, total_num_constraints);
+    }
+
+    fn compute_one_body_constraints(
+        &mut self,
+        bodies: &RigidBodySet,
+        manifolds_all: &[&mut ContactManifold],
+    ) {
+        let total_num_constraints = self
+            .one_body_interaction_groups
+            .nongrouped_interactions
+            .iter()
+            .map(|i| ConstraintsCounts::from_contacts(&manifolds_all[*i]).num_constraints)
+            .sum();
+
+        unsafe {
+            reset_buffer(
+                &mut self.velocity_one_body_constraints_builder,
+                total_num_constraints,
+            );
+            reset_buffer(
+                &mut self.velocity_one_body_constraints,
+                total_num_constraints,
+            );
+        }
+
+        let mut curr_start = 0;
+
+        for manifold_i in &self.one_body_interaction_groups.nongrouped_interactions {
+            let manifold = &manifolds_all[*manifold_i];
+            let num_to_add = ConstraintsCounts::from_contacts(manifold).num_constraints;
+
+            OneBodyConstraintBuilder::generate(
+                *manifold_i,
+                manifold,
+                bodies,
+                &mut self.velocity_one_body_constraints_builder[curr_start..],
+                &mut self.velocity_one_body_constraints[curr_start..],
+            );
+
+            curr_start += num_to_add;
+        }
+        assert_eq!(curr_start, total_num_constraints);
+    }
+
+    pub fn solve_restitution(
+        &mut self,
+        solver_vels: &mut [SolverVel<Real>],
+        generic_solver_vels: &mut DVector<Real>,
+    ) {
+        let (jac, constraints) = self.iter_constraints_mut();
+        for mut c in constraints {
+            c.solve_restitution(jac, solver_vels, generic_solver_vels);
+        }
+    }
+
+    pub fn solve_restitution_wo_bias(
+        &mut self,
+        solver_vels: &mut [SolverVel<Real>],
+        generic_solver_vels: &mut DVector<Real>,
+    ) {
+        let (jac, constraints) = self.iter_constraints_mut();
+        for mut c in constraints {
+            c.remove_bias();
+            c.solve_restitution(jac, solver_vels, generic_solver_vels);
+        }
+    }
+
+    pub fn solve_friction(
+        &mut self,
+        solver_vels: &mut [SolverVel<Real>],
+        generic_solver_vels: &mut DVector<Real>,
+    ) {
+        let (jac, constraints) = self.iter_constraints_mut();
+        for mut c in constraints {
+            c.solve_friction(jac, solver_vels, generic_solver_vels);
+        }
+    }
+
+    pub fn writeback_impulses(&mut self, manifolds_all: &mut [&mut ContactManifold]) {
+        let (_, constraints) = self.iter_constraints_mut();
+        for mut c in constraints {
+            c.writeback_impulses(manifolds_all);
+        }
+    }
+
+    pub fn update(
+        &mut self,
+        params: &IntegrationParameters,
+        small_step_id: usize,
+        multibodies: &MultibodyJointSet,
+        solver_bodies: &[SolverBody],
+    ) {
+        macro_rules! update_contacts(
+            ($builders: ident, $constraints: ident) => {
+                for (builder, constraint) in self.$builders.iter().zip(self.$constraints.iter_mut()) {
+                    builder.update(
+                        &params,
+                        small_step_id as Real * params.dt,
+                        solver_bodies,
+                        multibodies,
+                        constraint,
+                    );
+                }
+            }
+        );
+
+        update_contacts!(
+            generic_velocity_constraints_builder,
+            generic_velocity_constraints
+        );
+        update_contacts!(velocity_constraints_builder, velocity_constraints);
+        #[cfg(feature = "simd-is-enabled")]
+        update_contacts!(simd_velocity_constraints_builder, simd_velocity_constraints);
+
+        update_contacts!(
+            generic_velocity_one_body_constraints_builder,
+            generic_velocity_one_body_constraints
+        );
+        update_contacts!(
+            velocity_one_body_constraints_builder,
+            velocity_one_body_constraints
+        );
+        #[cfg(feature = "simd-is-enabled")]
+        update_contacts!(
+            simd_velocity_one_body_constraints_builder,
+            simd_velocity_one_body_constraints
+        );
+    }
+}
diff --git a/src/dynamics/solver/contact_constraint/generic_one_body_constraint.rs b/src/dynamics/solver/contact_constraint/generic_one_body_constraint.rs
new file mode 100644
index 0000000..7c39eab
--- /dev/null
+++ b/src/dynamics/solver/contact_constraint/generic_one_body_constraint.rs
@@ -0,0 +1,292 @@
+use crate::dynamics::solver::OneBodyConstraint;
+use crate::dynamics::{
+    IntegrationParameters, MultibodyJointSet, MultibodyLinkId, RigidBodySet, RigidBodyVelocity,
+};
+use crate::geometry::{ContactManifold, ContactManifoldIndex};
+use crate::math::{Point, Real, DIM, MAX_MANIFOLD_POINTS};
+use crate::utils::SimdCross;
+
+use super::{OneBodyConstraintElement, OneBodyConstraintNormalPart};
+use crate::dynamics::solver::solver_body::SolverBody;
+use crate::dynamics::solver::{ContactPointInfos, OneBodyConstraintBuilder};
+#[cfg(feature = "dim2")]
+use crate::utils::SimdBasis;
+use na::DVector;
+
+#[derive(Copy, Clone)]
+pub(crate) struct GenericOneBodyConstraintBuilder {
+    link2: MultibodyLinkId,
+    ccd_thickness: Real,
+    inner: OneBodyConstraintBuilder,
+}
+
+impl GenericOneBodyConstraintBuilder {
+    pub fn invalid() -> Self {
+        Self {
+            link2: MultibodyLinkId::default(),
+            ccd_thickness: 0.0,
+            inner: OneBodyConstraintBuilder::invalid(),
+        }
+    }
+
+    pub fn generate(
+        manifold_id: ContactManifoldIndex,
+        manifold: &ContactManifold,
+        bodies: &RigidBodySet,
+        multibodies: &MultibodyJointSet,
+        out_builders: &mut [GenericOneBodyConstraintBuilder],
+        out_constraints: &mut [GenericOneBodyConstraint],
+        jacobians: &mut DVector<Real>,
+        jacobian_id: &mut usize,
+    ) {
+        let mut handle1 = manifold.data.rigid_body1;
+        let mut handle2 = manifold.data.rigid_body2;
+        let flipped = manifold.data.relative_dominance <