box3d-rust 0.1.0

// Port of motor_joint.c: public accessors, force/torque reporting, and the
// prepare/warm-start/solve simulation functions.
//
// Unlike distance/prismatic/revolute, `b3SolveMotorJoint` takes no `useBias`
// parameter -- springs always use their own softness.
//
// SPDX-FileCopyrightText: 2025 Erin Catto
// SPDX-License-Identifier: MIT

use super::{get_joint_sim_check_type, get_joint_sim_check_type_ref, JointSim, JointType};
use crate::body::{body_flags, BodyState, IDENTITY_BODY_STATE};
use crate::core::NULL_INDEX;
use crate::id::JointId;
use crate::math_functions::{
    add, add_mm, cross, delta_quat_to_rotation, det, dot_quat, inv_mul_quat, invert_matrix,
    length_squared, max_float, mul_add, mul_mm, mul_mv, mul_quat, mul_sub, mul_sv, neg,
    negate_mat3, negate_quat, normalize, rotate_vector, skew, solve3, sub, sub_pos, Vec3,
    QUAT_IDENTITY, VEC3_ZERO,
};
use crate::solver::{make_soft, StepContext};
use crate::solver_set::AWAKE_SET;
use crate::world::World;

/// (b3MotorJoint_SetLinearVelocity)
pub fn motor_joint_set_linear_velocity(world: &mut World, joint_id: JointId, velocity: Vec3) {
    crate::recording::with_recording(world, |rec| {
        rec.write_motor_joint_set_linear_velocity(joint_id, velocity);
    });
    get_joint_sim_check_type(world, joint_id, JointType::Motor)
        .motor_mut()
        .linear_velocity = velocity;
}

/// (b3MotorJoint_GetLinearVelocity)
pub fn motor_joint_get_linear_velocity(world: &World, joint_id: JointId) -> Vec3 {
    get_joint_sim_check_type_ref(world, joint_id, JointType::Motor)
        .motor()
        .linear_velocity
}

/// (b3MotorJoint_SetAngularVelocity)
pub fn motor_joint_set_angular_velocity(world: &mut World, joint_id: JointId, velocity: Vec3) {
    crate::recording::with_recording(world, |rec| {
        rec.write_motor_joint_set_angular_velocity(joint_id, velocity);
    });
    get_joint_sim_check_type(world, joint_id, JointType::Motor)
        .motor_mut()
        .angular_velocity = velocity;
}

/// (b3MotorJoint_GetAngularVelocity)
pub fn motor_joint_get_angular_velocity(world: &World, joint_id: JointId) -> Vec3 {
    get_joint_sim_check_type_ref(world, joint_id, JointType::Motor)
        .motor()
        .angular_velocity
}

/// (b3MotorJoint_SetMaxVelocityTorque)
pub fn motor_joint_set_max_velocity_torque(world: &mut World, joint_id: JointId, max_torque: f32) {
    crate::recording::with_recording(world, |rec| {
        rec.write_motor_joint_set_max_velocity_torque(joint_id, max_torque);
    });
    get_joint_sim_check_type(world, joint_id, JointType::Motor)
        .motor_mut()
        .max_velocity_torque = max_torque;
}

/// (b3MotorJoint_GetMaxVelocityTorque)
pub fn motor_joint_get_max_velocity_torque(world: &World, joint_id: JointId) -> f32 {
    get_joint_sim_check_type_ref(world, joint_id, JointType::Motor)
        .motor()
        .max_velocity_torque
}

/// (b3MotorJoint_SetMaxVelocityForce)
pub fn motor_joint_set_max_velocity_force(world: &mut World, joint_id: JointId, max_force: f32) {
    crate::recording::with_recording(world, |rec| {
        rec.write_motor_joint_set_max_velocity_force(joint_id, max_force);
    });
    get_joint_sim_check_type(world, joint_id, JointType::Motor)
        .motor_mut()
        .max_velocity_force = max_force;
}

/// (b3MotorJoint_GetMaxVelocityForce)
pub fn motor_joint_get_max_velocity_force(world: &World, joint_id: JointId) -> f32 {
    get_joint_sim_check_type_ref(world, joint_id, JointType::Motor)
        .motor()
        .max_velocity_force
}

/// (b3MotorJoint_SetLinearHertz)
pub fn motor_joint_set_linear_hertz(world: &mut World, joint_id: JointId, hertz: f32) {
    crate::recording::with_recording(world, |rec| {
        rec.write_motor_joint_set_linear_hertz(joint_id, hertz);
    });
    get_joint_sim_check_type(world, joint_id, JointType::Motor)
        .motor_mut()
        .linear_hertz = hertz;
}

/// (b3MotorJoint_GetLinearHertz)
pub fn motor_joint_get_linear_hertz(world: &World, joint_id: JointId) -> f32 {
    get_joint_sim_check_type_ref(world, joint_id, JointType::Motor)
        .motor()
        .linear_hertz
}

/// (b3MotorJoint_SetLinearDampingRatio)
pub fn motor_joint_set_linear_damping_ratio(world: &mut World, joint_id: JointId, damping: f32) {
    crate::recording::with_recording(world, |rec| {
        rec.write_motor_joint_set_linear_damping_ratio(joint_id, damping);
    });
    get_joint_sim_check_type(world, joint_id, JointType::Motor)
        .motor_mut()
        .linear_damping_ratio = damping;
}

/// (b3MotorJoint_GetLinearDampingRatio)
pub fn motor_joint_get_linear_damping_ratio(world: &World, joint_id: JointId) -> f32 {
    get_joint_sim_check_type_ref(world, joint_id, JointType::Motor)
        .motor()
        .linear_damping_ratio
}

/// (b3MotorJoint_SetAngularHertz)
pub fn motor_joint_set_angular_hertz(world: &mut World, joint_id: JointId, hertz: f32) {
    crate::recording::with_recording(world, |rec| {
        rec.write_motor_joint_set_angular_hertz(joint_id, hertz);
    });
    get_joint_sim_check_type(world, joint_id, JointType::Motor)
        .motor_mut()
        .angular_hertz = hertz;
}

/// (b3MotorJoint_GetAngularHertz)
pub fn motor_joint_get_angular_hertz(world: &World, joint_id: JointId) -> f32 {
    get_joint_sim_check_type_ref(world, joint_id, JointType::Motor)
        .motor()
        .angular_hertz
}

/// (b3MotorJoint_SetAngularDampingRatio)
pub fn motor_joint_set_angular_damping_ratio(world: &mut World, joint_id: JointId, damping: f32) {
    crate::recording::with_recording(world, |rec| {
        rec.write_motor_joint_set_angular_damping_ratio(joint_id, damping);
    });
    get_joint_sim_check_type(world, joint_id, JointType::Motor)
        .motor_mut()
        .angular_damping_ratio = damping;
}

/// (b3MotorJoint_GetAngularDampingRatio)
pub fn motor_joint_get_angular_damping_ratio(world: &World, joint_id: JointId) -> f32 {
    get_joint_sim_check_type_ref(world, joint_id, JointType::Motor)
        .motor()
        .angular_damping_ratio
}

/// (b3MotorJoint_SetMaxSpringForce)
pub fn motor_joint_set_max_spring_force(world: &mut World, joint_id: JointId, max_force: f32) {
    crate::recording::with_recording(world, |rec| {
        rec.write_motor_joint_set_max_spring_force(joint_id, max_force);
    });
    get_joint_sim_check_type(world, joint_id, JointType::Motor)
        .motor_mut()
        .max_spring_force = max_float(0.0, max_force);
}

/// (b3MotorJoint_GetMaxSpringForce)
pub fn motor_joint_get_max_spring_force(world: &World, joint_id: JointId) -> f32 {
    get_joint_sim_check_type_ref(world, joint_id, JointType::Motor)
        .motor()
        .max_spring_force
}

/// (b3MotorJoint_SetMaxSpringTorque)
pub fn motor_joint_set_max_spring_torque(world: &mut World, joint_id: JointId, max_torque: f32) {
    crate::recording::with_recording(world, |rec| {
        rec.write_motor_joint_set_max_spring_torque(joint_id, max_torque);
    });
    get_joint_sim_check_type(world, joint_id, JointType::Motor)
        .motor_mut()
        .max_spring_torque = max_float(0.0, max_torque);
}

/// (b3MotorJoint_GetMaxSpringTorque)
pub fn motor_joint_get_max_spring_torque(world: &World, joint_id: JointId) -> f32 {
    get_joint_sim_check_type_ref(world, joint_id, JointType::Motor)
        .motor()
        .max_spring_torque
}

/// (b3GetMotorJointForce)
pub fn get_motor_joint_force(world: &World, base: &JointSim) -> Vec3 {
    mul_sv(
        world.inv_h,
        add(
            base.motor().linear_velocity_impulse,
            base.motor().linear_spring_impulse,
        ),
    )
}

