ksys/phys: Add remaining RigidBodyMotionProxy functions (except one)

4 files changed, 219 insertions, 19 deletions

diff --git a/src/KingSystem/Physics/RigidBody/physRigidBody.h b/src/KingSystem/Physics/RigidBody/physRigidBody.h
index d3879e82..c2a248ac 100644
--- a/src/KingSystem/Physics/RigidBody/physRigidBody.h
+++ b/src/KingSystem/Physics/RigidBody/physRigidBody.h
@@ -2,6 +2,7 @@
 
 #include <container/seadPtrArray.h>
 #include <heap/seadDisposer.h>
+#include <math/seadMathCalcCommon.h>
 #include <prim/seadRuntimeTypeInfo.h>
 #include <prim/seadTypedBitFlag.h>
 #include <thread/seadAtomic.h>
@@ -17,6 +18,7 @@ class hkpMotion;
 namespace ksys::phys {
 
 class MotionAccessor;
+class RigidBodyMotion;
 class RigidContactPoints;
 class UserTag;
 
@@ -131,8 +133,9 @@ public:
     bool x_6();
 
     // 0x0000007100f8d680
-    // FIXME: rename after we figure out what the two types of MotionAccessor are
-    MotionAccessor* getMotionAccessorType1();
+    RigidBodyMotion* getMotionAccessor() const;
+    // 0x0000007100f90f28 - for internal use
+    RigidBodyMotion* getMotionAccessorForProxy() const;
     // 0x0000007100f8d70c
     void* getMotionAccessorType2Stuff();
     // 0x0000007100f8d7a8
@@ -187,14 +190,14 @@ public:
     void setTransform(const sead::Matrix34f& mtx, bool propagate_to_linked_motions);
 
     // 0x0000007100f8ec3c
-    bool setLinearVelocity(const sead::Vector3f& velocity, float epsilon);
+    bool setLinearVelocity(const sead::Vector3f& velocity, float epsilon = sead::Mathf::epsilon());
     // 0x0000007100f9118c
     void getLinearVelocity(sead::Vector3f* velocity) const;
     // 0x0000007100f911ac
     sead::Vector3f getLinearVelocity() const;
 
     // 0x0000007100f8ed74
-    bool setAngularVelocity(const sead::Vector3f& velocity, float epsilon);
+    bool setAngularVelocity(const sead::Vector3f& velocity, float epsilon = sead::Mathf::epsilon());
     // 0x0000007100f911f8
     void getAngularVelocity(sead::Vector3f* velocity) const;
     // 0x0000007100f91218
@@ -237,6 +240,30 @@ public:
     Type getType() const { return mType; }
     bool isCharacterControllerType() const { return mType == Type::CharacterController; }
 
+    // 0x0000007100f969a0
+    void lock(bool also_lock_world);
+    // 0x0000007100f969e8
+    void unlock(bool also_unlock_world);
+
+    class ScopedLock {
+    public:
+        explicit ScopedLock(RigidBody* body, bool also_lock_world)
+            : mBody(body), mAlsoLockWorld(also_lock_world) {
+            mBody->lock(also_lock_world);
+        }
+        ~ScopedLock() { mBody->unlock(mAlsoLockWorld); }
+        ScopedLock(const ScopedLock&) = delete;
+        auto operator=(const ScopedLock&) = delete;
+
+    private:
+        RigidBody* mBody;
+        bool mAlsoLockWorld;
+    };
+
+    [[nodiscard]] auto makeScopedLock(bool also_lock_world) {
+        return ScopedLock(this, also_lock_world);
+    }
+
 private:
     sead::CriticalSection mCS;
     sead::TypedBitFlag<Flag, sead::Atomic<u32>> mFlags{};
diff --git a/src/KingSystem/Physics/RigidBody/physRigidBodyMotion.cpp b/src/KingSystem/Physics/RigidBody/physRigidBodyMotion.cpp
index 5d1dba4d..de49341a 100644
--- a/src/KingSystem/Physics/RigidBody/physRigidBodyMotion.cpp
+++ b/src/KingSystem/Physics/RigidBody/physRigidBodyMotion.cpp
@@ -52,7 +52,7 @@ void RigidBodyMotion::setTransform(const sead::Matrix34f& mtx, bool propagate_to
     if (propagate_to_linked_motions) {
         for (int i = 0, n = mLinkedAccessors.size(); i < n; ++i) {
             auto* accessor = mLinkedAccessors[i];
-            accessor->setTransformMaybe(mtx);
+            accessor->setTransformFromLinkedBody(mtx);
             accessor->setLinearVelocity(sead::Vector3f::zero, sead::Mathf::epsilon());
             accessor->setAngularVelocity(sead::Vector3f::zero, sead::Mathf::epsilon());
         }
@@ -76,7 +76,7 @@ void RigidBodyMotion::setPosition(const sead::Vector3f& position,
     if (propagate_to_linked_motions) {
         for (int i = 0, n = mLinkedAccessors.size(); i < n; ++i) {
             auto* accessor = mLinkedAccessors[i];
-            accessor->setTransformMaybe(hk_position, hk_rotate);
+            accessor->setTransformFromLinkedBody(hk_position, hk_rotate);
         }
     }
 }
@@ -578,9 +578,9 @@ void RigidBodyMotion::freeze(bool freeze, bool preserve_velocities, bool preserv
     mGravityFactor = mBody->getGravityFactor();
     mMaxImpulseCopy = mMaxImpulse;
 
-    mBody->setLinearVelocity(sead::Vector3f::zero, sead::Mathf::epsilon());
+    mBody->setLinearVelocity(sead::Vector3f::zero);
     mBody->setLinearDamping(1.0);
-    mBody->setAngularVelocity(sead::Vector3f::zero, sead::Mathf::epsilon());
+    mBody->setAngularVelocity(sead::Vector3f::zero);
     mBody->setAngularDamping(1.0);
 
     mBody->setInertiaLocal(mInertiaLocal * mass_factor);
diff --git a/src/KingSystem/Physics/RigidBody/physRigidBodyMotionProxy.cpp b/src/KingSystem/Physics/RigidBody/physRigidBodyMotionProxy.cpp
index aa0db452..e0ebcdc3 100644
--- a/src/KingSystem/Physics/RigidBody/physRigidBodyMotionProxy.cpp
+++ b/src/KingSystem/Physics/RigidBody/physRigidBodyMotionProxy.cpp
@@ -1,12 +1,17 @@
 #include "KingSystem/Physics/RigidBody/physRigidBodyMotionProxy.h"
 #include <Havok/Common/Base/hkBase.h>
 #include <Havok/Physics2012/Dynamics/Entity/hkpRigidBody.h>
+#include "KingSystem/Physics/RigidBody/physRigidBodyMotion.h"
 #include "KingSystem/Physics/physConversions.h"
 
 namespace ksys::phys {
 
 RigidBodyMotionProxy::RigidBodyMotionProxy(RigidBody* body) : MotionAccessor(body) {}
 
+RigidBodyMotionProxy::~RigidBodyMotionProxy() {
+    resetLinkedRigidBody();
+}
+
 bool RigidBodyMotionProxy::init(const RigidBodyInstanceParam& params, sead::Heap* heap) {
     mMaxLinearVelocity = params.max_linear_velocity;
     mMaxAngularVelocity = params.max_angular_velocity_rad;
@@ -41,6 +46,109 @@ void RigidBodyMotionProxy::setPosition(const sead::Vector3f& position,
     setTransformImpl(mTransform);
 }
 
+void RigidBodyMotionProxy::setTransformFromLinkedBody(const sead::Matrix34f& mtx) {
+    if (mFlags.isOff(Flag::_80000) && mFlags.isOff(Flag::_100000)) {
+        setTransform(mtx, false);
+        return;
+    }
+
+    sead::Matrix34f new_mtx = mtx;
+
+    sead::Vector3f translate;
+    if (mFlags.isOn(Flag::_80000)) {
+        translate.setMul(mtx, mLinkedBodyExtraTranslate);
+    } else {
+        mtx.getTranslation(translate);
+    }
+
+    if (mFlags.isOn(Flag::_100000)) {
+        sead::Quatf quat;
+        mtx.toQuat(quat);
+        quat *= mLinkedBodyExtraRotate;
+        new_mtx.fromQuat(quat);
+    }
+
+    // This must be done after fromQuat() because fromQuat resets the translation component.
+    new_mtx.setTranslation(translate);
+    setTransform(new_mtx, false);
+}
+
+static void makeSRT(sead::Matrix34f& o, const sead::Vector3f& s, const sead::Quatf& r,
+                    const sead::Vector3f& t) {
+    const float yy = 2 * r.y * r.y;
+    const float zz = 2 * r.z * r.z;
+    const float xx = 2 * r.x * r.x;
+    const float xy = 2 * r.x * r.y;
+    const float xz = 2 * r.x * r.z;
+    const float yz = 2 * r.y * r.z;
+    const float wz = 2 * r.w * r.z;
+    const float wx = 2 * r.w * r.x;
+    const float wy = 2 * r.w * r.y;
+
+    o.m[0][0] = s.x * (1 - yy - zz);
+    o.m[0][1] = s.y * (xy - wz);
+    o.m[0][2] = s.z * (xz + wy);
+
+    o.m[1][0] = s.x * (xy + wz);
+    o.m[1][1] = s.y * (1 - xx - zz);
+    o.m[1][2] = s.z * (yz - wx);
+
+    o.m[2][0] = s.x * (xz - wy);
+    o.m[2][1] = s.y * (yz + wx);
+    o.m[2][2] = s.z * (1 - xx - yy);
+
+    o.m[0][3] = t.x;
+    o.m[1][3] = t.y;
+    o.m[2][3] = t.z;
+}
+
+static sead::Matrix34f makeSRT(const hkVector4f& s, const hkQuaternionf& r, const hkVector4f& t) {
+    sead::Vector3f ss;
+    toVec3(&ss, s);
+
+    sead::Quatf rr;
+    toQuat(&rr, r);
+
+    sead::Vector3f tt;
+    toVec3(&tt, t);
+
+    sead::Matrix34f o;
+    makeSRT(o, ss, rr, tt);
+    return o;
+}
+
+static sead::Matrix34f makeRT(const hkQuaternionf& r, const hkVector4f& t) {
+    return makeSRT(hkVector4f::getConstant<HK_QUADREAL_1>(), r, t);
+}
+
+// NON_MATCHING: two fmul instructions reordered in sead::Matrix34f mtx = makeRT(...)
+void RigidBodyMotionProxy::setTransformFromLinkedBody(const hkVector4f& hk_translate,
+                                                      const hkQuaternionf& hk_rotate) {
+    if (mFlags.isOff(Flag::_80000) && mFlags.isOff(Flag::_100000)) {
+        setTransform(makeRT(hk_rotate, hk_translate), false);
+        return;
+    }
+
+    sead::Matrix34f mtx = makeRT(hk_rotate, hk_translate);
+
+    sead::Vector3f translate;
+    if (mFlags.isOn(Flag::_80000)) {
+        translate.setMul(mtx, mLinkedBodyExtraTranslate);
+    } else {
+        storeToVec3(&translate, hk_translate);
+    }
+
+    if (mFlags.isOn(Flag::_100000)) {
+        sead::Quatf quat;
+        toQuat(&quat, hk_rotate);
+        quat *= mLinkedBodyExtraRotate;
+        mtx.fromQuat(quat);
+    }
+
+    mtx.setTranslation(translate);
+    setTransform(mtx, false);
+}
+
 void RigidBodyMotionProxy::getPosition(sead::Vector3f* position) {
     if (hasMotionFlagSet(RigidBody::MotionFlag::_20)) {
         mTransform.getTranslation(*position);
@@ -162,6 +270,48 @@ float RigidBodyMotionProxy::getMaxAngularVelocity() {
     return getRigidBodyMotion()->getMotionState()->m_maxAngularVelocity;
 }
 
+void RigidBodyMotionProxy::setLinkedRigidBody(RigidBody* body) {
+    auto lock = mBody->makeScopedLock(mBody->isFlag8Set());
+
+    if (mLinkedRigidBody == body)
+        return;
+
+    if (mLinkedRigidBody) {
+        // If we already have a linked rigid body, unlink it first.
+        mLinkedRigidBody->getMotionAccessorForProxy()->deregisterAccessor(this);
+        mLinkedRigidBody = nullptr;
+    }
+
+    if (body) {
+        if (!body->hasFlag(RigidBody::Flag::MassScaling) &&
+            mFlags.isOff(Flag::HasLinkedRigidBodyWithoutFlag10)) {
+            RigidBodyMotion* accessor = body->getMotionAccessorForProxy();
+            if (accessor && accessor->registerAccessor(this)) {
+                mLinkedRigidBody = body;
+                if (mBody->hasFlag(RigidBody::Flag::_10))
+                    mFlags.reset(Flag::HasLinkedRigidBodyWithoutFlag10);
+                else
+                    mFlags.set(Flag::HasLinkedRigidBodyWithoutFlag10);
+            }
+        }
+    } else {
+        mLinkedRigidBody = nullptr;
+        mFlags.reset(Flag::HasLinkedRigidBodyWithoutFlag10);
+    }
+}
+
+void RigidBodyMotionProxy::resetLinkedRigidBody() {
+    if (!mLinkedRigidBody)
+        return;
+
+    auto lock = mBody->makeScopedLock(mBody->isFlag8Set());
+    if (mLinkedRigidBody) {
+        mLinkedRigidBody->getMotionAccessorForProxy()->deregisterAccessor(this);
+        mLinkedRigidBody = nullptr;
+        mFlags.reset(Flag::HasLinkedRigidBodyWithoutFlag10);
+    }
+}
+
 RigidBody* RigidBodyMotionProxy::getLinkedRigidBody() const {
     return mLinkedRigidBody;
 }
@@ -185,4 +335,24 @@ float RigidBodyMotionProxy::getTimeFactor() {
     return mTimeFactor;
 }
 
+void RigidBodyMotionProxy::freeze(bool freeze, bool preserve_velocities,
+                                  bool preserve_max_impulse) {
+    if (!freeze) {
+        mBody->setLinearVelocity(mFrozenLinearVelocity);
+        mBody->setAngularVelocity(mFrozenAngularVelocity);
+        return;
+    }
+
+    if (preserve_velocities) {
+        mFrozenLinearVelocity = mBody->getLinearVelocity();
+        mFrozenAngularVelocity = mBody->getAngularVelocity();
+    } else {
+        mFrozenLinearVelocity.set(0, 0, 0);
+        mFrozenAngularVelocity.set(0, 0, 0);
+    }
+
+    mBody->setLinearVelocity(sead::Vector3f::zero);
+    mBody->setAngularVelocity(sead::Vector3f::zero);
+}
+
 }  // namespace ksys::phys
diff --git a/src/KingSystem/Physics/RigidBody/physRigidBodyMotionProxy.h b/src/KingSystem/Physics/RigidBody/physRigidBodyMotionProxy.h
index 50fb1a59..d7d010c7 100644
--- a/src/KingSystem/Physics/RigidBody/physRigidBodyMotionProxy.h
+++ b/src/KingSystem/Physics/RigidBody/physRigidBodyMotionProxy.h
@@ -12,16 +12,17 @@ class RigidBodyMotionProxy : public MotionAccessor {
 public:
     enum class Flag {
         _40000 = 1 << 18,
+        _80000 = 1 << 19,
+        _100000 = 1 << 20,
+        HasLinkedRigidBodyWithoutFlag10 = 1 << 21,
     };
 
     explicit RigidBodyMotionProxy(RigidBody* body);
 
     void setTransform(const sead::Matrix34f& mtx, bool propagate_to_linked_motions) override;
     void setPosition(const sead::Vector3f& position, bool propagate_to_linked_motions) override;
-    // 0x0000007100fa4318
-    void setTransformMaybe(const sead::Matrix34f& mtx);
-    // 0x0000007100fa4594
-    void setTransformMaybe(const hkVector4f& translate, const hkQuaternionf& rotate);
+    void setTransformFromLinkedBody(const sead::Matrix34f& mtx);
+    void setTransformFromLinkedBody(const hkVector4f& hk_translate, const hkQuaternionf& hk_rotate);
     void getPosition(sead::Vector3f* position) override;
     void getRotation(sead::Quatf* rotation) override;
     void getTransform(sead::Matrix34f* mtx) override;
@@ -41,9 +42,7 @@ public:
     void setMaxAngularVelocity(float max) override;
     float getMaxAngularVelocity() override;
 
-    // 0x0000007100fa4f48 - called from RigidBody, sets a secondary rigid body
     void setLinkedRigidBody(RigidBody* body);
-    // 0x0000007100fa3f8c - called from RigidBody and inlined by the destructor
     void resetLinkedRigidBody();
     RigidBody* getLinkedRigidBody() const;
     bool isFlag40000Set() const;
@@ -56,14 +55,18 @@ public:
     void getRotation(hkQuaternionf* quat) override;
     void setTimeFactor(float factor) override;
     float getTimeFactor() override;
+
     void freeze(bool freeze, bool preserve_velocities, bool preserve_max_impulse) override;
-    void resetFrozenState() override;
+    void resetFrozenState() override {
+        mFrozenLinearVelocity.set(0, 0, 0);
+        mFrozenAngularVelocity.set(0, 0, 0);
+    }
 
 private:
     void setTransformImpl(const sead::Matrix34f& mtx);
 
-    sead::Vector3f _18 = sead::Vector3f::zero;
-    sead::Vector3f _24 = sead::Vector3f::zero;
+    sead::Vector3f mFrozenLinearVelocity = sead::Vector3f::zero;
+    sead::Vector3f mFrozenAngularVelocity = sead::Vector3f::zero;
     sead::Vector3f mLinearVelocity = sead::Vector3f::zero;
     float mMaxLinearVelocity{};
     sead::Vector3f mAngularVelocity = sead::Vector3f::zero;
@@ -72,8 +75,8 @@ private:
     RigidBody* mLinkedRigidBody{};
     sead::Matrix34f mTransform = sead::Matrix34f::ident;
     float mTimeFactor = 1.0f;
-    sead::Vector3f _9c = sead::Vector3f::zero;
-    sead::Quatf _a8 = sead::Quatf::unit;
+    sead::Vector3f mLinkedBodyExtraTranslate = sead::Vector3f::zero;
+    sead::Quatf mLinkedBodyExtraRotate = sead::Quatf::unit;
     sead::TypedBitFlag<Flag, sead::Atomic<u32>> mFlags{};
 };