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

5 files changed, 79 insertions, 108 deletions

diff --git a/src/dynamics/joint/multibody_joint/mod.rs b/src/dynamics/joint/multibody_joint/mod.rs
index a701350..c789004 100644
--- a/src/dynamics/joint/multibody_joint/mod.rs
+++ b/src/dynamics/joint/multibody_joint/mod.rs
@@ -2,7 +2,9 @@
 
 pub use self::multibody::Multibody;
 pub use self::multibody_joint::MultibodyJoint;
-pub use self::multibody_joint_set::{MultibodyIndex, MultibodyJointHandle, MultibodyJointSet};
+pub use self::multibody_joint_set::{
+    MultibodyIndex, MultibodyJointHandle, MultibodyJointSet, MultibodyLinkId,
+};
 pub use self::multibody_link::MultibodyLink;
 pub use self::unit_multibody_joint::{unit_joint_limit_constraint, unit_joint_motor_constraint};
 
diff --git a/src/dynamics/joint/multibody_joint/multibody.rs b/src/dynamics/joint/multibody_joint/multibody.rs
index 63e87e2..c3092ae 100644
--- a/src/dynamics/joint/multibody_joint/multibody.rs
+++ b/src/dynamics/joint/multibody_joint/multibody.rs
@@ -1,16 +1,13 @@
 use super::multibody_link::{MultibodyLink, MultibodyLinkVec};
 use super::multibody_workspace::MultibodyWorkspace;
-use crate::dynamics::{
-    solver::AnyJointVelocityConstraint, IntegrationParameters, RigidBodyHandle, RigidBodySet,
-    RigidBodyType, RigidBodyVelocity,
-};
+use crate::dynamics::{RigidBodyHandle, RigidBodySet, RigidBodyType, RigidBodyVelocity};
 #[cfg(feature = "dim3")]
 use crate::math::Matrix;
 use crate::math::{
     AngDim, AngVector, Dim, Isometry, Jacobian, Point, Real, Vector, ANG_DIM, DIM, SPATIAL_DIM,
 };
 use crate::prelude::MultibodyJoint;
-use crate::utils::{IndexMut2, WAngularInertia, WCross, WCrossMatrix};
+use crate::utils::{IndexMut2, SimdAngularInertia, SimdCross, SimdCrossMatrix};
 use na::{self, DMatrix, DVector, DVectorView, DVectorViewMut, Dyn, OMatrix, SMatrix, SVector, LU};
 
 #[repr(C)]
