Split rigid-bodies and colliders into multiple components

1 files changed, 67 insertions, 56 deletions

diff --git a/src/dynamics/solver/velocity_ground_constraint_wide.rs b/src/dynamics/solver/velocity_ground_constraint_wide.rs
index 4237d99..3aede0b 100644
--- a/src/dynamics/solver/velocity_ground_constraint_wide.rs
+++ b/src/dynamics/solver/velocity_ground_constraint_wide.rs
@@ -2,7 +2,8 @@ use super::{
     AnyVelocityConstraint, DeltaVel, VelocityGroundConstraintElement,
     VelocityGroundConstraintNormalPart,
 };
-use crate::dynamics::{IntegrationParameters, RigidBodySet};
+use crate::data::ComponentSet;
+use crate::dynamics::{IntegrationParameters, RigidBodyIds, RigidBodyMassProps, RigidBodyVelocity};
 use crate::geometry::{ContactManifold, ContactManifoldIndex};
 use crate::math::{
     AngVector, AngularInertia, Point, Real, SimdReal, Vector, DIM, MAX_MANIFOLD_POINTS, SIMD_WIDTH,
@@ -31,52 +32,71 @@ pub(crate) struct WVelocityGroundConstraint {
 }
 
 impl WVelocityGroundConstraint {
-    pub fn generate(
+    pub fn generate<Bodies>(
         params: &IntegrationParameters,
         manifold_id: [ContactManifoldIndex; SIMD_WIDTH],
         manifolds: [&ContactManifold; SIMD_WIDTH],
-        bodies: &RigidBodySet,
+        bodies: &Bodies,
         out_constraints: &mut Vec<AnyVelocityConstraint>,
         push: bool,
-    ) {
+    ) where
+        Bodies: ComponentSet<RigidBodyIds>
+            + ComponentSet<RigidBodyVelocity>
+            + ComponentSet<RigidBodyMassProps>,
+    {
         let inv_dt = SimdReal::splat(params.inv_dt());
         let velocity_solve_fraction = SimdReal::splat(params.velocity_solve_fraction);
         let velocity_based_erp_inv_dt = SimdReal::splat(params.velocity_based_erp_inv_dt());
 
-        let mut rbs1 = array![|ii| &bodies[manifolds[ii].data.body_pair.body1]; SIMD_WIDTH];
-        let mut rbs2 = array![|ii| &bodies[manifolds[ii].data.body_pair.body2]; SIMD_WIDTH];
+        let mut handles1 = gather![|ii| manifolds[ii].data.rigid_body1];
+        let mut handles2 = gather![|ii| manifolds[ii].data.rigid_body2];
         let mut flipped = [1.0; SIMD_WIDTH];
 
         for ii in 0..SIMD_WIDTH {
             if manifolds[ii].data.relative_dominance < 0 {
-                std::mem::swap(&mut rbs1[ii], &mut rbs2[ii]);
+                std::mem::swap(&mut handles1[ii], &mut handles2[ii]);
                 flipped[ii] = -1.0;
             }
         }
 
+        let vels1: [RigidBodyVelocity; SIMD_WIDTH] = gather![|ii| {
+            handles1[ii]
+                .map(|h| *bodies.index(h.0))
+                .unwrap_or_else(RigidBodyVelocity::zero)
+        }];
+        let world_com1 = Point::from(gather![|ii| {
+            handles1[ii]
+                .map(|h| ComponentSet::<RigidBodyMassProps>::index(bodies, h.0).world_com)
+                .unwrap_or_else(Point::origin)
+        }]);
+
+        let vels2: [&RigidBodyVelocity; SIMD_WIDTH] =
+            gather![|ii| bodies.index(handles2[ii].unwrap().0)];
+        let ids2: [&RigidBodyIds; SIMD_WIDTH] = gather![|ii| bodies.index(handles2[ii].unwrap().0)];
+        let mprops2: [&RigidBodyMassProps; SIMD_WIDTH] =
+            gather![|ii| bodies.index(handles2[ii].unwrap().0)];
+
         let flipped_sign = SimdReal::from(flipped);
 
-        let im2 = SimdReal::from(array![|ii| rbs2[ii].effective_inv_mass; SIMD_WIDTH]);
-        let ii2: AngularInertia<SimdReal> = AngularInertia::from(
-            array![|ii| rbs2[ii].effective_world_inv_inertia_sqrt; SIMD_WIDTH],
-        );
+        let im2 = SimdReal::from(gather![|ii| mprops2[ii].effective_inv_mass]);
+        let ii2: AngularInertia<SimdReal> =
+            AngularInertia::from(gather![|ii| mprops2[ii].effective_world_inv_inertia_sqrt]);
 
-        let linvel1 = Vector::from(array![|ii| rbs1[ii].linvel; SIMD_WIDTH]);
-        let angvel1 = AngVector::<SimdReal>::from(array![|ii| rbs1[ii].angvel; SIMD_WIDTH]);
+        let linvel1 = Vector::from(gather![|ii| vels1[ii].linvel]);
+        let angvel1 = AngVector::<SimdReal>::from(gather![|ii| vels1[ii].angvel]);
 
-        let linvel2 = Vector::from(array![|ii| rbs2[ii].linvel; SIMD_WIDTH]);
-        let angvel2 = AngVector::<SimdReal>::from(array![|ii| rbs2[ii].angvel; SIMD_WIDTH]);
+        let linvel2 = Vector::from(gather![|ii| vels2[ii].linvel]);
+        let angvel2 = AngVector::<SimdReal>::from(gather![|ii| vels2[ii].angvel]);
 
-        let world_com1 = Point::from(array![|ii| rbs1[ii].world_com; SIMD_WIDTH]);
-        let world_com2 = Point::from(array![|ii| rbs2[ii].world_com; SIMD_WIDTH]);
+        let world_com2 = Point::from(gather![|ii| mprops2[ii].world_com]);
 
-        let normal1 = Vector::from(array![|ii| manifolds[ii].data.normal; SIMD_WIDTH]);
+        let normal1 = Vector::from(gather![|ii| manifolds[ii].data.normal]);
         let force_dir1 = normal1 * -flipped_sign;
 
-        let mj_lambda2 = array![|ii| rbs2[ii].active_set_offset; SIMD_WIDTH];
+        let mj_lambda2 = gather![|ii| ids2[ii].active_set_offset];
 
         let warmstart_multiplier =
-            SimdReal::from(array![|ii| manifolds[ii].data.warmstart_multiplier; SIMD_WIDTH]);
+            SimdReal::from(gather![|ii| manifolds[ii].data.warmstart_multiplier]);
         let warmstart_coeff = warmstart_multiplier * SimdReal::splat(params.warmstart_coeff);
         let warmstart_correction_slope = SimdReal::splat(params.warmstart_correction_slope);
         let num_active_contacts = manifolds[0].data.num_active_contacts();
@@ -88,7 +108,7 @@ impl WVelocityGroundConstraint {
             super::compute_tangent_contact_directions(&force_dir1, &linvel1, &linvel2);
 
         for l in (0..num_active_contacts).step_by(MAX_MANIFOLD_POINTS) {
-            let manifold_points = array![|ii| &manifolds[ii].data.solver_contacts[l..]; SIMD_WIDTH];
+            let manifold_points = gather![|ii| &manifolds[ii].data.solver_contacts[l..]];
             let num_points = manifold_points[0].len().min(MAX_MANIFOLD_POINTS);
 
             let mut constraint = WVelocityGroundConstraint {
@@ -107,24 +127,20 @@ impl WVelocityGroundConstraint {
             };
 
             for k in 0..num_points {
-                let friction =
-                    SimdReal::from(array![|ii| manifold_points[ii][k].friction; SIMD_WIDTH]);
-                let restitution =
-                    SimdReal::from(array![|ii| manifold_points[ii][k].restitution; SIMD_WIDTH]);
-                let is_bouncy = SimdReal::from(
-                    array![|ii| manifold_points[ii][k].is_bouncy() as u32 as Real; SIMD_WIDTH],
-                );
+                let friction = SimdReal::from(gather![|ii| manifold_points[ii][k].friction]);
+                let restitution = SimdReal::from(gather![|ii| manifold_points[ii][k].restitution]);
+                let is_bouncy = SimdReal::from(gather![
+                    |ii| manifold_points[ii][k].is_bouncy() as u32 as Real
+                ]);
                 let is_resting = SimdReal::splat(1.0) - is_bouncy;
-                let point = Point::from(array![|ii| manifold_points[ii][k].point; SIMD_WIDTH]);
-                let dist = SimdReal::from(array![|ii| manifold_points[ii][k].dist; SIMD_WIDTH]);
+                let point = Point::from(gather![|ii| manifold_points[ii][k].point]);
+                let dist = SimdReal::from(gather![|ii| manifold_points[ii][k].dist]);
                 let tangent_velocity =
-                    Vector::from(array![|ii| manifold_points[ii][k].tangent_velocity; SIMD_WIDTH]);
+                    Vector::from(gather![|ii| manifold_points[ii][k].tangent_velocity]);
 
-                let impulse = SimdReal::from(
-                    array![|ii| manifold_points[ii][k].warmstart_impulse; SIMD_WIDTH],
-                );
-                let prev_rhs =
-                    SimdReal::from(array![|ii| manifold_points[ii][k].prev_rhs; SIMD_WIDTH]);
+                let impulse =
+                    SimdReal::from(gather![|ii| manifold_points[ii][k].warmstart_impulse]);
+                let prev_rhs = SimdReal::from(gather![|ii| manifold_points[ii][k].prev_rhs]);
                 let dp1 = point - world_com1;
                 let dp2 = point - world_com2;
 
@@ -133,8 +149,7 @@ impl WVelocityGroundConstraint {
                 let warmstart_correction;
 
                 constraint.limit = friction;
-                constraint.manifold_contact_id[k] =
-                    array![|ii| manifold_points[ii][k].contact_id; SIMD_WIDTH];
+                constraint.manifold_contact_id[k] = gather![|ii| manifold_points[ii][k].contact_id];
 
                 // Normal part.
                 {
@@ -162,14 +177,14 @@ impl WVelocityGroundConstraint {
 
                 // tangent parts.
                 #[cfg(feature = "dim2")]
-                let impulse = [SimdReal::from(
-                    array![|ii| manifold_points[ii][k].warmstart_tangent_impulse; SIMD_WIDTH],
-                ) * warmstart_correction];
+                let impulse = [SimdReal::from(gather![
+                    |ii| manifold_points[ii][k].warmstart_tangent_impulse
+                ]) * warmstart_correction];
                 #[cfg(feature = "dim3")]
                 let impulse = tangent_rot1
-                    * na::Vector2::from(
-                        array![|ii| manifold_points[ii][k].warmstart_tangent_impulse; SIMD_WIDTH],
-                    )
+                    * na::Vector2::from(gather![
+                        |ii| manifold_points[ii][k].warmstart_tangent_impulse
+                    ])
                     * warmstart_correction;
                 constraint.elements[k].tangent_part.impulse = impulse;
 
@@ -195,12 +210,10 @@ impl WVelocityGroundConstraint {
 
     pub fn warmstart(&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],
-            ),
+            linear: Vector::from(gather![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear]),
+            angular: AngVector::from(gather![
+                |ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular
+            ]),
         };
 
         VelocityGroundConstraintElement::warmstart_group(
@@ -220,12 +233,10 @@ impl WVelocityGroundConstraint {
 
     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],
-            ),
+            linear: Vector::from(gather![|ii| mj_lambdas[self.mj_lambda2[ii] as usize].linear]),
+            angular: AngVector::from(gather![
+                |ii| mj_lambdas[self.mj_lambda2[ii] as usize].angular
+            ]),
         };
 
         VelocityGroundConstraintElement::solve_group(