Add restitution.

4 files changed, 34 insertions, 10 deletions

diff --git a/src/dynamics/solver/velocity_constraint.rs b/src/dynamics/solver/velocity_constraint.rs
index 190076d..948081d 100644
--- a/src/dynamics/solver/velocity_constraint.rs
+++ b/src/dynamics/solver/velocity_constraint.rs
@@ -238,8 +238,13 @@ impl VelocityConstraint {
                             + gcross1.gdot(gcross1)
                             + gcross2.gdot(gcross2));
 
-                    let rhs = (vel1 - vel2).dot(&force_dir1)
-                        + manifold_point.dist.max(0.0) * params.inv_dt();
+                    let mut rhs = (vel1 - vel2).dot(&force_dir1);
+
+                    if rhs <= -params.restitution_velocity_threshold {
+                        rhs += manifold.restitution * rhs
+                    }
+
+                    rhs += manifold_point.dist.max(0.0) * params.inv_dt();
 
                     let impulse = manifold_points[k].impulse * warmstart_coeff;
 
diff --git a/src/dynamics/solver/velocity_constraint_wide.rs b/src/dynamics/solver/velocity_constraint_wide.rs
index f515d5e..5d8078a 100644
--- a/src/dynamics/solver/velocity_constraint_wide.rs
+++ b/src/dynamics/solver/velocity_constraint_wide.rs
@@ -104,6 +104,10 @@ impl WVelocityConstraint {
         let mj_lambda2 = array![|ii| rbs2[ii].active_set_offset; SIMD_WIDTH];
 
         let friction = SimdFloat::from(array![|ii| manifolds[ii].friction; SIMD_WIDTH]);
+        let restitution = SimdFloat::from(array![|ii| manifolds[ii].restitution; SIMD_WIDTH]);
+        let restitution_velocity_threshold =
+            SimdFloat::splat(params.restitution_velocity_threshold);
+
         let warmstart_multiplier =
             SimdFloat::from(array![|ii| manifolds[ii].warmstart_multiplier; SIMD_WIDTH]);
         let warmstart_coeff = warmstart_multiplier * SimdFloat::splat(params.warmstart_coeff);
@@ -151,8 +155,11 @@ impl WVelocityConstraint {
 
                     let r = SimdFloat::splat(1.0)
                         / (im1 + im2 + gcross1.gdot(gcross1) + gcross2.gdot(gcross2));
-                    let rhs =
-                        (vel1 - vel2).dot(&force_dir1) + dist.simd_max(SimdFloat::zero()) * inv_dt;
+                    let mut rhs = (vel1 - vel2).dot(&force_dir1);
+                    let use_restitution = rhs.simd_le(-restitution_velocity_threshold);
+                    let rhs_with_restitution = rhs + rhs * restitution;
+                    rhs = rhs_with_restitution.select(use_restitution, rhs);
+                    rhs += dist.simd_max(SimdFloat::zero()) * inv_dt;
 
                     constraint.elements[k].normal_part = WVelocityConstraintElementPart {
                         gcross1,
diff --git a/src/dynamics/solver/velocity_ground_constraint.rs b/src/dynamics/solver/velocity_ground_constraint.rs
index d9229ff..d8ef8be 100644
--- a/src/dynamics/solver/velocity_ground_constraint.rs
+++ b/src/dynamics/solver/velocity_ground_constraint.rs
@@ -169,9 +169,15 @@ impl VelocityGroundConstraint {
                         .transform_vector(dp2.gcross(-force_dir1));
 
                     let r = 1.0 / (rb2.mass_properties.inv_mass + gcross2.gdot(gcross2));
-                    let rhs = -vel2.dot(&force_dir1)
-                        + vel1.dot(&force_dir1)
-                        + manifold_point.dist.max(0.0) * params.inv_dt();
+
+                    let mut rhs = (vel1 - vel2).dot(&force_dir1);
+
+                    if rhs <= -params.restitution_velocity_threshold {
+                        rhs += manifold.restitution * rhs
+                    }
+
+                    rhs += manifold_point.dist.max(0.0) * params.inv_dt();
+
                     let impulse = manifold_points[k].impulse * warmstart_coeff;
 
                     constraint.elements[k].normal_part = VelocityGroundConstraintElementPart {
diff --git a/src/dynamics/solver/velocity_ground_constraint_wide.rs b/src/dynamics/solver/velocity_ground_constraint_wide.rs
index 18e9257..be01944 100644
--- a/src/dynamics/solver/velocity_ground_constraint_wide.rs
+++ b/src/dynamics/solver/velocity_ground_constraint_wide.rs
@@ -110,6 +110,10 @@ impl WVelocityGroundConstraint {
         let mj_lambda2 = array![|ii| rbs2[ii].active_set_offset; SIMD_WIDTH];
 
         let friction = SimdFloat::from(array![|ii| manifolds[ii].friction; SIMD_WIDTH]);
+        let restitution = SimdFloat::from(array![|ii| manifolds[ii].restitution; SIMD_WIDTH]);
+        let restitution_velocity_threshold =
+            SimdFloat::splat(params.restitution_velocity_threshold);
+
         let warmstart_multiplier =
             SimdFloat::from(array![|ii| manifolds[ii].warmstart_multiplier; SIMD_WIDTH]);
         let warmstart_coeff = warmstart_multiplier * SimdFloat::splat(params.warmstart_coeff);
@@ -154,9 +158,11 @@ impl WVelocityGroundConstraint {
                     let gcross2 = ii2.transform_vector(dp2.gcross(-force_dir1));
 
                     let r = SimdFloat::splat(1.0) / (im2 + gcross2.gdot(gcross2));
-                    let rhs = -vel2.dot(&force_dir1)
-                        + vel1.dot(&force_dir1)
-                        + dist.simd_max(SimdFloat::zero()) * inv_dt;
+                    let mut rhs = (vel1 - vel2).dot(&force_dir1);
+                    let use_restitution = rhs.simd_le(-restitution_velocity_threshold);
+                    let rhs_with_restitution = rhs + rhs * restitution;
+                    rhs = rhs_with_restitution.select(use_restitution, rhs);
+                    rhs += dist.simd_max(SimdFloat::zero()) * inv_dt;
 
                     constraint.elements[k].normal_part = WVelocityGroundConstraintElementPart {
                         gcross2,