/// (b3GetMotorJointTorque)
pub fn get_motor_joint_torque(world: &World, base: &JointSim) -> Vec3 {
    mul_sv(
        world.inv_h,
        add(
            base.motor().angular_velocity_impulse,
            base.motor().angular_spring_impulse,
        ),
    )
}

/// (b3PrepareMotorJoint)
pub fn prepare_motor_joint(world: &World, base: &mut JointSim, context: &StepContext) {
    debug_assert!(base.type_ == JointType::Motor);

    let id_a = base.body_id_a;
    let id_b = base.body_id_b;

    let body_a = &world.bodies[id_a as usize];
    let body_b = &world.bodies[id_b as usize];

    debug_assert!(body_a.set_index == AWAKE_SET || body_b.set_index == AWAKE_SET);

    let body_sim_a =
        &world.solver_sets[body_a.set_index as usize].body_sims[body_a.local_index as usize];
    let body_sim_b =
        &world.solver_sets[body_b.set_index as usize].body_sims[body_b.local_index as usize];

    base.inv_mass_a = body_sim_a.inv_mass;
    base.inv_mass_b = body_sim_b.inv_mass;
    base.inv_i_a = body_sim_a.inv_inertia_world;
    base.inv_i_b = body_sim_b.inv_inertia_world;

    let inv_inertia_sum = add_mm(base.inv_i_a, base.inv_i_b);
    base.fixed_rotation = det(inv_inertia_sum) < 1000.0 * f32::MIN_POSITIVE;

    let local_frame_a = base.local_frame_a;
    let local_frame_b = base.local_frame_b;

    let index_a = if body_a.set_index == AWAKE_SET {
        body_a.local_index
    } else {
        NULL_INDEX
    };
    let index_b = if body_b.set_index == AWAKE_SET {
        body_b.local_index
    } else {
        NULL_INDEX
    };

    let frame_a_q = mul_quat(body_sim_a.transform.q, local_frame_a.q);
    let frame_a_p = rotate_vector(
        body_sim_a.transform.q,
        sub(local_frame_a.p, body_sim_a.local_center),
    );
    let frame_b_q = mul_quat(body_sim_b.transform.q, local_frame_b.q);
    let frame_b_p = rotate_vector(
        body_sim_b.transform.q,
        sub(local_frame_b.p, body_sim_b.local_center),
    );
    let delta_center = sub_pos(body_sim_b.center, body_sim_a.center);
    let angular_mass = invert_matrix(inv_inertia_sum);

    let joint = base.motor_mut();
    joint.index_a = index_a;
    joint.index_b = index_b;
    joint.frame_a.q = frame_a_q;
    joint.frame_a.p = frame_a_p;
    joint.frame_b.q = frame_b_q;
    joint.frame_b.p = frame_b_p;
    joint.delta_center = delta_center;

    joint.linear_spring = make_soft(joint.linear_hertz, joint.linear_damping_ratio, context.h);
    joint.angular_spring = make_soft(joint.angular_hertz, joint.angular_damping_ratio, context.h);
    joint.angular_mass = angular_mass;

    if !context.enable_warm_starting {
        joint.linear_velocity_impulse = VEC3_ZERO;
        joint.angular_velocity_impulse = VEC3_ZERO;
        joint.linear_spring_impulse = VEC3_ZERO;
        joint.angular_spring_impulse = VEC3_ZERO;
    }
}

/// (b3WarmStartMotorJoint)
pub fn warm_start_motor_joint(base: &mut JointSim, states: &mut [BodyState]) {
    debug_assert!(base.type_ == JointType::Motor);

    let m_a = base.inv_mass_a;
    let m_b = base.inv_mass_b;
    let i_a = base.inv_i_a;
    let i_b = base.inv_i_b;

    let joint = base.motor_mut();

    let mut state_a = if joint.index_a == NULL_INDEX {
        IDENTITY_BODY_STATE
    } else {
        states[joint.index_a as usize]
    };
    let mut state_b = if joint.index_b == NULL_INDEX {
        IDENTITY_BODY_STATE
    } else {
        states[joint.index_b as usize]
    };

    let r_a = rotate_vector(state_a.delta_rotation, joint.frame_a.p);
    let r_b = rotate_vector(state_b.delta_rotation, joint.frame_b.p);

    let linear_impulse = add(joint.linear_velocity_impulse, joint.linear_spring_impulse);
    let angular_impulse = add(joint.angular_velocity_impulse, joint.angular_spring_impulse);

    state_a.linear_velocity = mul_sub(state_a.linear_velocity, m_a, linear_impulse);
    state_a.angular_velocity = sub(
        state_a.angular_velocity,
        mul_mv(i_a, add(cross(r_a, linear_impulse), angular_impulse)),
    );
    state_b.linear_velocity = mul_add(state_b.linear_velocity, m_b, linear_impulse);
    state_b.angular_velocity = add(
        state_b.angular_velocity,
        mul_mv(i_b, add(cross(r_b, linear_impulse), angular_impulse)),
    );

    // C writes unconditionally (no dynamic-flag guard) on warm start.
    if joint.index_a != NULL_INDEX {
        states[joint.index_a as usize] = state_a;
    }
    if joint.index_b != NULL_INDEX {
        states[joint.index_b as usize] = state_b;
    }
}

