First public release of Rapier.

src/dynamics/solver/velocity_solver.rs

@@ -0,0 +1,405 @@
+use super::{
+    AnyJointVelocityConstraint, InteractionGroups, VelocityConstraint, VelocityGroundConstraint,
+};
+#[cfg(feature = "simd-is-enabled")]
+use super::{WVelocityConstraint, WVelocityGroundConstraint};
+use crate::dynamics::solver::categorization::{categorize_joints, categorize_velocity_contacts};
+use crate::dynamics::{
+    solver::{AnyVelocityConstraint, DeltaVel},
+    IntegrationParameters, JointGraphEdge, JointIndex, RigidBodySet,
+};
+use crate::geometry::{ContactManifold, ContactManifoldIndex};
+#[cfg(feature = "simd-is-enabled")]
+use crate::math::SIMD_WIDTH;
+use crate::utils::WAngularInertia;
+
+pub(crate) struct VelocitySolver {
+    pub mj_lambdas: Vec<DeltaVel<f32>>,
+    pub contact_part: VelocitySolverPart<AnyVelocityConstraint>,
+    pub joint_part: VelocitySolverPart<AnyJointVelocityConstraint>,
+}
+
+impl VelocitySolver {
+    pub fn new() -> Self {
+        Self {
+            mj_lambdas: Vec::new(),
+            contact_part: VelocitySolverPart::new(),
+            joint_part: VelocitySolverPart::new(),
+        }
+    }
+
+    pub fn init_constraints(
+        &mut self,
+        island_id: usize,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        manifolds: &[&mut ContactManifold],
+        manifold_indices: &[ContactManifoldIndex],
+        joints: &[JointGraphEdge],
+        joint_constraint_indices: &[JointIndex],
+    ) {
+        self.contact_part
+            .init_constraints(island_id, params, bodies, manifolds, manifold_indices);
+        self.joint_part.init_constraints(
+            island_id,
+            params,
+            bodies,
+            joints,
+            joint_constraint_indices,
+        )
+    }
+
+    pub fn solve_constraints(
+        &mut self,
+        island_id: usize,
+        params: &IntegrationParameters,
+        bodies: &mut RigidBodySet,
+        manifolds_all: &mut [&mut ContactManifold],
+        joints_all: &mut [JointGraphEdge],
+    ) {
+        self.mj_lambdas.clear();
+        self.mj_lambdas
+            .resize(bodies.active_island(island_id).len(), DeltaVel::zero());
+
+        /*
+         * Warmstart constraints.
+         */
+        for constraint in &self.joint_part.constraints {
+            constraint.warmstart(&mut self.mj_lambdas[..]);
+        }
+
+        for constraint in &self.contact_part.constraints {
+            constraint.warmstart(&mut self.mj_lambdas[..]);
+        }
+
+        /*
+         * Solve constraints.
+         */
+        for _ in 0..params.max_velocity_iterations {
+            for constraint in &mut self.joint_part.constraints {
+                constraint.solve(&mut self.mj_lambdas[..]);
+            }
+
+            for constraint in &mut self.contact_part.constraints {
+                constraint.solve(&mut self.mj_lambdas[..]);
+            }
+        }
+
+        // Update velocities.
+        bodies.foreach_active_island_body_mut_internal(island_id, |_, rb| {
+            let dvel = self.mj_lambdas[rb.active_set_offset];
+            rb.linvel += dvel.linear;
+            rb.angvel += rb.world_inv_inertia_sqrt.transform_vector(dvel.angular);
+        });
+
+        // Write impulses back into the manifold structures.
+        for constraint in &self.joint_part.constraints {
+            constraint.writeback_impulses(joints_all);
+        }
+
+        for constraint in &self.contact_part.constraints {
+            constraint.writeback_impulses(manifolds_all);
+        }
+    }
+}
+
+pub(crate) struct VelocitySolverPart<Constraint> {
+    pub not_ground_interactions: Vec<usize>,
+    pub ground_interactions: Vec<usize>,
+    pub interaction_groups: InteractionGroups,
+    pub ground_interaction_groups: InteractionGroups,
+    pub constraints: Vec<Constraint>,
+}
+
+impl<Constraint> VelocitySolverPart<Constraint> {
+    pub fn new() -> Self {
+        Self {
+            not_ground_interactions: Vec::new(),
+            ground_interactions: Vec::new(),
+            interaction_groups: InteractionGroups::new(),
+            ground_interaction_groups: InteractionGroups::new(),
+            constraints: Vec::new(),
+        }
+    }
+}
+
+impl VelocitySolverPart<AnyVelocityConstraint> {
+    pub fn init_constraint_groups(
+        &mut self,
+        island_id: usize,
+        bodies: &RigidBodySet,
+        manifolds: &[&mut ContactManifold],
+        manifold_indices: &[ContactManifoldIndex],
+    ) {
+        self.not_ground_interactions.clear();
+        self.ground_interactions.clear();
+        categorize_velocity_contacts(
+            bodies,
+            manifolds,
+            manifold_indices,
+            &mut self.ground_interactions,
+            &mut self.not_ground_interactions,
+        );
+
+        self.interaction_groups.clear_groups();
+        self.interaction_groups.group_manifolds(
+            island_id,
+            bodies,
+            manifolds,
+            &self.not_ground_interactions,
+        );
+
+        self.ground_interaction_groups.clear_groups();
+        self.ground_interaction_groups.group_manifolds(
+            island_id,
+            bodies,
+            manifolds,
+            &self.ground_interactions,
+        );
+
+        // NOTE: uncomment this do disable SIMD contact resolution.
+        //        self.interaction_groups
+        //            .nongrouped_interactions
+        //            .append(&mut self.interaction_groups.grouped_interactions);
+        //        self.ground_interaction_groups
+        //            .nongrouped_interactions
+        //            .append(&mut self.ground_interaction_groups.grouped_interactions);
+    }
+
+    pub fn init_constraints(
+        &mut self,
+        island_id: usize,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        manifolds: &[&mut ContactManifold],
+        manifold_indices: &[ContactManifoldIndex],
+    ) {
+        self.init_constraint_groups(island_id, bodies, manifolds, manifold_indices);
+        self.constraints.clear();
+        #[cfg(feature = "simd-is-enabled")]
+        {
+            self.compute_grouped_constraints(params, bodies, manifolds);
+        }
+        self.compute_nongrouped_constraints(params, bodies, manifolds);
+        #[cfg(feature = "simd-is-enabled")]
+        {
+            self.compute_grouped_ground_constraints(params, bodies, manifolds);
+        }
+        self.compute_nongrouped_ground_constraints(params, bodies, manifolds);
+    }
+
+    #[cfg(feature = "simd-is-enabled")]
+    fn compute_grouped_constraints(
+        &mut self,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        manifolds_all: &[&mut ContactManifold],
+    ) {
+        for manifolds_i in self
+            .interaction_groups
+            .grouped_interactions
+            .chunks_exact(SIMD_WIDTH)
+        {
+            let manifold_id = array![|ii| manifolds_i[ii]; SIMD_WIDTH];
+            let manifolds = array![|ii| &*manifolds_all[manifolds_i[ii]]; SIMD_WIDTH];
+            WVelocityConstraint::generate(
+                params,
+                manifold_id,
+                manifolds,
+                bodies,
+                &mut self.constraints,
+                true,
+            );
+        }
+    }
+
+    fn compute_nongrouped_constraints(
+        &mut self,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        manifolds_all: &[&mut ContactManifold],
+    ) {
+        for manifold_i in &self.interaction_groups.nongrouped_interactions {
+            let manifold = &manifolds_all[*manifold_i];
+            VelocityConstraint::generate(
+                params,
+                *manifold_i,
+                manifold,
+                bodies,
+                &mut self.constraints,
+                true,
+            );
+        }
+    }
+
+    #[cfg(feature = "simd-is-enabled")]
+    fn compute_grouped_ground_constraints(
+        &mut self,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        manifolds_all: &[&mut ContactManifold],
+    ) {
+        for manifolds_i in self
+            .ground_interaction_groups
+            .grouped_interactions
+            .chunks_exact(SIMD_WIDTH)
+        {
+            let manifold_id = array![|ii| manifolds_i[ii]; SIMD_WIDTH];
+            let manifolds = array![|ii| &*manifolds_all[manifolds_i[ii]]; SIMD_WIDTH];
+            WVelocityGroundConstraint::generate(
+                params,
+                manifold_id,
+                manifolds,
+                bodies,
+                &mut self.constraints,
+                true,
+            );
+        }
+    }
+
+    fn compute_nongrouped_ground_constraints(
+        &mut self,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        manifolds_all: &[&mut ContactManifold],
+    ) {
+        for manifold_i in &self.ground_interaction_groups.nongrouped_interactions {
+            let manifold = &manifolds_all[*manifold_i];
+            VelocityGroundConstraint::generate(
+                params,
+                *manifold_i,
+                manifold,
+                bodies,
+                &mut self.constraints,
+                true,
+            )
+        }
+    }
+}
+
+impl VelocitySolverPart<AnyJointVelocityConstraint> {
+    pub fn init_constraints(
+        &mut self,
+        island_id: usize,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        joints: &[JointGraphEdge],
+        joint_constraint_indices: &[JointIndex],
+    ) {
+        // Generate constraints for joints.
+        self.not_ground_interactions.clear();
+        self.ground_interactions.clear();
+        categorize_joints(
+            bodies,
+            joints,
+            joint_constraint_indices,
+            &mut self.ground_interactions,
+            &mut self.not_ground_interactions,
+        );
+
+        self.constraints.clear();
+
+        self.interaction_groups.clear_groups();
+        self.interaction_groups.group_joints(
+            island_id,
+            bodies,
+            joints,
+            &self.not_ground_interactions,
+        );
+
+        self.ground_interaction_groups.clear_groups();
+        self.ground_interaction_groups.group_joints(
+            island_id,
+            bodies,
+            joints,
+            &self.ground_interactions,
+        );
+        // NOTE: uncomment this do disable SIMD joint resolution.
+        // self.interaction_groups
+        //     .nongrouped_interactions
+        //     .append(&mut self.interaction_groups.grouped_interactions);
+        // self.ground_interaction_groups
+        //     .nongrouped_interactions
+        //     .append(&mut self.ground_interaction_groups.grouped_interactions);
+
+        self.compute_nongrouped_joint_ground_constraints(params, bodies, joints);
+        #[cfg(feature = "simd-is-enabled")]
+        {
+            self.compute_grouped_joint_ground_constraints(params, bodies, joints);
+        }
+        self.compute_nongrouped_joint_constraints(params, bodies, joints);
+        #[cfg(feature = "simd-is-enabled")]
+        {
+            self.compute_grouped_joint_constraints(params, bodies, joints);
+        }
+    }
+
+    fn compute_nongrouped_joint_ground_constraints(
+        &mut self,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        joints_all: &[JointGraphEdge],
+    ) {
+        for joint_i in &self.ground_interaction_groups.nongrouped_interactions {
+            let joint = &joints_all[*joint_i].weight;
+            let vel_constraint =
+                AnyJointVelocityConstraint::from_joint_ground(params, *joint_i, joint, bodies);
+            self.constraints.push(vel_constraint);
+        }
+    }
+
+    #[cfg(feature = "simd-is-enabled")]
+    fn compute_grouped_joint_ground_constraints(
+        &mut self,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        joints_all: &[JointGraphEdge],
+    ) {
+        for joints_i in self
+            .ground_interaction_groups
+            .grouped_interactions
+            .chunks_exact(SIMD_WIDTH)
+        {
+            let joints_id = array![|ii| joints_i[ii]; SIMD_WIDTH];
+            let joints = array![|ii| &joints_all[joints_i[ii]].weight; SIMD_WIDTH];
+            let vel_constraint = AnyJointVelocityConstraint::from_wide_joint_ground(
+                params, joints_id, joints, bodies,
+            );
+            self.constraints.push(vel_constraint);
+        }
+    }
+
+    fn compute_nongrouped_joint_constraints(
+        &mut self,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        joints_all: &[JointGraphEdge],
+    ) {
+        for joint_i in &self.interaction_groups.nongrouped_interactions {
+            let joint = &joints_all[*joint_i].weight;
+            let vel_constraint =
+                AnyJointVelocityConstraint::from_joint(params, *joint_i, joint, bodies);
+            self.constraints.push(vel_constraint);
+        }
+    }
+
+    #[cfg(feature = "simd-is-enabled")]
+    fn compute_grouped_joint_constraints(
+        &mut self,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        joints_all: &[JointGraphEdge],
+    ) {
+        for joints_i in self
+            .interaction_groups
+            .grouped_interactions
+            .chunks_exact(SIMD_WIDTH)
+        {
+            let joints_id = array![|ii| joints_i[ii]; SIMD_WIDTH];
+            let joints = array![|ii| &joints_all[joints_i[ii]].weight; SIMD_WIDTH];
+            let vel_constraint =
+                AnyJointVelocityConstraint::from_wide_joint(params, joints_id, joints, bodies);
+            self.constraints.push(vel_constraint);
+        }
+    }
+}