Implement multibody joints and the new solver

1 files changed, 29 insertions, 82 deletions

diff --git a/src/dynamics/solver/parallel_velocity_solver.rs b/src/dynamics/solver/parallel_velocity_solver.rs
index 54792f8..69ceb03 100644
--- a/src/dynamics/solver/parallel_velocity_solver.rs
+++ b/src/dynamics/solver/parallel_velocity_solver.rs
@@ -1,29 +1,36 @@
 use super::{AnyJointVelocityConstraint, AnyVelocityConstraint, DeltaVel, ThreadContext};
-use crate::dynamics::solver::{
-    AnyJointPositionConstraint, AnyPositionConstraint, ParallelSolverConstraints,
-};
+use crate::dynamics::solver::generic_velocity_constraint::GenericVelocityConstraint;
+use crate::dynamics::solver::ParallelSolverConstraints;
 use crate::dynamics::{IntegrationParameters, JointGraphEdge};
 use crate::geometry::ContactManifold;
 use crate::math::Real;
+use na::DVector;
 use std::sync::atomic::Ordering;
 
-pub(crate) struct ParallelVelocitySolver {}
+pub(crate) struct ParallelVelocitySolver {
+    pub mj_lambdas: Vec<DeltaVel<Real>>,
+    pub generic_mj_lambdas: DVector<Real>,
+}
 
 impl ParallelVelocitySolver {
+    pub fn new() -> Self {
+        Self {
+            mj_lambdas: Vec::new(),
+            generic_mj_lambdas: DVector::zeros(0),
+        }
+    }
+
     pub fn solve(
+        &mut self,
         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,
+            GenericVelocityConstraint,
         >,
+        joint_constraints: &mut ParallelSolverConstraints<AnyJointVelocityConstraint, ()>,
     ) {
         if contact_constraints.constraint_descs.is_empty()
             && joint_constraints.constraint_descs.is_empty()
@@ -32,69 +39,6 @@ impl ParallelVelocitySolver {
         }
 
         /*
-         * Warmstart constraints.
-         */
-        {
-            // Each thread will concurrently grab thread.batch_size constraint desc to
-            // solve. If the batch size is large enough for to cross the boundary of
-            // a parallel_desc_group, we have to wait util the current group is finished
-            // before starting the next one.
-            let mut target_num_desc = 0;
-            let mut start_index = thread
-                .warmstart_constraint_index
-                .fetch_add(thread.batch_size, Ordering::SeqCst);
-            let mut batch_size = thread.batch_size;
-            let mut shift = 0;
-
-            macro_rules! warmstart(
-                ($part: expr) => {
-                    for group in $part.parallel_desc_groups.windows(2) {
-                        let num_descs_in_group = group[1] - group[0];
-                        target_num_desc += num_descs_in_group;
-
-                        while start_index < group[1] {
-                            let end_index = (start_index + batch_size).min(group[1]);
-
-                            let constraints = if end_index == $part.constraint_descs.len() {
-                                &mut $part.velocity_constraints[$part.constraint_descs[start_index].0..]
-                            } else {
-                                &mut $part.velocity_constraints[$part.constraint_descs[start_index].0..$part.constraint_descs[end_index].0]
-                            };
-
-                            for constraint in constraints {
-                                constraint.warmstart(mj_lambdas);
-                            }
-
-                            let num_solved = end_index - start_index;
-                            batch_size -= num_solved;
-
-                            thread
-                                .num_warmstarted_constraints
-                                .fetch_add(num_solved, Ordering::SeqCst);
-
-                            if batch_size == 0 {
-                                start_index = thread
-                                    .warmstart_constraint_index
-                                    .fetch_add(thread.batch_size, Ordering::SeqCst);
-                                start_index -= shift;
-                                batch_size = thread.batch_size;
-                            } else {
-                                start_index += num_solved;
-                            }
-                        }
-
-                        ThreadContext::lock_until_ge(&thread.num_warmstarted_constraints, target_num_desc);
-                    }
-                }
-            );
-
-            warmstart!(joint_constraints);
-            shift = joint_constraints.constraint_descs.len();
-            start_index -= shift;
-            warmstart!(contact_constraints);
-        }
-
-        /*
          * Solve constraints.
          */
         {
@@ -113,8 +57,8 @@ impl ParallelVelocitySolver {
 
             for _ in 0..params.max_velocity_iterations {
                 macro_rules! solve {
-                    ($part: expr) => {
-                        // Joint groups.
+                    ($part: expr, $($solve_args: expr),*) => {
+                        // ImpulseJoint groups.
                         for group in $part.parallel_desc_groups.windows(2) {
                             let num_descs_in_group = group[1] - group[0];
 
@@ -133,12 +77,10 @@ impl ParallelVelocitySolver {
                                         ..$part.constraint_descs[end_index].0]
                                 };
 
-                                //                                println!(
-                                //                                    "Solving a constraint {:?}.",
-                                //                                    rayon::current_thread_index()
-                                //                                );
                                 for constraint in constraints {
-                                    constraint.solve(mj_lambdas);
+                                    constraint.solve(
+                                        $($solve_args),*
+                                    );
                                 }
 
                                 let num_solved = end_index - start_index;
@@ -166,10 +108,15 @@ impl ParallelVelocitySolver {
                     };
                 }
 
-                solve!(joint_constraints);
+                solve!(
+                    joint_constraints,
+                    &joint_constraints.generic_jacobians,
+                    &mut self.mj_lambdas,
+                    &mut self.generic_mj_lambdas
+                );
                 shift += joint_descs.len();
                 start_index -= joint_descs.len();
-                solve!(contact_constraints);
+                solve!(contact_constraints, &mut self.mj_lambdas, true, true);
                 shift += contact_descs.len();
                 start_index -= contact_descs.len();
             }