Some minor cleanup and joint constraint refactoring.

3 files changed, 96 insertions, 83 deletions

diff --git a/src/dynamics/solver/joint_constraint/ball_velocity_constraint.rs b/src/dynamics/solver/joint_constraint/ball_velocity_constraint.rs
index 88e7c83..cb95532 100644
--- a/src/dynamics/solver/joint_constraint/ball_velocity_constraint.rs
+++ b/src/dynamics/solver/joint_constraint/ball_velocity_constraint.rs
@@ -183,10 +183,7 @@ impl BallVelocityConstraint {
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }
 
-    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
-        let mut mj_lambda1 = mj_lambdas[self.mj_lambda1 as usize];
-        let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
-
+    fn solve_dofs(&mut self, mj_lambda1: &mut DeltaVel<Real>, mj_lambda2: &mut DeltaVel<Real>) {
         let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
         let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
         let vel1 = mj_lambda1.linear + ang_vel1.gcross(self.r1);
@@ -201,10 +198,9 @@ impl BallVelocityConstraint {
 
         mj_lambda2.linear -= self.im2 * impulse;
         mj_lambda2.angular -= self.ii2_sqrt.transform_vector(self.r2.gcross(impulse));
+    }
 
-        /*
-         * Motor part.
-         */
+    fn solve_motors(&mut self, mj_lambda1: &mut DeltaVel<Real>, mj_lambda2: &mut DeltaVel<Real>) {
         if let Some(motor_inv_lhs) = &self.motor_inv_lhs {
             let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
             let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
@@ -224,6 +220,14 @@ impl BallVelocityConstraint {
             mj_lambda1.angular += self.ii1_sqrt.transform_vector(effective_impulse);
             mj_lambda2.angular -= self.ii2_sqrt.transform_vector(effective_impulse);
         }
+    }
+
+    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
+        let mut mj_lambda1 = mj_lambdas[self.mj_lambda1 as usize];
+        let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
+
+        self.solve_dofs(&mut mj_lambda1, &mut mj_lambda2);
+        self.solve_motors(&mut mj_lambda1, &mut mj_lambda2);
 
         mj_lambdas[self.mj_lambda1 as usize] = mj_lambda1;
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
@@ -232,7 +236,8 @@ impl BallVelocityConstraint {
     pub fn writeback_impulses(&self, joints_all: &mut [JointGraphEdge]) {
         let joint = &mut joints_all[self.joint_id].weight;
         if let JointParams::BallJoint(ball) = &mut joint.params {
-            ball.impulse = self.impulse
+            ball.impulse = self.impulse;
+            ball.motor_impulse = self.motor_impulse;
         }
     }
 }
@@ -393,9 +398,7 @@ impl BallVelocityGroundConstraint {
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }
 
-    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
-        let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
-
+    fn solve_dofs(&mut self, mj_lambda2: &mut DeltaVel<Real>) {
         let angvel = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
         let vel2 = mj_lambda2.linear + angvel.gcross(self.r2);
         let dvel = vel2 + self.rhs;
@@ -405,10 +408,9 @@ impl BallVelocityGroundConstraint {
 
         mj_lambda2.linear -= self.im2 * impulse;
         mj_lambda2.angular -= self.ii2_sqrt.transform_vector(self.r2.gcross(impulse));
+    }
 
-        /*
-         * Motor part.
-         */
+    fn solve_motors(&mut self, mj_lambda2: &mut DeltaVel<Real>) {
         if let Some(motor_inv_lhs) = &self.motor_inv_lhs {
             let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
 
@@ -425,6 +427,13 @@ impl BallVelocityGroundConstraint {
 
             mj_lambda2.angular -= self.ii2_sqrt.transform_vector(effective_impulse);
         }
+    }
+
+    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
+        let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
+
+        self.solve_dofs(&mut mj_lambda2);
+        self.solve_motors(&mut mj_lambda2);
 
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }
@@ -433,7 +442,8 @@ impl BallVelocityGroundConstraint {
     pub fn writeback_impulses(&self, joints_all: &mut [JointGraphEdge]) {
         let joint = &mut joints_all[self.joint_id].weight;
         if let JointParams::BallJoint(ball) = &mut joint.params {
-            ball.impulse = self.impulse
+            ball.impulse = self.impulse;
+            ball.motor_impulse = self.motor_impulse;
         }
     }
 }
diff --git a/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint_wide.rs b/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint_wide.rs
index bba9cb6..be85024 100644
--- a/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint_wide.rs
+++ b/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint_wide.rs
@@ -308,27 +308,11 @@ impl WPrismaticVelocityConstraint {
         }
     }
 
-    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
-        let mut mj_lambda1 = DeltaVel {
-            linear: Vector::from(
-                array![|ii| mj_lambdas[self.mj_lambda1[ii] as usize].linear; SIMD_WIDTH],
-            ),
-            angular: AngVector::from(
-                array![|ii| mj_lambdas[self.mj_lambda1[ii] as usize].angular; SIMD_WIDTH],
-            ),
-        };
-        let mut mj_lambda2 = DeltaVel {
-            linear: Vector::from(
-                array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear; SIMD_WIDTH],
-            ),
-            angular: AngVector::from(
-                array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular; SIMD_WIDTH],
-            ),
-        };
-
-        /*
-         * Joint consraint.
-         */
+    fn solve_dofs(
+        &mut self,
+        mj_lambda1: &mut DeltaVel<SimdReal>,
+        mj_lambda2: &mut DeltaVel<SimdReal>,
+    ) {
         let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
         let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
         let lin_vel1 = mj_lambda1.linear + ang_vel1.gcross(self.r1);
@@ -357,10 +341,13 @@ impl WPrismaticVelocityConstraint {
         mj_lambda2.angular -= self
             .ii2_sqrt
             .transform_vector(ang_impulse + self.r2.gcross(lin_impulse));
+    }
 
-        /*
-         * Joint limits.
-         */
+    fn solve_limits(
+        &mut self,
+        mj_lambda1: &mut DeltaVel<SimdReal>,
+        mj_lambda2: &mut DeltaVel<SimdReal>,
+    ) {
         if let Some((limits_forcedir1, limits_forcedir2)) = self.limits_forcedirs {
             let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
             let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
@@ -381,6 +368,28 @@ impl WPrismaticVelocityConstraint {
             mj_lambda2.linear += lin_impulse2 * self.im2;
             mj_lambda2.angular += self.ii2_sqrt.transform_vector(self.r2.gcross(lin_impulse2));
         }
+    }
+
+    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
+        let mut mj_lambda1 = DeltaVel {
+            linear: Vector::from(
+                array![|ii| mj_lambdas[self.mj_lambda1[ii] as usize].linear; SIMD_WIDTH],
+            ),
+            angular: AngVector::from(
+                array![|ii| mj_lambdas[self.mj_lambda1[ii] as usize].angular; SIMD_WIDTH],
+            ),
+        };
+        let mut mj_lambda2 = DeltaVel {
+            linear: Vector::from(
+                array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear; SIMD_WIDTH],
+            ),
+            angular: AngVector::from(
+                array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular; SIMD_WIDTH],
+            ),
+        };
+
+        self.solve_dofs(&mut mj_lambda1, &mut mj_lambda2);
+        self.solve_limits(&mut mj_lambda1, &mut mj_lambda2);
 
         for ii in 0..SIMD_WIDTH {
             mj_lambdas[self.mj_lambda1[ii] as usize].linear = mj_lambda1.linear.extract(ii);
@@ -503,19 +512,6 @@ impl WPrismaticVelocityGroundConstraint {
         let axis1 = position1 * local_axis1;
         let axis2 = position2 * local_axis2;
 
-        // #[cfg(feature = "dim2")]
-        // let r21 = Rotation::rotation_between_axis(&axis1, &axis2)
-        //     .to_rotation_matrix()
-        //     .into_inner();
-        // #[cfg(feature = "dim3")]
-        // let r21 = Rotation::rotation_between_axis(&axis1, &axis2)
-        //     .unwrap_or_else(Rotation::identity)
-        //     .to_rotation_matrix()
-        //     .into_inner();
-        // let basis2 = r21 * basis1;
-        // NOTE: we use basis2 := basis1 for now is that allows
-        // simplifications of the computation without introducing
-        // much instabilities.
         let ii2 = ii2_sqrt.squared();
         let r1 = anchor1 - world_com1;
         let r2 = anchor2 - world_com2;
@@ -642,19 +638,7 @@ impl WPrismaticVelocityGroundConstraint {
         }
     }
 
-    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
-        let mut mj_lambda2 = DeltaVel {
-            linear: Vector::from(
-                array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear; SIMD_WIDTH],
-            ),
-            angular: AngVector::from(
-                array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular; SIMD_WIDTH],
-            ),
-        };
-
-        /*
-         * Joint consraint.
-         */
+    fn solve_dofs(&mut self, mj_lambda2: &mut DeltaVel<SimdReal>) {
         let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
         let lin_vel2 = mj_lambda2.linear + ang_vel2.gcross(self.r2);
         let lin_dvel = self.basis1.tr_mul(&lin_vel2);
@@ -676,10 +660,9 @@ impl WPrismaticVelocityGroundConstraint {
         mj_lambda2.angular -= self
             .ii2_sqrt
             .transform_vector(ang_impulse + self.r2.gcross(lin_impulse));
+    }
 
-        /*
-         * Joint limits.
-         */
+    fn solve_limits(&mut self, mj_lambda2: &mut DeltaVel<SimdReal>) {
         if let Some(limits_forcedir2) = self.limits_forcedir2 {
             // FIXME: the transformation by ii2_sqrt could be avoided by
             // reusing some computations above.
@@ -693,6 +676,20 @@ impl WPrismaticVelocityGroundConstraint {
 
             mj_lambda2.linear += limits_forcedir2 * (self.im2 * dimpulse);
         }
+    }
+
+    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
+        let mut mj_lambda2 = DeltaVel {
+            linear: Vector::from(
+                array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear; SIMD_WIDTH],
+            ),
+            angular: AngVector::from(
+                array![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular; SIMD_WIDTH],
+            ),
+        };
+
+        self.solve_dofs(&mut mj_lambda2);
+        self.solve_limits(&mut mj_lambda2);
 
         for ii in 0..SIMD_WIDTH {
             mj_lambdas[self.mj_lambda2[ii] as usize].linear = mj_lambda2.linear.extract(ii);
@@ -700,7 +697,6 @@ impl WPrismaticVelocityGroundConstraint {
         }
     }
 
-    // FIXME: duplicated code with the non-ground constraint.
     pub fn writeback_impulses(&self, joints_all: &mut [JointGraphEdge]) {
         for ii in 0..SIMD_WIDTH {
             let joint = &mut joints_all[self.joint_id[ii]].weight;
diff --git a/src/dynamics/solver/joint_constraint/revolute_velocity_constraint.rs b/src/dynamics/solver/joint_constraint/revolute_velocity_constraint.rs
index ef3a1ba..c1975da 100644
--- a/src/dynamics/solver/joint_constraint/revolute_velocity_constraint.rs
+++ b/src/dynamics/solver/joint_constraint/revolute_velocity_constraint.rs
@@ -209,10 +209,7 @@ impl RevoluteVelocityConstraint {
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }
 
-    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
-        let mut mj_lambda1 = mj_lambdas[self.mj_lambda1 as usize];
-        let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
-
+    fn solve_dofs(&mut self, mj_lambda1: &mut DeltaVel<Real>, mj_lambda2: &mut DeltaVel<Real>) {
         let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
         let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
 
@@ -237,10 +234,9 @@ impl RevoluteVelocityConstraint {
         mj_lambda2.angular -= self
             .ii2_sqrt
             .transform_vector(ang_impulse2 + self.r2.gcross(lin_impulse2));
+    }
 
-        /*
-         * Motor.
-         */
+    fn solve_motors(&mut self, mj_lambda1: &mut DeltaVel<Real>, mj_lambda2: &mut DeltaVel<Real>) {
         if self.motor_inv_lhs != 0.0 {
             let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
             let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
@@ -258,6 +254,14 @@ impl RevoluteVelocityConstraint {
             mj_lambda1.angular += self.ii1_sqrt.transform_vector(self.motor_axis1 * impulse);
             mj_lambda2.angular -= self.ii2_sqrt.transform_vector(self.motor_axis2 * impulse);
         }
+    }
+
+    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
+        let mut mj_lambda1 = mj_lambdas[self.mj_lambda1 as usize];
+        let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
+
+        self.solve_dofs(&mut mj_lambda1, &mut mj_lambda2);
+        self.solve_motors(&mut mj_lambda1, &mut mj_lambda2);
 
         mj_lambdas[self.mj_lambda1 as usize] = mj_lambda1;
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
@@ -452,9 +456,7 @@ impl RevoluteVelocityGroundConstraint {
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }
 
-    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
-        let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
-
+    fn solve_dofs(&mut self, mj_lambda2: &mut DeltaVel<Real>) {
         let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
 
         let lin_dvel = mj_lambda2.linear + ang_vel2.gcross(self.r2);
@@ -470,10 +472,8 @@ impl RevoluteVelocityGroundConstraint {
         mj_lambda2.angular -= self
             .ii2_sqrt
             .transform_vector(ang_impulse + self.r2.gcross(lin_impulse));
-
-        /*
-         * Motor.
-         */
+    }
+    fn solve_motors(&mut self, mj_lambda2: &mut DeltaVel<Real>) {
         if self.motor_inv_lhs != 0.0 {
             let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
             let ang_dvel = ang_vel2.dot(&self.motor_axis2);
@@ -489,6 +489,13 @@ impl RevoluteVelocityGroundConstraint {
 
             mj_lambda2.angular -= self.ii2_sqrt.transform_vector(self.motor_axis2 * impulse);
         }
+    }
+
+    pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
+        let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
+
+        self.solve_dofs(&mut mj_lambda2);
+        self.solve_motors(&mut mj_lambda2);
 
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }