Implement multibody joints and the new solver

1 files changed, 34 insertions, 122 deletions

diff --git a/src/dynamics/solver/parallel_island_solver.rs b/src/dynamics/solver/parallel_island_solver.rs
index bd6fc5d..e695b3d 100644
--- a/src/dynamics/solver/parallel_island_solver.rs
+++ b/src/dynamics/solver/parallel_island_solver.rs
@@ -3,14 +3,14 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use rayon::Scope;
 
 use crate::data::{BundleSet, ComponentSet, ComponentSetMut};
+use crate::dynamics::solver::generic_velocity_constraint::GenericVelocityConstraint;
 use crate::dynamics::solver::{
-    AnyJointPositionConstraint, AnyJointVelocityConstraint, AnyPositionConstraint,
-    AnyVelocityConstraint, ParallelPositionSolver, ParallelSolverConstraints,
+    AnyJointVelocityConstraint, AnyVelocityConstraint, ParallelSolverConstraints,
 };
 use crate::dynamics::{
-    IntegrationParameters, IslandManager, JointGraphEdge, JointIndex, RigidBodyDamping,
-    RigidBodyForces, RigidBodyIds, RigidBodyMassProps, RigidBodyPosition, RigidBodyType,
-    RigidBodyVelocity,
+    IntegrationParameters, IslandManager, JointGraphEdge, JointIndex, MultibodyJointSet,
+    RigidBodyDamping, RigidBodyForces, RigidBodyIds, RigidBodyMassProps, RigidBodyPosition,
+    RigidBodyType, RigidBodyVelocity,
 };
 use crate::geometry::{ContactManifold, ContactManifoldIndex};
 use crate::math::{Isometry, Real};
