Complete the parallel solver fix

5 files changed, 55 insertions, 19 deletions

diff --git a/src/dynamics/joint/impulse_joint/impulse_joint.rs b/src/dynamics/joint/impulse_joint/impulse_joint.rs
index 4677a1f..f80380e 100644
--- a/src/dynamics/joint/impulse_joint/impulse_joint.rs
+++ b/src/dynamics/joint/impulse_joint/impulse_joint.rs
@@ -15,5 +15,4 @@ pub struct ImpulseJoint {
 
     // A joint needs to know its handle to simplify its removal.
     pub(crate) handle: ImpulseJointHandle,
-    pub(crate) constraint_index: usize,
 }
diff --git a/src/dynamics/joint/impulse_joint/impulse_joint_set.rs b/src/dynamics/joint/impulse_joint/impulse_joint_set.rs
index cd11a61..448b87d 100644
--- a/src/dynamics/joint/impulse_joint/impulse_joint_set.rs
+++ b/src/dynamics/joint/impulse_joint/impulse_joint_set.rs
@@ -187,7 +187,6 @@ impl ImpulseJointSet {
             data,
             impulses: na::zero(),
             handle: ImpulseJointHandle(handle),
-            constraint_index: 0,
         };
 
         let default_id = InteractionGraph::<(), ()>::invalid_graph_index();
diff --git a/src/dynamics/joint/multibody_joint/multibody.rs b/src/dynamics/joint/multibody_joint/multibody.rs
index f2fa623..58d11b6 100644
--- a/src/dynamics/joint/multibody_joint/multibody.rs
+++ b/src/dynamics/joint/multibody_joint/multibody.rs
@@ -971,27 +971,48 @@ impl Multibody {
     }
 
     #[inline]
+    pub fn num_active_internal_constraints_and_jacobian_lines(&self) -> (usize, usize) {
+        let num_constraints: usize = self
+            .links
+            .iter()
+            .map(|l| l.joint().num_velocity_constraints())
+            .sum();
+        (num_constraints, num_constraints)
+    }
+
+    #[inline]
     pub fn generate_internal_constraints(
         &self,
         params: &IntegrationParameters,
         j_id: &mut usize,
         jacobians: &mut DVector<Real>,
         out: &mut Vec<AnyJointVelocityConstraint>,
+        mut insert_at: Option<usize>,
     ) {
-        let num_constraints: usize = self
-            .links
-            .iter()
-            .map(|l| l.joint().num_velocity_constraints())
-            .sum();
-
-        let required_jacobian_len = *j_id + num_constraints * self.ndofs * 2;
-        if jacobians.nrows() < required_jacobian_len {
-            jacobians.resize_vertically_mut(required_jacobian_len, 0.0);
+        if !cfg!(feature = "parallel") {
+            let num_constraints: usize = self
+                .links
+                .iter()
+                .map(|l| l.joint().num_velocity_constraints())
+                .sum();
+
+            let required_jacobian_len = *j_id + num_constraints * self.ndofs * 2;
+            if jacobians.nrows() < required_jacobian_len {
+                jacobians.resize_vertically_mut(required_jacobian_len, 0.0);
+            }
         }
 
         for link in self.links.iter() {
-            link.joint()
-                .velocity_constraints(params, self, link, 0, j_id, jacobians, out);
+            link.joint().velocity_constraints(
+                params,
+                self,
+                link,
+                0,
+                j_id,
+                jacobians,
+                out,
+                &mut insert_at,
+            );
         }
     }
 }
diff --git a/src/dynamics/joint/multibody_joint/multibody_joint.rs b/src/dynamics/joint/multibody_joint/multibody_joint.rs
index 2f7a71e..5a04070 100644
--- a/src/dynamics/joint/multibody_joint/multibody_joint.rs
+++ b/src/dynamics/joint/multibody_joint/multibody_joint.rs
@@ -255,6 +255,7 @@ impl MultibodyJoint {
         j_id: &mut usize,
         jacobians: &mut DVector<Real>,
         constraints: &mut Vec<AnyJointVelocityConstraint>,
+        insert_at: &mut Option<usize>,
     ) {
         let locked_bits = self.data.locked_axes.bits();
         let limit_bits = self.data.limit_axes.bits();
@@ -281,6 +282,7 @@ impl MultibodyJoint {
                         j_id,
                         jacobians,
                         constraints,
+                        insert_at,
                     );
                 }
 
@@ -295,6 +297,7 @@ impl MultibodyJoint {
                         j_id,
                         jacobians,
                         constraints,
+                        insert_at,
                     );
                 }
                 curr_free_dof += 1;
@@ -329,6 +332,7 @@ impl MultibodyJoint {
                         j_id,
                         jacobians,
                         constraints,
+                        insert_at,
                     );
                     Some(limits)
                 } else {
@@ -347,6 +351,7 @@ impl MultibodyJoint {
                         j_id,
                         jacobians,
                         constraints,
+                        insert_at,
                     );
                 }
                 curr_free_dof += 1;
diff --git a/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs b/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs
index 42212be..a1ec483 100644
--- a/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs
+++ b/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs
@@ -20,6 +20,7 @@ pub fn unit_joint_limit_constraint(
     j_id: &mut usize,
     jacobians: &mut DVector<Real>,
     constraints: &mut Vec<AnyJointVelocityConstraint>,
+    insert_at: &mut Option<usize>,
 ) {
     let ndofs = multibody.ndofs();
     let joint_velocity = multibody.joint_velocity(link);
@@ -60,9 +61,14 @@ pub fn unit_joint_limit_constraint(
         writeback_id: WritebackId::Limit(dof_id),
     };
 
-    constraints.push(AnyJointVelocityConstraint::JointGenericGroundConstraint(
-        constraint,
-    ));
+    if let Some(at) = insert_at {
+        constraints[*at] = AnyJointVelocityConstraint::JointGenericGroundConstraint(constraint);
+        *at += 1;
+    } else {
+        constraints.push(AnyJointVelocityConstraint::JointGenericGroundConstraint(
+            constraint,
+        ));
+    }
     *j_id += 2 * ndofs;
 }
 
@@ -79,6 +85,7 @@ pub fn unit_joint_motor_constraint(
     j_id: &mut usize,
     jacobians: &mut DVector<Real>,
     constraints: &mut Vec<AnyJointVelocityConstraint>,
+    insert_at: &mut Option<usize>,
 ) {
     let inv_dt = params.inv_dt();
     let ndofs = multibody.ndofs();
@@ -128,8 +135,13 @@ pub fn unit_joint_motor_constraint(
         writeback_id: WritebackId::Limit(dof_id),
     };
 
-    constraints.push(AnyJointVelocityConstraint::JointGenericGroundConstraint(
-        constraint,
-    ));
+    if let Some(at) = insert_at {
+        constraints[*at] = AnyJointVelocityConstraint::JointGenericGroundConstraint(constraint);
+        *at += 1;
+    } else {
+        constraints.push(AnyJointVelocityConstraint::JointGenericGroundConstraint(
+            constraint,
+        ));
+    }
     *j_id += 2 * ndofs;
 }