@@ -372,6 +369,7 @@ impl Multibody {
 
         self.accelerations.fill(0.0);
 
+        // Eqn 42 to 45
         for i in 0..self.links.len() {
             let link = &self.links[i];
             let rb = &bodies[link.rigid_body];
@@ -400,7 +398,7 @@ impl Multibody {
             }
 
             acc.linvel += rb.vels.angvel.gcross(rb.vels.angvel.gcross(link.shift23));
-            acc.linvel += self.workspace.accs[i].angvel.gcross(link.shift23);
+            acc.linvel += acc.angvel.gcross(link.shift23);
 
             self.workspace.accs[i] = acc;
 
@@ -728,7 +726,7 @@ impl Multibody {
 
     /// The generalized velocity at the multibody_joint of the given link.
     #[inline]
-    pub(crate) fn joint_velocity(&self, link: &MultibodyLink) -> DVectorView<Real> {
+    pub fn joint_velocity(&self, link: &MultibodyLink) -> DVectorView<Real> {
         let ndofs = link.joint().ndofs();
         DVectorView::from_slice(
             &self.velocities.as_slice()[link.assembly_id..link.assembly_id + ndofs],
@@ -829,8 +827,10 @@ impl Multibody {
         }
     }
 
+    // TODO: make a version that doesn’t write back to bodies and doesn’t update the jacobians
+    //       (i.e. just something used by the velocity solver’s small steps).
     /// Apply forward-kinematics to this multibody and its related rigid-bodies.
-    pub fn forward_kinematics(&mut self, bodies: &mut RigidBodySet, update_mass_props: bool) {
+    pub fn forward_kinematics(&mut self, bodies: &mut RigidBodySet, update_rb_mass_props: bool) {
         // Special case for the root, which has no parent.
         {
             let link = &mut self.links[0];
@@ -839,7 +839,7 @@ impl Multibody {
 
             if let Some(rb) = bodies.get_mut_internal(link.rigid_body) {
                 rb.pos.next_position = link.local_to_world;
-                if update_mass_props {
+                if update_rb_mass_props {
                     rb.mprops.update_world_mass_properties(&link.local_to_world);
                 }
             }
@@ -873,7 +873,7 @@ impl Multibody {
                 "A rigid-body that is not at the root of a multibody must be dynamic."
             );
 
-            if update_mass_props {
+            if update_rb_mass_props {
                 link_rb
                     .mprops
                     .update_world_mass_properties(&link.local_to_world);
@@ -951,40 +951,4 @@ impl Multibody {
             .sum();
         (num_constraints, num_constraints)
     }
-
-    #[inline]
-    pub(crate) fn generate_internal_constraints(
-        &self,
-        params: &IntegrationParameters,
-        j_id: &mut usize,
-        jacobians: &mut DVector<Real>,
-        out: &mut Vec<AnyJointVelocityConstraint>,
-        mut insert_at: Option<usize>,
-    ) {
-        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,
-                &mut insert_at,
-            );
-        }
-    }
 }
diff --git a/src/dynamics/joint/multibody_joint/multibody_joint.rs b/src/dynamics/joint/multibody_joint/multibody_joint.rs
index da650e6..62fc434 100644
--- a/src/dynamics/joint/multibody_joint/multibody_joint.rs
+++ b/src/dynamics/joint/multibody_joint/multibody_joint.rs
@@ -1,4 +1,4 @@
-use crate::dynamics::solver::AnyJointVelocityConstraint;
+use crate::dynamics::solver::JointGenericOneBodyConstraint;
 use crate::dynamics::{
     joint, FixedJointBuilder, GenericJoint, IntegrationParameters, Multibody, MultibodyLink,
     RigidBodyVelocity,
@@ -254,15 +254,15 @@ impl MultibodyJoint {
         params: &IntegrationParameters,
         multibody: &Multibody,
         link: &MultibodyLink,
-        dof_id: usize,
-        j_id: &mut usize,
+        mut j_id: usize,
         jacobians: &mut DVector<Real>,
-        constraints: &mut Vec<AnyJointVelocityConstraint>,
-        insert_at: &mut Option<usize>,
-    ) {
+        constraints: &mut [JointGenericOneBodyConstraint],
+    ) -> usize {
+        let j_id = &mut j_id;
         let locked_bits = self.data.locked_axes.bits();
         let limit_bits = self.data.limit_axes.bits();
         let motor_bits = self.data.motor_axes.bits();
+        let mut num_constraints = 0;
         let mut curr_free_dof = 0;
 
         for i in 0..DIM {
@@ -281,11 +281,11 @@ impl MultibodyJoint {
                         &self.data.motors[i],
                         self.coords[i],
                         limits,
-                        dof_id + curr_free_dof,
+                        curr_free_dof,
                         j_id,
                         jacobians,
                         constraints,
-                        insert_at,
+                        &mut num_constraints,
                     );
                 }
 
@@ -296,11 +296,11 @@ impl MultibodyJoint {
                         link,
                         [self.data.limits[i].min, self.data.limits[i].max],
                         self.coords[i],
-                        dof_id + curr_free_dof,
+                        curr_free_dof,
                         j_id,
                         jacobians,
                         constraints,
-                        insert_at,
+                        &mut num_constraints,
                     );
                 }
                 curr_free_dof += 1;
@@ -331,11 +331,11 @@ impl MultibodyJoint {
                         link,
                         limits,
                         self.coords[i],
-                        dof_id + curr_free_dof,
+                        curr_free_dof,
                         j_id,
                         jacobians,
                         constraints,
-                        insert_at,
+                        &mut num_constraints,
                     );
                     Some(limits)
                 } else {
@@ -350,15 +350,17 @@ impl MultibodyJoint {
                         &self.data.motors[i],
                         self.coords[i],
                         limits,
-                        dof_id + curr_free_dof,
+                        curr_free_dof,
                         j_id,
                         jacobians,
                         constraints,
-                        insert_at,
+                        &mut num_constraints,
                     );
                 }
                 curr_free_dof += 1;
             }
         }
+
+        num_constraints
     }
 }
diff --git a/src/dynamics/joint/multibody_joint/multibody_joint_set.rs b/src/dynamics/joint/multibody_joint/multibody_joint_set.rs
index a4b338a..a428c8b 100644
--- a/src/dynamics/joint/multibody_joint/multibody_joint_set.rs
+++ b/src/dynamics/joint/multibody_joint/multibody_joint_set.rs
@@ -53,13 +53,24 @@ impl IndexedData for MultibodyJointHandle {
 
 #[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
-pub struct MultibodyJointLink {
-    pub graph_id: RigidBodyGraphIndex,
+/// Indexes usable to get a multibody link from a `MultibodyJointSet`.
+///
+/// ```rust
+/// // With:
+/// //     multibody_joint_set: MultibodyJointSet
+/// //     multibody_link_id: MultibodyLinkId
+/// let multibody = &multibody_joint_set[multibody_link_id.multibody];
+/// let link = multibody.link(multibody_link_id.id).expect("Link not found.");
+pub struct MultibodyLinkId {
+    pub(crate) graph_id: RigidBodyGraphIndex,
+    /// The multibody index to be used as `&multibody_joint_set[multibody]` to
+    /// retrieve the multibody reference.
     pub multibody: MultibodyIndex,
+    /// The multibody link index to be given to [`Multibody::link`].
     pub id: usize,
 }
 
-impl Default for MultibodyJointLink {
+impl Default for MultibodyLinkId {
     fn default() -> Self {
         Self {
             graph_id: RigidBodyGraphIndex::new(crate::INVALID_U32),
@@ -78,7 +89,7 @@ impl Default for MultibodyJointLink {
 #[derive(Clone)]
 pub struct MultibodyJointSet {
     pub(crate) multibodies: Arena<Multibody>, // NOTE: a Slab would be sufficient.
-    pub(crate) rb2mb: Coarena<MultibodyJointLink>,
+    pub(crate) rb2mb: Coarena<MultibodyLinkId>,
     // NOTE: this is mostly for the island extraction. So perhaps we won’t need
     //       that any more in the future when we improve our island builder.
     pub(crate) connectivity_graph: InteractionGraph<RigidBodyHandle, ()>,
@@ -97,13 +108,22 @@ impl MultibodyJointSet {
     }
 
     /// Iterates through all the multibody joints from this set.
-    pub fn iter(&self) -> impl Iterator<Item = (MultibodyJointHandle, &Multibody, &MultibodyLink)> {
+    pub fn iter(
+        &self,
+    ) -> impl Iterator<
+        Item = (
+            MultibodyJointHandle,
+            &MultibodyLinkId,
+            &Multibody,
+            &MultibodyLink,
+        ),
+    > {
         self.rb2mb
             .iter()
             .filter(|(_, link)| link.id > 0) // The first link of a rigid-body hasn’t been added by the user.
             .map(|(h, link)| {
                 let mb = &self.multibodies[link.multibody.0];
-                (MultibodyJointHandle(h), mb, mb.link(link.id).unwrap())
+                (MultibodyJointHandle(h), link, mb, mb.link(link.id).unwrap())
             })
     }
 
@@ -118,7 +138,7 @@ impl MultibodyJointSet {
         let data = data.into();
         let link1 = self.rb2mb.get(body1.0).copied().unwrap_or_else(|| {
             let mb_handle = self.multibodies.insert(Multibody::with_root(body1));
-            MultibodyJointLink {
+            MultibodyLinkId {
                 graph_id: self.connectivity_graph.graph.add_node(body1),
                 multibody: MultibodyIndex(mb_handle),
                 id: 0,
@@ -127,7 +147,7 @@ impl MultibodyJointSet {
 
         let link2 = self.rb2mb.get(body2.0).copied().unwrap_or_else(|| {
             let mb_handle = self.multibodies.insert(Multibody::with_root(body2));
-            MultibodyJointLink {
+            MultibodyLinkId {
                 graph_id: self.connectivity_graph.graph.add_node(body2),
                 multibody: MultibodyIndex(mb_handle),
                 id: 0,
@@ -257,7 +277,7 @@ impl MultibodyJointSet {
     /// Returns the link of this multibody attached to the given rigid-body.
     ///
     /// Returns `None` if `rb` isn’t part of any rigid-body.
-    pub fn rigid_body_link(&self, rb: RigidBodyHandle) -> Option<&MultibodyJointLink> {
+    pub fn rigid_body_link(&self, rb: RigidBodyHandle) -> Option<&MultibodyLinkId> {
         self.rb2mb.get(rb.0)
     }
 
@@ -340,15 +360,15 @@ impl MultibodyJointSet {
         // NOTE: if there is a joint between these two bodies, then
         //       one of the bodies must be the parent of the other.
         let link1 = mb.link(id1.id)?;
-        let parent1 = link1.parent_id()?;
+        let parent1 = link1.parent_id();
 
-        if parent1 == id2.id {
+        if parent1 == Some(id2.id) {
             Some((MultibodyJointHandle(rb1.0), mb, &link1))
         } else {
             let link2 = mb.link(id2.id)?;
-            let parent2 = link2.parent_id()?;
+            let parent2 = link2.parent_id();
 
-            if parent2 == id1.id {
+            if parent2 == Some(id1.id) {
                 Some((MultibodyJointHandle(rb2.0), mb, &link2))
             } else {
                 None
diff --git a/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs b/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs
index a1ec483..d6efd12 100644
--- a/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs
+++ b/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs
@@ -1,9 +1,7 @@
 #![allow(missing_docs)] // For downcast.
 
 use crate::dynamics::joint::MultibodyLink;
-use crate::dynamics::solver::{
-    AnyJointVelocityConstraint, JointGenericVelocityGroundConstraint, WritebackId,
-};
+use crate::dynamics::solver::{JointGenericOneBodyConstraint, WritebackId};
 use crate::dynamics::{IntegrationParameters, JointMotor, Multibody};
 use crate::math::Real;
 use na::DVector;
@@ -19,18 +17,16 @@ pub fn unit_joint_limit_constraint(
     dof_id: usize,
     j_id: &mut usize,
     jacobians: &mut DVector<Real>,
-    constraints: &mut Vec<AnyJointVelocityConstraint>,
-    insert_at: &mut Option<usize>,
+    constraints: &mut [JointGenericOneBodyConstraint],
+    insert_at: &mut usize,
 ) {
     let ndofs = multibody.ndofs();
-    let joint_velocity = multibody.joint_velocity(link);
-
     let min_enabled = curr_pos < limits[0];
     let max_enabled = limits[1] < curr_pos;
     let erp_inv_dt = params.joint_erp_inv_dt();
     let cfm_coeff = params.joint_cfm_coeff();
     let rhs_bias = ((curr_pos - limits[1]).max(0.0) - (limits[0] - curr_pos).max(0.0)) * erp_inv_dt;
-    let rhs_wo_bias = joint_velocity[dof_id];
+    let rhs_wo_bias = 0.0;
 
     let dof_j_id = *j_id + dof_id + link.assembly_id;
     jacobians.rows_mut(*j_id, ndofs * 2).fill(0.0);
@@ -46,8 +42,8 @@ pub fn unit_joint_limit_constraint(
         max_enabled as u32 as Real * Real::MAX,
     ];
 
-    let constraint = JointGenericVelocityGroundConstraint {
-        mj_lambda2: multibody.solver_id,
+    let constraint = JointGenericOneBodyConstraint {
+        solver_vel2: multibody.solver_id,
         ndofs2: ndofs,
         j_id2: *j_id,
         joint_id: usize::MAX,
@@ -61,14 +57,9 @@ pub fn unit_joint_limit_constraint(
         writeback_id: WritebackId::Limit(dof_id),
     };
 
-    if let Some(at) = insert_at {
-        constraints[*at] = AnyJointVelocityConstraint::JointGenericGroundConstraint(constraint);
-        *at += 1;
-    } else {
-        constraints.push(AnyJointVelocityConstraint::JointGenericGroundConstraint(
-            constraint,
-        ));
-    }
+    constraints[*insert_at] = constraint;
+    *insert_at += 1;
+
     *j_id += 2 * ndofs;
 }
 
@@ -84,13 +75,11 @@ pub fn unit_joint_motor_constraint(
     dof_id: usize,
     j_id: &mut usize,
     jacobians: &mut DVector<Real>,
-    constraints: &mut Vec<AnyJointVelocityConstraint>,
-    insert_at: &mut Option<usize>,
+    constraints: &mut [JointGenericOneBodyConstraint],
+    insert_at: &mut usize,
 ) {
     let inv_dt = params.inv_dt();
     let ndofs = multibody.ndofs();
-    let joint_velocity = multibody.joint_velocity(link);
-
     let motor_params = motor.motor_params(params.dt);
 
     let dof_j_id = *j_id + dof_id + link.assembly_id;
@@ -117,11 +106,10 @@ pub fn unit_joint_motor_constraint(
         );
     };
 
-    let dvel = joint_velocity[dof_id];
-    rhs_wo_bias += dvel - target_vel;
+    rhs_wo_bias += -target_vel;
 
-    let constraint = JointGenericVelocityGroundConstraint {
-        mj_lambda2: multibody.solver_id,
+    let constraint = JointGenericOneBodyConstraint {
+        solver_vel2: multibody.solver_id,
         ndofs2: ndofs,
         j_id2: *j_id,
         joint_id: usize::MAX,
@@ -135,13 +123,8 @@ pub fn unit_joint_motor_constraint(
         writeback_id: WritebackId::Limit(dof_id),
     };
 
-    if let Some(at) = insert_at {
-        constraints[*at] = AnyJointVelocityConstraint::JointGenericGroundConstraint(constraint);
-        *at += 1;
-    } else {
-        constraints.push(AnyJointVelocityConstraint::JointGenericGroundConstraint(
-            constraint,
-        ));
-    }
+    constraints[*insert_at] = constraint;
+    *insert_at += 1;
+
     *j_id += 2 * ndofs;
 }