feat: implement new "small-steps" solver + joint improvements

1 files changed, 446 insertions, 0 deletions

diff --git a/src/dynamics/solver/joint_constraint/joint_constraints_set.rs b/src/dynamics/solver/joint_constraint/joint_constraints_set.rs
new file mode 100644
index 0000000..21b6459
--- /dev/null
+++ b/src/dynamics/solver/joint_constraint/joint_constraints_set.rs
@@ -0,0 +1,446 @@
+use crate::dynamics::solver::categorization::categorize_joints;
+use crate::dynamics::solver::solver_body::SolverBody;
+use crate::dynamics::solver::solver_vel::SolverVel;
+use crate::dynamics::solver::{
+    reset_buffer, JointConstraintTypes, JointGenericOneBodyConstraint,
+    JointGenericOneBodyConstraintBuilder, JointGenericTwoBodyConstraint,
+    JointGenericTwoBodyConstraintBuilder, JointGenericVelocityOneBodyExternalConstraintBuilder,
+    JointGenericVelocityOneBodyInternalConstraintBuilder, SolverConstraintsSet,
+};
+use crate::dynamics::{
+    IntegrationParameters, IslandManager, JointGraphEdge, JointIndex, MultibodyJointSet,
+    RigidBodySet,
+};
+use na::DVector;
+use parry::math::Real;
+
+use crate::dynamics::solver::joint_constraint::joint_constraint_builder::{
+    JointOneBodyConstraintBuilder, JointTwoBodyConstraintBuilder,
+};
+#[cfg(feature = "simd-is-enabled")]
+use {
+    crate::dynamics::solver::joint_constraint::joint_constraint_builder::{
+        JointOneBodyConstraintBuilderSimd, JointTwoBodyConstraintBuilderSimd,
+    },
+    crate::math::SIMD_WIDTH,
+};
+
+pub type JointConstraintsSet = SolverConstraintsSet<JointConstraintTypes>;
+
+impl JointConstraintsSet {
+    pub fn init(
+        &mut self,
+        island_id: usize,
+        islands: &IslandManager,
+        bodies: &RigidBodySet,
+        multibody_joints: &MultibodyJointSet,
+        impulse_joints: &[JointGraphEdge],
+        joint_constraint_indices: &[JointIndex],
+    ) {
+        // Generate constraints for impulse_joints.
+        self.two_body_interactions.clear();
+        self.one_body_interactions.clear();
+        self.generic_two_body_interactions.clear();
+        self.generic_one_body_interactions.clear();
+
+        categorize_joints(
+            bodies,
+            multibody_joints,
+            impulse_joints,
+            joint_constraint_indices,
+            &mut self.one_body_interactions,
+            &mut self.two_body_interactions,
+            &mut self.generic_one_body_interactions,
+            &mut self.generic_two_body_interactions,
+        );
+
+        self.clear_constraints();
+        self.clear_builders();
+
+        self.interaction_groups.clear_groups();
+        self.interaction_groups.group_joints(
+            island_id,
+            islands,
+            bodies,
+            impulse_joints,
+            &self.two_body_interactions,
+        );
+
+        self.one_body_interaction_groups.clear_groups();
+        self.one_body_interaction_groups.group_joints(
+            island_id,
+            islands,
+            bodies,
+            impulse_joints,
+            &self.one_body_interactions,
+        );
+        // NOTE: uncomment this do disable SIMD joint 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);
+
+        let mut j_id = 0;
+        self.compute_joint_constraints(bodies, impulse_joints);
+        #[cfg(feature = "simd-is-enabled")]
+        {
+            self.simd_compute_joint_constraints(bodies, impulse_joints);
+        }
+        self.compute_generic_joint_constraints(bodies, multibody_joints, impulse_joints, &mut j_id);
+
+        self.compute_joint_one_body_constraints(bodies, impulse_joints);
+        #[cfg(feature = "simd-is-enabled")]
+        {
+            self.simd_compute_joint_one_body_constraints(bodies, impulse_joints);
+        }
+        self.compute_generic_one_body_joint_constraints(
+            island_id,
+            islands,
+            bodies,
+            multibody_joints,
+            impulse_joints,
+            &mut j_id,
+        );
+    }
+
+    fn compute_joint_one_body_constraints(
+        &mut self,
+        bodies: &RigidBodySet,
+        joints_all: &[JointGraphEdge],
+    ) {
+        let total_num_builders = self
+            .one_body_interaction_groups
+            .nongrouped_interactions
+            .len();
+
+        unsafe {
+            reset_buffer(
+                &mut self.velocity_one_body_constraints_builder,
+                total_num_builders,
+            );
+        }
+
+        let mut num_constraints = 0;
+        for (joint_i, builder) in self
+            .one_body_interaction_groups
+            .nongrouped_interactions
+            .iter()
+            .zip(self.velocity_one_body_constraints_builder.iter_mut())
+        {
+            let joint = &joints_all[*joint_i].weight;
+            JointOneBodyConstraintBuilder::generate(
+                joint,
+                bodies,
+                *joint_i,
+                builder,
+                &mut num_constraints,
+            );
+        }
+
+        unsafe {
+            reset_buffer(&mut self.velocity_one_body_constraints, num_constraints);
+        }
+    }
+
+    #[cfg(feature = "simd-is-enabled")]
+    fn simd_compute_joint_one_body_constraints(
+        &mut self,
+        bodies: &RigidBodySet,
+        joints_all: &[JointGraphEdge],
+    ) {
+        let total_num_builders =
+            self.one_body_interaction_groups.simd_interactions.len() / SIMD_WIDTH;
+
+        unsafe {
+            reset_buffer(
+                &mut self.simd_velocity_one_body_constraints_builder,
+                total_num_builders,
+            );
+        }
+
+        let mut num_constraints = 0;
+        for (joints_i, builder) in self
+            .one_body_interaction_groups
+            .simd_interactions
+            .chunks_exact(SIMD_WIDTH)
+            .zip(self.simd_velocity_one_body_constraints_builder.iter_mut())
+        {
+            let joints_id = gather![|ii| joints_i[ii]];
+            let impulse_joints = gather![|ii| &joints_all[joints_i[ii]].weight];
+            JointOneBodyConstraintBuilderSimd::generate(
+                impulse_joints,
+                bodies,
+                joints_id,
+                builder,
+                &mut num_constraints,
+            );
+        }
+
+        unsafe {
+            reset_buffer(
+                &mut self.simd_velocity_one_body_constraints,
+                num_constraints,
+            );
+        }
+    }
+
+    fn compute_joint_constraints(&mut self, bodies: &RigidBodySet, joints_all: &[JointGraphEdge]) {
+        let total_num_builders = self.interaction_groups.nongrouped_interactions.len();
+
+        unsafe {
+            reset_buffer(&mut self.velocity_constraints_builder, total_num_builders);
+        }
+
+        let mut num_constraints = 0;
+        for (joint_i, builder) in self
+            .interaction_groups
+            .nongrouped_interactions
+            .iter()
+            .zip(self.velocity_constraints_builder.iter_mut())
+        {
+            let joint = &joints_all[*joint_i].weight;
+            JointTwoBodyConstraintBuilder::generate(
+                joint,
+                bodies,
+                *joint_i,
+                builder,
+                &mut num_constraints,
+            );
+        }
+
+        unsafe {
+            reset_buffer(&mut self.velocity_constraints, num_constraints);
+        }
+    }
+
+    fn compute_generic_joint_constraints(
+        &mut self,
+        bodies: &RigidBodySet,
+        multibodies: &MultibodyJointSet,
+        joints_all: &[JointGraphEdge],
+        j_id: &mut usize,
+    ) {
+        let total_num_builders = self.generic_two_body_interactions.len();
+        self.generic_velocity_constraints_builder.resize(
+            total_num_builders,
+            JointGenericTwoBodyConstraintBuilder::invalid(),
+        );
+
+        let mut num_constraints = 0;
+        for (joint_i, builder) in self
+            .generic_two_body_interactions
+            .iter()
+            .zip(self.generic_velocity_constraints_builder.iter_mut())
+        {
+            let joint = &joints_all[*joint_i].weight;
+            JointGenericTwoBodyConstraintBuilder::generate(
+                *joint_i,
+                joint,
+                bodies,
+                multibodies,
+                builder,
+                j_id,
+                &mut self.generic_jacobians,
+                &mut num_constraints,
+            );
+        }
+
+        self.generic_velocity_constraints
+            .resize(num_constraints, JointGenericTwoBodyConstraint::invalid());
+    }
+
+    fn compute_generic_one_body_joint_constraints(
+        &mut self,
+        island_id: usize,
+        islands: &IslandManager,
+        bodies: &RigidBodySet,
+        multibodies: &MultibodyJointSet,
+        joints_all: &[JointGraphEdge],
+        j_id: &mut usize,
+    ) {
+        let mut total_num_builders = self.generic_one_body_interactions.len();
+
+        for handle in islands.active_island(island_id) {
+            if let Some(link_id) = multibodies.rigid_body_link(*handle) {
+                if JointGenericVelocityOneBodyInternalConstraintBuilder::num_constraints(
+                    multibodies,
+                    link_id,
+                ) > 0
+                {
+                    total_num_builders += 1;
+                }
+            }
+        }
+
+        self.generic_velocity_one_body_constraints_builder.resize(
+            total_num_builders,
+            JointGenericOneBodyConstraintBuilder::invalid(),
+        );
+
+        let mut num_constraints = 0;
+        for (joint_i, builder) in self.generic_one_body_interactions.iter().zip(
+            self.generic_velocity_one_body_constraints_builder
+                .iter_mut(),
+        ) {
+            let joint = &joints_all[*joint_i].weight;
+            JointGenericVelocityOneBodyExternalConstraintBuilder::generate(
+                *joint_i,
+                joint,
+                bodies,
+                multibodies,
+                builder,
+                j_id,
+                &mut self.generic_jacobians,
+                &mut num_constraints,
+            );
+        }
+
+        let mut curr_builder = self.generic_one_body_interactions.len();
+        for handle in islands.active_island(island_id) {
+            if curr_builder >= self.generic_velocity_one_body_constraints_builder.len() {
+                break; // No more builder need to be generated.
+            }
+
+            if let Some(link_id) = multibodies.rigid_body_link(*handle) {
+                let prev_num_constraints = num_constraints;
+                JointGenericVelocityOneBodyInternalConstraintBuilder::generate(
+                    multibodies,
+                    link_id,
+                    &mut self.generic_velocity_one_body_constraints_builder[curr_builder],
+                    j_id,
+                    &mut self.generic_jacobians,
+                    &mut num_constraints,
+                );
+                if num_constraints != prev_num_constraints {
+                    curr_builder += 1;
+                }
+            }
+        }
+
+        self.generic_velocity_one_body_constraints
+            .resize(num_constraints, JointGenericOneBodyConstraint::invalid());
+    }
+
+    #[cfg(feature = "simd-is-enabled")]
+    fn simd_compute_joint_constraints(
+        &mut self,
+        bodies: &RigidBodySet,
+        joints_all: &[JointGraphEdge],
+    ) {
+        let total_num_builders = self.interaction_groups.simd_interactions.len() / SIMD_WIDTH;
+
+        unsafe {
+            reset_buffer(
+                &mut self.simd_velocity_constraints_builder,
+                total_num_builders,
+            );
+        }
+
+        let mut num_constraints = 0;
+        for (joints_i, builder) in self
+            .interaction_groups
+            .simd_interactions
+            .chunks_exact(SIMD_WIDTH)
+            .zip(self.simd_velocity_constraints_builder.iter_mut())
+        {
+            let joints_id = gather![|ii| joints_i[ii]];
+            let impulse_joints = gather![|ii| &joints_all[joints_i[ii]].weight];
+            JointTwoBodyConstraintBuilderSimd::generate(
+                impulse_joints,
+                bodies,
+                joints_id,
+                builder,
+                &mut num_constraints,
+            );
+        }
+
+        unsafe {
+            reset_buffer(&mut self.simd_velocity_constraints, num_constraints);
+        }
+    }
+
+    pub fn solve(
+        &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(jac, solver_vels, generic_solver_vels);
+        }
+    }
+
+    pub fn solve_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(jac, solver_vels, generic_solver_vels);
+        }
+    }
+
+    pub fn writeback_impulses(&mut self, joints_all: &mut [JointGraphEdge]) {
+        let (_, constraints) = self.iter_constraints_mut();
+        for mut c in constraints {
+            c.writeback_impulses(joints_all);
+        }
+    }
+
+    pub fn update(
+        &mut self,
+        params: &IntegrationParameters,
+        multibodies: &MultibodyJointSet,
+        solver_bodies: &[SolverBody],
+    ) {
+        for builder in &mut self.generic_velocity_constraints_builder {
+            builder.update(
+                &params,
+                multibodies,
+                solver_bodies,
+                &mut self.generic_jacobians,
+                &mut self.generic_velocity_constraints,
+            );
+        }
+
+        for builder in &mut self.generic_velocity_one_body_constraints_builder {
+            builder.update(
+                &params,
+                multibodies,
+                solver_bodies,
+                &mut self.generic_jacobians,
+                &mut self.generic_velocity_one_body_constraints,
+            );
+        }
+
+        for builder in &mut self.velocity_constraints_builder {
+            builder.update(&params, solver_bodies, &mut self.velocity_constraints);
+        }
+
+        #[cfg(feature = "simd-is-enabled")]
+        for builder in &mut self.simd_velocity_constraints_builder {
+            builder.update(&params, solver_bodies, &mut self.simd_velocity_constraints);
+        }
+
+        for builder in &mut self.velocity_one_body_constraints_builder {
+            builder.update(
+                &params,
+                solver_bodies,
+                &mut self.velocity_one_body_constraints,
+            );
+        }
+
+        #[cfg(feature = "simd-is-enabled")]
+        for builder in &mut self.simd_velocity_one_body_constraints_builder {
+            builder.update(
+                &params,
+                solver_bodies,
+                &mut self.simd_velocity_one_body_constraints,
+            );
+        }
+    }
+}