Refactor the parallel solver code the same way we did with the non-parallel solver.

1 files changed, 36 insertions, 316 deletions

diff --git a/src/dynamics/solver/parallel_velocity_solver.rs b/src/dynamics/solver/parallel_velocity_solver.rs
index ca2a987..9701547 100644
--- a/src/dynamics/solver/parallel_velocity_solver.rs
+++ b/src/dynamics/solver/parallel_velocity_solver.rs
@@ -1,325 +1,41 @@
 use super::ParallelInteractionGroups;
 use super::{AnyJointVelocityConstraint, AnyVelocityConstraint, DeltaVel, ThreadContext};
 use crate::dynamics::solver::categorization::{categorize_contacts, categorize_joints};
-use crate::dynamics::solver::{InteractionGroups, VelocityConstraint, VelocityGroundConstraint};
+use crate::dynamics::solver::parallel_solver_constraints::ConstraintDesc;
+use crate::dynamics::solver::{
+    AnyJointPositionConstraint, AnyPositionConstraint, InteractionGroups,
+    ParallelSolverConstraints, VelocityConstraint, VelocityGroundConstraint,
+};
 use crate::dynamics::{IntegrationParameters, JointGraphEdge, RigidBodySet};
 use crate::geometry::ContactManifold;
 #[cfg(feature = "simd-is-enabled")]
 use crate::{
     dynamics::solver::{WVelocityConstraint, WVelocityGroundConstraint},
-    simd::SIMD_WIDTH,
+    math::{Real, SIMD_WIDTH},
 };
 use std::sync::atomic::Ordering;
 