/// (b3SolveMotorJoint)
pub fn solve_motor_joint(base: &mut JointSim, context: &StepContext, states: &mut [BodyState]) {
    debug_assert!(base.type_ == JointType::Motor);

    let m_a = base.inv_mass_a;
    let m_b = base.inv_mass_b;
    let i_a = base.inv_i_a;
    let i_b = base.inv_i_b;

    let joint = base.motor_mut();
    let mut state_a = if joint.index_a == NULL_INDEX {
        IDENTITY_BODY_STATE
    } else {
        states[joint.index_a as usize]
    };
    let mut state_b = if joint.index_b == NULL_INDEX {
        IDENTITY_BODY_STATE
    } else {
        states[joint.index_b as usize]
    };

    let mut v_a = state_a.linear_velocity;
    let mut w_a = state_a.angular_velocity;
    let mut v_b = state_b.linear_velocity;
    let mut w_b = state_b.angular_velocity;

    let quat_a = mul_quat(state_a.delta_rotation, joint.frame_a.q);
    let mut quat_b = mul_quat(state_b.delta_rotation, joint.frame_b.q);

    if dot_quat(quat_a, quat_b) < 0.0 {
        // this keeps the rotation angle in the range [-pi, pi]
        quat_b = negate_quat(quat_b);
    }

    let rel_q = inv_mul_quat(quat_a, quat_b);

    // angular spring
    if joint.max_spring_torque > 0.0 && joint.angular_hertz > 0.0 {
        let target_quat = QUAT_IDENTITY;
        let delta_rotation = delta_quat_to_rotation(rel_q, target_quat);
        let c = neg(rotate_vector(quat_a, delta_rotation));

        let bias = mul_sv(joint.angular_spring.bias_rate, c);
        let mass_scale = joint.angular_spring.mass_scale;
        let impulse_scale = joint.angular_spring.impulse_scale;

        let cdot = sub(w_b, w_a);

        let max_impulse = context.h * joint.max_spring_torque;
        let old_impulse = joint.angular_spring_impulse;
        let mut impulse = mul_sub(
            mul_sv(-mass_scale, mul_mv(joint.angular_mass, add(cdot, bias))),
            impulse_scale,
            old_impulse,
        );
        joint.angular_spring_impulse = add(old_impulse, impulse);
        if length_squared(joint.angular_spring_impulse) > max_impulse * max_impulse {
            joint.angular_spring_impulse =
                mul_sv(max_impulse, normalize(joint.angular_spring_impulse));
        }
        impulse = sub(joint.angular_spring_impulse, old_impulse);

        w_a = sub(w_a, mul_mv(i_a, impulse));
        w_b = add(w_b, mul_mv(i_b, impulse));
    }

    // angular velocity
    if joint.max_velocity_torque > 0.0 {
        let cdot = sub(sub(w_b, w_a), joint.angular_velocity);
        let mut impulse = neg(mul_mv(joint.angular_mass, cdot));

        let max_impulse = context.h * joint.max_velocity_torque;
        let old_impulse = joint.angular_velocity_impulse;
        joint.angular_velocity_impulse = add(old_impulse, impulse);
        if length_squared(joint.angular_velocity_impulse) > max_impulse * max_impulse {
            joint.angular_velocity_impulse =
                mul_sv(max_impulse, normalize(joint.angular_velocity_impulse));
        }
        impulse = sub(joint.angular_velocity_impulse, old_impulse);

        w_a = sub(w_a, mul_mv(i_a, impulse));
        w_b = add(w_b, mul_mv(i_b, impulse));
    }

    let r_a = rotate_vector(state_a.delta_rotation, joint.frame_a.p);
    let r_b = rotate_vector(state_b.delta_rotation, joint.frame_b.p);

    // linear spring
    if joint.max_spring_force > 0.0 && joint.linear_hertz > 0.0 {
        let dc_a = state_a.delta_position;
        let dc_b = state_b.delta_position;
        let c = add(add(sub(dc_b, dc_a), sub(r_b, r_a)), joint.delta_center);

        let bias = mul_sv(joint.linear_spring.bias_rate, c);
        let mass_scale = joint.linear_spring.mass_scale;
        let impulse_scale = joint.linear_spring.impulse_scale;

        let cdot = sub(add(v_b, cross(w_b, r_b)), add(v_a, cross(w_a, r_a)));

        // K = [(1/m1 + 1/m2) * eye(2) - skew(r1) * invI1 * skew(r1) - skew(r2) * invI2 * skew(r2)]
        let s_a = skew(r_a);
        let s_b = skew(r_b);
        let k_a = mul_mm(s_a, mul_mm(i_a, s_a));
        let k_b = mul_mm(s_b, mul_mm(i_b, s_b));
        let mut k = negate_mat3(add_mm(k_a, k_b));
        k.cx.x += m_a + m_b;
        k.cy.y += m_a + m_b;
        k.cz.z += m_a + m_b;

        let b = solve3(k, add(cdot, bias));

        let old_impulse = joint.linear_spring_impulse;
        let mut impulse = mul_sub(mul_sv(-mass_scale, b), impulse_scale, old_impulse);
        let max_impulse = context.h * joint.max_spring_force;
        joint.linear_spring_impulse = add(joint.linear_spring_impulse, impulse);

        if length_squared(joint.linear_spring_impulse) > max_impulse * max_impulse {
            joint.linear_spring_impulse =
                mul_sv(max_impulse, normalize(joint.linear_spring_impulse));
        }

        impulse = sub(joint.linear_spring_impulse, old_impulse);

        v_a = mul_sub(v_a, m_a, impulse);
        w_a = sub(w_a, mul_mv(i_a, cross(r_a, impulse)));
        v_b = mul_add(v_b, m_b, impulse);
        w_b = add(w_b, mul_mv(i_b, cross(r_b, impulse)));
    }

    // linear velocity
    if joint.max_velocity_force > 0.0 {
        let mut cdot = sub(add(v_b, cross(w_b, r_b)), add(v_a, cross(w_a, r_a)));
        cdot = sub(cdot, joint.linear_velocity);
        // K = [(1/m1 + 1/m2) * eye(2) - skew(r1) * invI1 * skew(r1) - skew(r2) * invI2 * skew(r2)]
        let s_a = skew(r_a);
        let s_b = skew(r_b);
        let k_a = mul_mm(s_a, mul_mm(i_a, s_a));
        let k_b = mul_mm(s_b, mul_mm(i_b, s_b));
        let mut k = negate_mat3(add_mm(k_a, k_b));
        k.cx.x += m_a + m_b;
        k.cy.y += m_a + m_b;
        k.cz.z += m_a + m_b;

        let b = solve3(k, cdot);
        let mut impulse = neg(b);

        let old_impulse = joint.linear_velocity_impulse;
        let max_impulse = context.h * joint.max_velocity_force;
        joint.linear_velocity_impulse = add(joint.linear_velocity_impulse, impulse);

        if length_squared(joint.linear_velocity_impulse) > max_impulse * max_impulse {
            joint.linear_velocity_impulse =
                mul_sv(max_impulse, normalize(joint.linear_velocity_impulse));
        }

        impulse = sub(joint.linear_velocity_impulse, old_impulse);

        v_a = mul_sub(v_a, m_a, impulse);
        w_a = sub(w_a, mul_mv(i_a, cross(r_a, impulse)));
        v_b = mul_add(v_b, m_b, impulse);
        w_b = add(w_b, mul_mv(i_b, cross(r_b, impulse)));
    }

    if state_a.flags & body_flags::DYNAMIC_FLAG != 0 {
        state_a.linear_velocity = v_a;
        state_a.angular_velocity = w_a;
        states[joint.index_a as usize] = state_a;
    }

    if state_b.flags & body_flags::DYNAMIC_FLAG != 0 {
        state_b.linear_velocity = v_b;
        state_b.angular_velocity = w_b;
        states[joint.index_b as usize] = state_b;
    }
}