@@ -69,8 +69,6 @@ pub(crate) struct ThreadContext {
     pub num_initialized_constraints: AtomicUsize,
     pub joint_constraint_initialization_index: AtomicUsize,
     pub num_initialized_joint_constraints: AtomicUsize,
-    pub warmstart_constraint_index: AtomicUsize,
-    pub num_warmstarted_constraints: AtomicUsize,
     pub solve_interaction_index: AtomicUsize,
     pub num_solved_interactions: AtomicUsize,
     pub impulse_writeback_index: AtomicUsize,
@@ -79,10 +77,6 @@ pub(crate) struct ThreadContext {
     pub body_force_integration_index: AtomicUsize,
     pub num_force_integrated_bodies: AtomicUsize,
     pub num_integrated_bodies: AtomicUsize,
-    // Position solver.
-    pub solve_position_interaction_index: AtomicUsize,
-    pub num_solved_position_interactions: AtomicUsize,
-    pub position_writeback_index: AtomicUsize,
 }
 
 impl ThreadContext {
@@ -93,8 +87,6 @@ impl ThreadContext {
             num_initialized_constraints: AtomicUsize::new(0),
             joint_constraint_initialization_index: AtomicUsize::new(0),
             num_initialized_joint_constraints: AtomicUsize::new(0),
-            num_warmstarted_constraints: AtomicUsize::new(0),
-            warmstart_constraint_index: AtomicUsize::new(0),
             solve_interaction_index: AtomicUsize::new(0),
             num_solved_interactions: AtomicUsize::new(0),
             impulse_writeback_index: AtomicUsize::new(0),
@@ -103,9 +95,6 @@ impl ThreadContext {
             num_force_integrated_bodies: AtomicUsize::new(0),
             body_integration_index: AtomicUsize::new(0),
             num_integrated_bodies: AtomicUsize::new(0),
-            solve_position_interaction_index: AtomicUsize::new(0),
-            num_solved_position_interactions: AtomicUsize::new(0),
-            position_writeback_index: AtomicUsize::new(0),
         }
     }
 
@@ -122,14 +111,13 @@ impl ThreadContext {
 }
 
 pub struct ParallelIslandSolver {
-    mj_lambdas: Vec<DeltaVel<Real>>,
+    velocity_solver: ParallelVelocitySolver,
     positions: Vec<Isometry<Real>>,
     parallel_groups: ParallelInteractionGroups,
     parallel_joint_groups: ParallelInteractionGroups,
     parallel_contact_constraints:
-        ParallelSolverConstraints<AnyVelocityConstraint, AnyPositionConstraint>,
-    parallel_joint_constraints:
-        ParallelSolverConstraints<AnyJointVelocityConstraint, AnyJointPositionConstraint>,
+        ParallelSolverConstraints<AnyVelocityConstraint, GenericVelocityConstraint>,
+    parallel_joint_constraints: ParallelSolverConstraints<AnyJointVelocityConstraint, ()>,
     thread: ThreadContext,
 }
 
@@ -142,7 +130,7 @@ impl Default for ParallelIslandSolver {
 impl ParallelIslandSolver {
     pub fn new() -> Self {
         Self {
-            mj_lambdas: Vec::new(),
+            velocity_solver: ParallelVelocitySolver::new(),
             positions: Vec::new(),
             parallel_groups: ParallelInteractionGroups::new(),
             parallel_joint_groups: ParallelInteractionGroups::new(),
@@ -152,84 +140,7 @@ impl ParallelIslandSolver {
         }
     }
 
-    pub fn solve_position_constraints<'s, Bodies>(
-        &'s mut self,
-        scope: &Scope<'s>,
-        island_id: usize,
-        islands: &'s IslandManager,
-        params: &'s IntegrationParameters,
-        bodies: &'s mut Bodies,
-    ) where
-        Bodies: ComponentSetMut<RigidBodyPosition> + ComponentSet<RigidBodyIds>,
-    {
-        let num_threads = rayon::current_num_threads();
-        let num_task_per_island = num_threads; // (num_threads / num_islands).max(1); // TODO: not sure this is the best value. Also, perhaps it is better to interleave tasks of each island?
-        self.thread = ThreadContext::new(8); // TODO: could we compute some kind of optimal value here?
-        self.positions.clear();
-        self.positions
-            .resize(islands.active_island(island_id).len(), Isometry::identity());
-
-        for _ in 0..num_task_per_island {
-            // We use AtomicPtr because it is Send+Sync while *mut is not.
-            // See https://internals.rust-lang.org/t/shouldnt-pointers-be-send-sync-or/8818
-            let thread = &self.thread;
-            let positions = std::sync::atomic::AtomicPtr::new(&mut self.positions as *mut _);
-            let bodies = std::sync::atomic::AtomicPtr::new(bodies as *mut _);
-            let parallel_contact_constraints =
-                std::sync::atomic::AtomicPtr::new(&mut self.parallel_contact_constraints as *mut _);
-            let parallel_joint_constraints =
-                std::sync::atomic::AtomicPtr::new(&mut self.parallel_joint_constraints as *mut _);
-
-            scope.spawn(move |_| {
-                // Transmute *mut -> &mut
-                let positions: &mut Vec<Isometry<Real>> =
-                    unsafe { std::mem::transmute(positions.load(Ordering::Relaxed)) };
-                let bodies: &mut Bodies =
-                    unsafe { std::mem::transmute(bodies.load(Ordering::Relaxed)) };
-                let parallel_contact_constraints: &mut ParallelSolverConstraints<AnyVelocityConstraint, AnyPositionConstraint> = unsafe {
-                    std::mem::transmute(parallel_contact_constraints.load(Ordering::Relaxed))
-                };
-                let parallel_joint_constraints: &mut ParallelSolverConstraints<AnyJointVelocityConstraint, AnyJointPositionConstraint> = unsafe {
-                    std::mem::transmute(parallel_joint_constraints.load(Ordering::Relaxed))
-                };
-
-                enable_flush_to_zero!(); // Ensure this is enabled on each thread.
-
-                // Write results back to rigid bodies and integrate velocities.
-                let island_range = islands.active_island_range(island_id);
-                let active_bodies = &islands.active_dynamic_set[island_range];
-
-                concurrent_loop! {
-                    let batch_size = thread.batch_size;
-                    for handle in active_bodies[thread.body_integration_index, thread.num_integrated_bodies] {
-                        let (rb_ids, rb_pos): (&RigidBodyIds, &RigidBodyPosition) = bodies.index_bundle(handle.0);
-                        positions[rb_ids.active_set_offset] = rb_pos.next_position;
-                    }
-                }
-
-                ThreadContext::lock_until_ge(&thread.num_integrated_bodies, active_bodies.len());
-
-                ParallelPositionSolver::solve(
-                    &thread,
-                    params,
-                    positions,
-                    parallel_contact_constraints,
-                    parallel_joint_constraints
-                );
-
-                // Write results back to rigid bodies.
-                concurrent_loop! {
-                    let batch_size = thread.batch_size;
-                    for handle in active_bodies[thread.position_writeback_index] {
-                        let rb_ids: RigidBodyIds = *bodies.index(handle.0);
-                        bodies.map_mut_internal(handle.0, |rb_pos: &mut RigidBodyPosition| rb_pos.next_position = positions[rb_ids.active_set_offset]);
-                    }
-                }
-            })
-        }
-    }
-
-    pub fn init_constraints_and_solve_velocity_constraints<'s, Bodies>(
+    pub fn init_and_solve<'s, Bodies>(
         &'s mut self,
         scope: &Scope<'s>,
         island_id: usize,
@@ -238,8 +149,9 @@ impl ParallelIslandSolver {
         bodies: &'s mut Bodies,
         manifolds: &'s mut Vec<&'s mut ContactManifold>,
         manifold_indices: &'s [ContactManifoldIndex],
-        joints: &'s mut Vec<JointGraphEdge>,
+        impulse_joints: &'s mut Vec<JointGraphEdge>,
         joint_indices: &[JointIndex],
+        multibody_joints: &mut MultibodyJointSet,
     ) where
         Bodies: ComponentSet<RigidBodyForces>
             + ComponentSetMut<RigidBodyPosition>
@@ -263,13 +175,14 @@ impl ParallelIslandSolver {
             island_id,
             islands,
             bodies,
-            joints,
+            impulse_joints,
             joint_indices,
         );
         self.parallel_contact_constraints.init_constraint_groups(
             island_id,
             islands,
             bodies,
+            multibody_joints,
             manifolds,
             &self.parallel_groups,
         );
@@ -277,25 +190,25 @@ impl ParallelIslandSolver {
             island_id,
             islands,
             bodies,
-            joints,
+            multibody_joints,
+            impulse_joints,
             &self.parallel_joint_groups,
         );
 
-        self.mj_lambdas.clear();
-        self.mj_lambdas
+        self.velocity_solver.mj_lambdas.clear();
+        self.velocity_solver
+            .mj_lambdas
             .resize(islands.active_island(island_id).len(), DeltaVel::zero());
-        self.positions.clear();
-        self.positions
-            .resize(islands.active_island(island_id).len(), Isometry::identity());
 
         for _ in 0..num_task_per_island {
             // We use AtomicPtr because it is Send+Sync while *mut is not.
             // See https://internals.rust-lang.org/t/shouldnt-pointers-be-send-sync-or/8818
             let thread = &self.thread;
-            let mj_lambdas = std::sync::atomic::AtomicPtr::new(&mut self.mj_lambdas as *mut _);
+            let velocity_solver =
+                std::sync::atomic::AtomicPtr::new(&mut self.velocity_solver as *mut _);
             let bodies = std::sync::atomic::AtomicPtr::new(bodies as *mut _);
             let manifolds = std::sync::atomic::AtomicPtr::new(manifolds as *mut _);
-            let joints = std::sync::atomic::AtomicPtr::new(joints as *mut _);
+            let impulse_joints = std::sync::atomic::AtomicPtr::new(impulse_joints as *mut _);
             let parallel_contact_constraints =
                 std::sync::atomic::AtomicPtr::new(&mut self.parallel_contact_constraints as *mut _);
             let parallel_joint_constraints =
@@ -303,18 +216,18 @@ impl ParallelIslandSolver {
 
             scope.spawn(move |_| {
                 // Transmute *mut -> &mut
-                let mj_lambdas: &mut Vec<DeltaVel<Real>> =
-                    unsafe { std::mem::transmute(mj_lambdas.load(Ordering::Relaxed)) };
+                let velocity_solver: &mut ParallelVelocitySolver =
+                    unsafe { std::mem::transmute(velocity_solver.load(Ordering::Relaxed)) };
                 let bodies: &mut Bodies =
                     unsafe { std::mem::transmute(bodies.load(Ordering::Relaxed)) };
                 let manifolds: &mut Vec<&mut ContactManifold> =
                     unsafe { std::mem::transmute(manifolds.load(Ordering::Relaxed)) };
-                let joints: &mut Vec<JointGraphEdge> =
-                    unsafe { std::mem::transmute(joints.load(Ordering::Relaxed)) };
-                let parallel_contact_constraints: &mut ParallelSolverConstraints<AnyVelocityConstraint, AnyPositionConstraint> = unsafe {
+                let impulse_joints: &mut Vec<JointGraphEdge> =
+                    unsafe { std::mem::transmute(impulse_joints.load(Ordering::Relaxed)) };
+                let parallel_contact_constraints: &mut ParallelSolverConstraints<AnyVelocityConstraint, GenericVelocityConstraint> = unsafe {
                     std::mem::transmute(parallel_contact_constraints.load(Ordering::Relaxed))
                 };
-                let parallel_joint_constraints: &mut ParallelSolverConstraints<AnyJointVelocityConstraint, AnyJointPositionConstraint> = unsafe {
+                let parallel_joint_constraints: &mut ParallelSolverConstraints<AnyJointVelocityConstraint, ()> = unsafe {
                     std::mem::transmute(parallel_joint_constraints.load(Ordering::Relaxed))
                 };
 
@@ -329,7 +242,7 @@ impl ParallelIslandSolver {
                         let batch_size = thread.batch_size;
                         for handle in active_bodies[thread.body_force_integration_index, thread.num_force_integrated_bodies] {
                             let (rb_ids, rb_forces, rb_mass_props): (&RigidBodyIds, &RigidBodyForces, &RigidBodyMassProps) = bodies.index_bundle(handle.0);
-                            let dvel = &mut mj_lambdas[rb_ids.active_set_offset];
+                            let dvel = &mut velocity_solver.mj_lambdas[rb_ids.active_set_offset];
 
                             // NOTE: `dvel.angular` is actually storing angular velocity delta multiplied
                             //       by the square root of the inertia tensor:
@@ -345,7 +258,7 @@ impl ParallelIslandSolver {
 
 
                 parallel_contact_constraints.fill_constraints(&thread, params, bodies, manifolds);
-                parallel_joint_constraints.fill_constraints(&thread, params, bodies, joints);
+                parallel_joint_constraints.fill_constraints(&thread, params, bodies, impulse_joints);
                 ThreadContext::lock_until_ge(
                     &thread.num_initialized_constraints,
                     parallel_contact_constraints.constraint_descs.len(),
@@ -355,14 +268,13 @@ impl ParallelIslandSolver {
                     parallel_joint_constraints.constraint_descs.len(),
                 );
 
-                ParallelVelocitySolver::solve(
+                velocity_solver.solve(
                         &thread,
                         params,
                         manifolds,
-                        joints,
-                        mj_lambdas,
+                        impulse_joints,
                         parallel_contact_constraints,
-                        parallel_joint_constraints
+                        parallel_joint_constraints,
                 );
 
                 // Write results back to rigid bodies and integrate velocities.
@@ -383,7 +295,7 @@ impl ParallelIslandSolver {
                         let mut new_rb_pos = *rb_pos;
                         let mut new_rb_vels = *rb_vels;
 
-                        let dvels = mj_lambdas[rb_ids.active_set_offset];
+                        let dvels = velocity_solver.mj_lambdas[rb_ids.active_set_offset];
                         new_rb_vels.linvel += dvels.linear;
                         new_rb_vels.angvel += rb_mprops.effective_world_inv_inertia_sqrt.transform_vector(dvels.angular);