-pub(crate) enum VelocityConstraintDesc {
-    NongroundNongrouped(usize),
-    GroundNongrouped(usize),
-    #[cfg(feature = "simd-is-enabled")]
-    NongroundGrouped([usize; SIMD_WIDTH]),
-    #[cfg(feature = "simd-is-enabled")]
-    GroundGrouped([usize; SIMD_WIDTH]),
-}
-
-pub(crate) struct ParallelSolverConstraints<Constraint> {
-    pub not_ground_interactions: Vec<usize>,
-    pub ground_interactions: Vec<usize>,
-    pub interaction_groups: InteractionGroups,
-    pub ground_interaction_groups: InteractionGroups,
-    pub constraints: Vec<Constraint>,
-    pub constraint_descs: Vec<(usize, VelocityConstraintDesc)>,
-    pub parallel_desc_groups: Vec<usize>,
-}
-
-impl<Constraint> ParallelSolverConstraints<Constraint> {
-    pub fn new() -> Self {
-        Self {
-            not_ground_interactions: Vec::new(),
-            ground_interactions: Vec::new(),
-            interaction_groups: InteractionGroups::new(),
-            ground_interaction_groups: InteractionGroups::new(),
-            constraints: Vec::new(),
-            constraint_descs: Vec::new(),
-            parallel_desc_groups: Vec::new(),
-        }
-    }
-}
-
-macro_rules! impl_init_constraints_group {
-    ($Constraint: ty, $Interaction: ty, $categorize: ident, $group: ident, $num_active_constraints: path, $empty_constraint: expr $(, $weight: ident)*) => {
-        impl ParallelSolverConstraints<$Constraint> {
-            pub fn init_constraints_groups(
-                &mut self,
-                island_id: usize,
-                bodies: &RigidBodySet,
-                interactions: &mut [$Interaction],
-                interaction_groups: &ParallelInteractionGroups,
-            ) {
-                let mut total_num_constraints = 0;
-                let num_groups = interaction_groups.num_groups();
-
-                self.interaction_groups.clear_groups();
-                self.ground_interaction_groups.clear_groups();
-                self.parallel_desc_groups.clear();
-                self.constraint_descs.clear();
-                self.parallel_desc_groups.push(0);
-
-                for i in 0..num_groups {
-                    let group = interaction_groups.group(i);
-
-                    self.not_ground_interactions.clear();
-                    self.ground_interactions.clear();
-                    $categorize(
-                        bodies,
-                        interactions,
-                        group,
-                        &mut self.ground_interactions,
-                        &mut self.not_ground_interactions,
-                    );
-
-                    #[cfg(feature = "simd-is-enabled")]
-                    let start_grouped = self.interaction_groups.grouped_interactions.len();
-                    let start_nongrouped = self.interaction_groups.nongrouped_interactions.len();
-                    #[cfg(feature = "simd-is-enabled")]
-                    let start_grouped_ground = self.ground_interaction_groups.grouped_interactions.len();
-                    let start_nongrouped_ground = self.ground_interaction_groups.nongrouped_interactions.len();
-
-                    self.interaction_groups.$group(
-                        island_id,
-                        bodies,
-                        interactions,
-                        &self.not_ground_interactions,
-                    );
-                    self.ground_interaction_groups.$group(
-                        island_id,
-                        bodies,
-                        interactions,
-                        &self.ground_interactions,
-                    );
-
-                    // Compute constraint indices.
-                    for interaction_i in &self.interaction_groups.nongrouped_interactions[start_nongrouped..] {
-                        let interaction = &mut interactions[*interaction_i]$(.$weight)*;
-                        interaction.constraint_index = total_num_constraints;
-                        self.constraint_descs.push((
-                            total_num_constraints,
-                            VelocityConstraintDesc::NongroundNongrouped(*interaction_i),
-                        ));
-                        total_num_constraints += $num_active_constraints(interaction);
-                    }
-
-                    #[cfg(feature = "simd-is-enabled")]
-                    for interaction_i in
-                        self.interaction_groups.grouped_interactions[start_grouped..].chunks(SIMD_WIDTH)
-                    {
-                        let interaction = &mut interactions[interaction_i[0]]$(.$weight)*;
-                        interaction.constraint_index = total_num_constraints;
-                        self.constraint_descs.push((
-                            total_num_constraints,
-                            VelocityConstraintDesc::NongroundGrouped(
-                                array![|ii| interaction_i[ii]; SIMD_WIDTH],
-                            ),
-                        ));
-                        total_num_constraints += $num_active_constraints(interaction);
-                    }
-
-                    for interaction_i in
-                        &self.ground_interaction_groups.nongrouped_interactions[start_nongrouped_ground..]
-                    {
-                        let interaction = &mut interactions[*interaction_i]$(.$weight)*;
-                        interaction.constraint_index = total_num_constraints;
-                        self.constraint_descs.push((
-                            total_num_constraints,
-                            VelocityConstraintDesc::GroundNongrouped(*interaction_i),
-                        ));
-                        total_num_constraints += $num_active_constraints(interaction);
-                    }
-
-                    #[cfg(feature = "simd-is-enabled")]
-                    for interaction_i in self.ground_interaction_groups.grouped_interactions
-                        [start_grouped_ground..]
-                        .chunks(SIMD_WIDTH)
-                    {
-                        let interaction = &mut interactions[interaction_i[0]]$(.$weight)*;
-                        interaction.constraint_index = total_num_constraints;
-                        self.constraint_descs.push((
-                            total_num_constraints,
-                            VelocityConstraintDesc::GroundGrouped(
-                                array![|ii| interaction_i[ii]; SIMD_WIDTH],
-                            ),
-                        ));
-                        total_num_constraints += $num_active_constraints(interaction);
-                    }
-
-                    self.parallel_desc_groups.push(self.constraint_descs.len());
-                }
-
-                // Resize the constraints set.
-                self.constraints.clear();
-                self.constraints
-                    .resize_with(total_num_constraints, || $empty_constraint)
-            }
-        }
-    }
-}
-
-impl_init_constraints_group!(
-    AnyVelocityConstraint,
-    &mut ContactManifold,
-    categorize_contacts,
-    group_manifolds,
-    VelocityConstraint::num_active_constraints,
-    AnyVelocityConstraint::Empty
-);
-
-impl_init_constraints_group!(
-    AnyJointVelocityConstraint,
-    JointGraphEdge,
-    categorize_joints,
-    group_joints,
-    AnyJointVelocityConstraint::num_active_constraints,
-    AnyJointVelocityConstraint::Empty,
-    weight
-);
-
-impl ParallelSolverConstraints<AnyVelocityConstraint> {
-    fn fill_constraints(
-        &mut self,
-        thread: &ThreadContext,
-        params: &IntegrationParameters,
-        bodies: &RigidBodySet,
-        manifolds_all: &[&mut ContactManifold],
-    ) {
-        let descs = &self.constraint_descs;
-
-        crate::concurrent_loop! {
-            let batch_size = thread.batch_size;
-            for desc in descs[thread.constraint_initialization_index, thread.num_initialized_constraints] {
-                match &desc.1 {
-                    VelocityConstraintDesc::NongroundNongrouped(manifold_id) => {
-                        let manifold = &*manifolds_all[*manifold_id];
-                        VelocityConstraint::generate(params, *manifold_id, manifold, bodies, &mut self.constraints, false);
-                    }
-                    VelocityConstraintDesc::GroundNongrouped(manifold_id) => {
-                        let manifold = &*manifolds_all[*manifold_id];
-                        VelocityGroundConstraint::generate(params, *manifold_id, manifold, bodies, &mut self.constraints, false);
-                    }
-                    #[cfg(feature = "simd-is-enabled")]
-                    VelocityConstraintDesc::NongroundGrouped(manifold_id) => {
-                        let manifolds = array![|ii| &*manifolds_all[manifold_id[ii]]; SIMD_WIDTH];
-                        WVelocityConstraint::generate(params, *manifold_id, manifolds, bodies, &mut self.constraints, false);
-                    }
-                    #[cfg(feature = "simd-is-enabled")]
-                    VelocityConstraintDesc::GroundGrouped(manifold_id) => {
-                        let manifolds = array![|ii| &*manifolds_all[manifold_id[ii]]; SIMD_WIDTH];
-                        WVelocityGroundConstraint::generate(params, *manifold_id, manifolds, bodies, &mut self.constraints, false);
-                    }
-                }
-            }
-        }
-    }
-}
-
-impl ParallelSolverConstraints<AnyJointVelocityConstraint> {
-    fn fill_constraints(
-        &mut self,
-        thread: &ThreadContext,
-        params: &IntegrationParameters,
-        bodies: &RigidBodySet,
-        joints_all: &[JointGraphEdge],
-    ) {
-        let descs = &self.constraint_descs;
-
-        crate::concurrent_loop! {
-            let batch_size = thread.batch_size;
-            for desc in descs[thread.joint_constraint_initialization_index, thread.num_initialized_joint_constraints] {
-                match &desc.1 {
-                    VelocityConstraintDesc::NongroundNongrouped(joint_id) => {
-                        let joint = &joints_all[*joint_id].weight;
-                        let constraint = AnyJointVelocityConstraint::from_joint(params, *joint_id, joint, bodies);
-                        self.constraints[joint.constraint_index] = constraint;
-                    }
-                    VelocityConstraintDesc::GroundNongrouped(joint_id) => {
-                        let joint = &joints_all[*joint_id].weight;
-                        let constraint = AnyJointVelocityConstraint::from_joint_ground(params, *joint_id, joint, bodies);
-                        self.constraints[joint.constraint_index] = constraint;
-                    }
-                    #[cfg(feature = "simd-is-enabled")]
-                    VelocityConstraintDesc::NongroundGrouped(joint_id) => {
-                        let joints = array![|ii| &joints_all[joint_id[ii]].weight; SIMD_WIDTH];
-                        let constraint = AnyJointVelocityConstraint::from_wide_joint(params, *joint_id, joints, bodies);
-                        self.constraints[joints[0].data.constraint_index] = constraint;
-                    }
-                    #[cfg(feature = "simd-is-enabled")]
-                    VelocityConstraintDesc::GroundGrouped(joint_id) => {
-                        let joints = array![|ii| &joints_all[joint_id[ii]].weight; SIMD_WIDTH];
-                        let constraint = AnyJointVelocityConstraint::from_wide_joint_ground(params, *joint_id, joints, bodies);
-                        self.constraints[joints[0].data.constraint_index] = constraint;
-                    }
-                }
-            }
-        }
-    }
-}
-
-pub(crate) struct ParallelVelocitySolver {
-    part: ParallelSolverConstraints<AnyVelocityConstraint>,
-    joint_part: ParallelSolverConstraints<AnyJointVelocityConstraint>,
-}
+pub(crate) struct ParallelVelocitySolver {}
 
 impl ParallelVelocitySolver {
-    pub fn new() -> Self {
-        Self {
-            part: ParallelSolverConstraints::new(),
-            joint_part: ParallelSolverConstraints::new(),
-        }
-    }
-
-    pub fn init_constraint_groups(
-        &mut self,
-        island_id: usize,
-        bodies: &RigidBodySet,
-        manifolds: &mut [&mut ContactManifold],
-        manifold_groups: &ParallelInteractionGroups,
-        joints: &mut [JointGraphEdge],
-        joint_groups: &ParallelInteractionGroups,
-    ) {
-        self.part
-            .init_constraints_groups(island_id, bodies, manifolds, manifold_groups);
-        self.joint_part
-            .init_constraints_groups(island_id, bodies, joints, joint_groups);
-    }
-
-    pub fn fill_constraints(
-        &mut self,
-        thread: &ThreadContext,
-        params: &IntegrationParameters,
-        bodies: &RigidBodySet,
-        manifolds: &[&mut ContactManifold],
-        joints: &[JointGraphEdge],
-    ) {
-        self.part
-            .fill_constraints(thread, params, bodies, manifolds);
-        self.joint_part
-            .fill_constraints(thread, params, bodies, joints);
-        ThreadContext::lock_until_ge(
-            &thread.num_initialized_constraints,
-            self.part.constraint_descs.len(),
-        );
-        ThreadContext::lock_until_ge(
-            &thread.num_initialized_joint_constraints,
-            self.joint_part.constraint_descs.len(),
-        );
-    }
-
-    pub fn solve_constraints(
-        &mut self,
+    pub fn solve(
         thread: &ThreadContext,
         params: &IntegrationParameters,
         manifolds_all: &mut [&mut ContactManifold],
         joints_all: &mut [JointGraphEdge],
         mj_lambdas: &mut [DeltaVel<Real>],
+        contact_constraints: &mut ParallelSolverConstraints<
+            AnyVelocityConstraint,
+            AnyPositionConstraint,
+        >,
+        joint_constraints: &mut ParallelSolverConstraints<
+            AnyJointVelocityConstraint,
+            AnyJointPositionConstraint,
+        >,
     ) {
-        if self.part.constraint_descs.len() == 0 && self.joint_part.constraint_descs.len() == 0 {
+        if contact_constraints.constraint_descs.is_empty()
+            && joint_constraints.constraint_descs.is_empty()
+        {
             return;
         }
 
@@ -348,9 +64,9 @@ impl ParallelVelocitySolver {
                             let end_index = (start_index + batch_size).min(group[1]);
 
                             let constraints = if end_index == $part.constraint_descs.len() {
-                                &mut $part.constraints[$part.constraint_descs[start_index].0..]
+                                &mut $part.velocity_constraints[$part.constraint_descs[start_index].0..]
                             } else {
-                                &mut $part.constraints[$part.constraint_descs[start_index].0..$part.constraint_descs[end_index].0]
+                                &mut $part.velocity_constraints[$part.constraint_descs[start_index].0..$part.constraint_descs[end_index].0]
                             };
 
                             for constraint in constraints {
@@ -380,10 +96,10 @@ impl ParallelVelocitySolver {
                 }
             );
 
-            warmstart!(self.joint_part);
-            shift = self.joint_part.constraint_descs.len();
+            warmstart!(joint_constraints);
+            shift = joint_constraints.constraint_descs.len();
             start_index -= shift;
-            warmstart!(self.part);
+            warmstart!(contact_constraints);
         }
 
         /*
@@ -398,8 +114,8 @@ impl ParallelVelocitySolver {
                 .solve_interaction_index
                 .fetch_add(thread.batch_size, Ordering::SeqCst);
             let mut batch_size = thread.batch_size;
-            let contact_descs = &self.part.constraint_descs[..];
-            let joint_descs = &self.joint_part.constraint_descs[..];
+            let contact_descs = &contact_constraints.constraint_descs[..];
+            let joint_descs = &joint_constraints.constraint_descs[..];
             let mut target_num_desc = 0;
             let mut shift = 0;
 
@@ -416,9 +132,12 @@ impl ParallelVelocitySolver {
                                 let end_index = (start_index + batch_size).min(group[1]);
 
                                 let constraints = if end_index == $part.constraint_descs.len() {
-                                    &mut $part.constraints[$part.constraint_descs[start_index].0..]
+                                    &mut $part.velocity_constraints
+                                        [$part.constraint_descs[start_index].0..]
                                 } else {
-                                    &mut $part.constraints[$part.constraint_descs[start_index].0
+                                    &mut $part.velocity_constraints[$part.constraint_descs
+                                        [start_index]
+                                        .0
                                         ..$part.constraint_descs[end_index].0]
                                 };
 
@@ -455,10 +174,10 @@ impl ParallelVelocitySolver {
                     };
                 }
 
-                solve!(self.joint_part);
+                solve!(joint_constraints);
                 shift += joint_descs.len();
                 start_index -= joint_descs.len();
-                solve!(self.part);
+                solve!(contact_constraints);
                 shift += contact_descs.len();
                 start_index -= contact_descs.len();
             }
@@ -467,8 +186,9 @@ impl ParallelVelocitySolver {
         /*
          * Writeback impulses.
          */
-        let joint_constraints = &self.joint_part.constraints;
-        let contact_constraints = &self.part.constraints;
+        let joint_constraints = &joint_constraints.velocity_constraints;
+        let contact_constraints = &contact_constraints.velocity_constraints;
+
         crate::concurrent_loop! {
              let batch_size = thread.batch_size;
              for constraint in joint_constraints[thread.joint_writeback_index] {