ksys/phys: Add RagdollController::getTransform and more

Credit to @Dragorn421 for helping me figure out the rotation matrix
thing in the getTransform function

5 files changed, 191 insertions, 23 deletions

diff --git a/src/KingSystem/Physics/Ragdoll/physRagdollController.cpp b/src/KingSystem/Physics/Ragdoll/physRagdollController.cpp
index 51b59047..93611e42 100644
--- a/src/KingSystem/Physics/Ragdoll/physRagdollController.cpp
+++ b/src/KingSystem/Physics/Ragdoll/physRagdollController.cpp
@@ -9,7 +9,10 @@
 #include <Havok/Physics2012/Dynamics/World/hkpPhysicsSystem.h>
 #include <Havok/Physics2012/Dynamics/World/hkpWorld.h>
 #include <Havok/Physics2012/Utilities/Dynamics/ScaleSystem/hkpSystemScalingUtility.h>
+#include <basis/seadRawPrint.h>
+#include <gsys/gsysModel.h>
 #include <math/seadMathCalcCommon.h>
+#include <memory>
 #include <resource/seadResource.h>
 #include "KingSystem/Physics/Ragdoll/physRagdollControllerMgr.h"
 #include "KingSystem/Physics/Ragdoll/physRagdollRigidBody.h"
@@ -22,6 +25,7 @@
 #include "KingSystem/Physics/physConversions.h"
 #include "KingSystem/Utils/Debug.h"
 #include "KingSystem/Utils/IteratorUtil.h"
+#include "KingSystem/Utils/MathUtil.h"
 #include "KingSystem/Utils/SafeDelete.h"
 
 namespace ksys::phys {
@@ -87,17 +91,7 @@ bool RagdollController::doInit(const RagdollParam* param, sead::DirectResource*
     }
 
     mBoneRigidBodies.allocBufferAssert(num_bones, heap);
-
-    {
-        static constexpr size_t Alignment = 0x20;
-        _b0 = sead::Mathu::roundUpPow2(sizeof(hkQsTransformf) * num_bones, Alignment);
-        auto* buffer = new (heap, Alignment) u8[_b0];
-        if (num_bones > 0 && buffer != nullptr) {
-            _a0 = num_bones;
-            _a8 = buffer;
-        }
-    }
-
+    allocateBoneTransforms(num_bones, heap);
     mBoneVectors.allocBufferAssert(num_bones, heap);
     mBoneStuff.allocBufferAssert(num_bones, heap);
 
@@ -158,6 +152,14 @@ bool RagdollController::doInit(const RagdollParam* param, sead::DirectResource*
     return true;
 }
 
+void RagdollController::allocateBoneTransforms(int num_bones, sead::Heap* heap) {
+    static constexpr size_t Alignment = 0x20;
+    mBoneTransformsByteSize =
+        sead::Mathu::roundUpPow2(sizeof(hkQsTransformf) * num_bones, Alignment);
+    auto* buffer = new (heap, Alignment) u8[mBoneTransformsByteSize];
+    mBoneTransforms.setBuffer(num_bones, new (buffer) hkQsTransformf[num_bones]);
+}
+
 void RagdollController::registerSelf() {
     const bool registered = System::instance()->getRagdollControllerMgr()->addController(this);
     mFlags.change(Flag::IsRegistered, registered);
@@ -208,9 +210,10 @@ void RagdollController::finalize() {
     mBoneStuff.freeBuffer();
     mBoneStuff2.freeBuffer();
 
-    if (_b0) {
-        delete[] _a8;
-        _b0 = 0;
+    if (mBoneTransformsByteSize != 0) {
+        std::destroy_n(mBoneTransforms.getBufferPtr(), mBoneTransforms.size());
+        delete[] reinterpret_cast<u8*>(mBoneTransforms.getBufferPtr());
+        mBoneTransformsByteSize = 0;
     }
 }
 
@@ -414,6 +417,64 @@ void RagdollController::setContactNone(int bone_index) {
     mBoneRigidBodies[bone_index]->setContactNone(false);
 }
 
+void RagdollController::changeWorldState(RagdollController::WorldState state) {
+    if (getWorldState() == state)
+        return;
+
+    const int num_rigid_bodies = mRagdollInstance->m_rigidBodies.size();
+
+    ScopedPhysicsLock lock{this};
+
+    if (state == WorldState::AddedToWorld) {
+        if (mModel) {
+            if (auto* accessor = getModelBoneAccessor())
+                accessor->copyModelPoseToHavok(ModelBoneAccessor::EnableScale::No);
+
+            setTransform(mModel->getMatrix());
+        }
+
+        for (int i = 0, n = num_rigid_bodies; i < n; ++i) {
+            mBoneRigidBodies[i]->addToWorld();
+
+            mBoneRigidBodies[i]->setLinearVelocity(sead::Vector3f::zero);
+            mBoneRigidBodies[i]->setAngularVelocity(sead::Vector3f::zero);
+
+            if (mKeyframedBones.isOnBit(i)) {
+                mBoneRigidBodies[i]->changeMotionType(MotionType::Keyframed);
+                mBoneRigidBodies[i]->setContactLayer(ContactLayer::EntityNoHit);
+            } else {
+                mBoneRigidBodies[i]->changeMotionType(MotionType::Dynamic);
+                mBoneRigidBodies[i]->setContactLayer(mContactLayer);
+            }
+        }
+
+        for (auto& x : mBoneStuff)
+            x = 1;
+
+        mFlags.reset(Flag::_8);
+        _e8 = _e9;
+
+        for (auto* body : mBoneRigidBodies)
+            body->setContactNone(true);
+
+        mFlags.set(Flag::_20);
+        mFlags.set(Flag::AddedToWorld);
+        mFlags.reset(Flag::_40);
+        mFlags.reset(Flag::_2);
+        mFlags.reset(Flag::_4);
+    } else {
+        for (int i = 0, n = num_rigid_bodies; i < n; ++i)
+            mBoneRigidBodies[i]->removeFromWorld();
+
+        mFlags.reset(Flag::AddedToWorld);
+        mFlags.reset(Flag::_40);
+    }
+}
+
+RagdollController::WorldState RagdollController::getWorldState() const {
+    return mFlags.isOn(Flag::AddedToWorld) ? WorldState::AddedToWorld : WorldState::NotAddedToWorld;
+}
+
 void RagdollController::setExtraRigidBody(RigidBody* body, int bone_index) {
     mExtraRigidBody = body;
     mBoneIndexForExtraRigidBody = bone_index;
@@ -468,6 +529,88 @@ RagdollRigidBody* RagdollController::getChildBoneRigidBody(const RigidBody* body
     return sead::DynamicCast<const RagdollRigidBody>(body)->getChildBodies_()[index];
 }
 
+sead::Matrix34f RagdollController::getTransform(int bone_index) const {
+    if (mFlags.isOff(Flag::_40)) {
+        if (util::isMatrixInvalid(mModel->getMatrix()))
+            return sead::Matrix34f::ident;
+
+        return mModel->getMatrix();
+    }
+
+    const sead::Matrix34f parent_transform = mBoneRigidBodies[bone_index]->getTransform();
+    hkQsTransformf hk_parent_transform;
+    toHkQsTransform(&hk_parent_transform, parent_transform, sead::Vector3f::ones);
+
+    hkQsTransformf result;
+    result.setMulMulInverse(hk_parent_transform, mBoneTransforms[bone_index]);
+
+    const bool parent_ok = hk_parent_transform.isOk();
+    SEAD_ASSERT(parent_ok);
+    const bool bone_ok = mBoneTransforms[bone_index].isOk();
+    SEAD_ASSERT(bone_ok);
+    const bool result_ok = result.isOk();
+    SEAD_ASSERT(result_ok);
+
+    sead::Matrix34f out;
+    toMtx34(&out, result);
+    return out;
+}
+
+sead::Matrix34f RagdollController::getTransformWithCustomYAxis(int bone_index,
+                                                               const sead::Vector3f& y_axis) const {
+    if (mFlags.isOff(Flag::_40)) {
+        if (util::isMatrixInvalid(mModel->getMatrix()))
+            return sead::Matrix34f::ident;
+
+        return mModel->getMatrix();
+    }
+
+    sead::Matrix34f transform = mBoneRigidBodies[bone_index]->getTransform();
+
+    {
+        sead::Vector3f v1 = (transform(1, 1) > 0.0f) ? y_axis : -y_axis;
+        v1.normalize();
+
+        sead::Vector3f c2 = util::getCol(transform, 2);
+        c2.normalize();
+
+        sead::Vector3f v0;
+        // Compute the cross product of v1 and c2.
+        // If v1 and c2 are collinear, then their cross product will be the zero vector;
+        // compute the cross product of v1 and the first column instead in that case.
+        // (That product should never be zero because the first column and u2 are not collinear.)
+        const auto dot = sead::Mathf::abs(v1.dot(c2));
+        // XXX: this should probably check if the dot product is approximately equal to 1.0 instead
+        v0.setCross(v1, dot >= 1.0f ? util::getCol(transform, 0) : c2);
+        v0.normalize();
+
+        sead::Vector3f v2;
+        v2.setCross(v0, v1);
+        v2.normalize();
+
+        transform.setBase(0, v0);
+        transform.setBase(1, v1);
+        transform.setBase(2, v2);
+    }
+
+    hkQsTransformf hk_parent_transform;
+    toHkQsTransform(&hk_parent_transform, transform, sead::Vector3f::ones);
+
+    hkQsTransformf result;
+    result.setMulMulInverse(hk_parent_transform, mBoneTransforms[bone_index]);
+
+    const bool parent_ok = hk_parent_transform.isOk();
+    SEAD_ASSERT(parent_ok);
+    const bool bone_ok = mBoneTransforms[bone_index].isOk();
+    SEAD_ASSERT(bone_ok);
+    const bool result_ok = result.isOk();
+    SEAD_ASSERT(result_ok);
+
+    sead::Matrix34f out;
+    toMtx34(&out, result);
+    return out;
+}
+
 int RagdollController::getConstraintIndexByName(const sead::SafeString& name) const {
     for (int i = 0, n = getNumConstraints(); i < n; ++i) {
         if (name == mRagdollInstance->m_constraints[i]->getName())
@@ -492,7 +635,7 @@ void RagdollController::setContactLayer(ContactLayer layer) {
     if (mContactLayer.value() == layer)
         return;
 
-    if (mFlags.isOn(Flag::_10)) {
+    if (mFlags.isOn(Flag::AddedToWorld)) {
         for (int bone = 0, num_bones = mBoneRigidBodies.size(); bone < num_bones; ++bone) {
             if (mKeyframedBones.isOffBit(bone))
                 mBoneRigidBodies[bone]->setContactLayer(layer);
@@ -528,7 +671,7 @@ void RagdollController::setMaximumUnk1(u8 value) {
 }
 
 void RagdollController::stopForcingKeyframing() {
-    if (mFlags.isOn(Flag::_10)) {
+    if (mFlags.isOn(Flag::AddedToWorld)) {
         for (int i = 0, n = mBoneRigidBodies.size(); i < n; ++i) {
             setKeyframed(i, false, {});
         }
diff --git a/src/KingSystem/Physics/Ragdoll/physRagdollController.h b/src/KingSystem/Physics/Ragdoll/physRagdollController.h
index c7524abd..7528e59e 100644
--- a/src/KingSystem/Physics/Ragdoll/physRagdollController.h
+++ b/src/KingSystem/Physics/Ragdoll/physRagdollController.h
@@ -90,7 +90,14 @@ public:
     void setContactNone();
     void setContactAll(int bone_index);
     void setContactNone(int bone_index);
-    u32 sub_7101221CC4();
+
+    enum class WorldState {
+        AddedToWorld = 0,
+        NotAddedToWorld = 2,
+    };
+    void changeWorldState(WorldState state);
+    WorldState getWorldState() const;
+
     void setExtraRigidBody(RigidBody* body, int bone_index);
     void setGravityFactor(float factor);
 
@@ -103,6 +110,9 @@ public:
     int getNumChildBones(const RigidBody* body) const;
     RagdollRigidBody* getChildBoneRigidBody(const RigidBody* body, int index) const;
 
+    sead::Matrix34f getTransform(int bone_index) const;
+    sead::Matrix34f getTransformWithCustomYAxis(int bone_index, const sead::Vector3f& y_axis) const;
+
     int getConstraintIndexByName(const sead::SafeString& name) const;
     int getNumConstraints() const;
     void enableConstraint(int index, bool enable);
@@ -133,7 +143,12 @@ private:
     };
 
     enum class Flag {
-        _10 = 0x10,
+        _2 = 0x2,
+        _4 = 0x4,
+        _8 = 0x8,
+        AddedToWorld = 0x10,
+        _20 = 0x20,
+        _40 = 0x40,
         _80 = 0x80,
         /// Whether this controller has been registered with the RagdollControllerMgr.
         IsRegistered = 0x100,
@@ -149,6 +164,7 @@ private:
 
     bool doInit(const RagdollParam* param, sead::DirectResource* res, gsys::Model* model,
                 sead::Heap* heap);
+    void allocateBoneTransforms(int num_bones, sead::Heap* heap);
     void finalize();
     void removeConstraints();
     void setTransform(const hkQsTransformf& transform);
@@ -179,10 +195,8 @@ private:
     sead::Buffer<float> mBoneStuff2;
     float _98 = 0.1;
     float mGravityFactorOverride = 1.0;
-    u32 _a0 = 0;
-    // TODO: type
-    u8* _a8 = nullptr;
-    u32 _b0 = 0;
+    sead::Buffer<hkQsTransformf> mBoneTransforms;
+    u32 mBoneTransformsByteSize = 0;
     const RagdollParam* mRagdollParam = nullptr;
     sead::TypedBitFlag<Flag> mFlags;
     sead::BitFlag32 mDisabledConstraints;
diff --git a/src/KingSystem/Physics/System/physInstanceSet.cpp b/src/KingSystem/Physics/System/physInstanceSet.cpp
index 872cd713..6e5269c8 100644
--- a/src/KingSystem/Physics/System/physInstanceSet.cpp
+++ b/src/KingSystem/Physics/System/physInstanceSet.cpp
@@ -192,7 +192,7 @@ void InstanceSet::sub_7100FBD284(const sead::Matrix34f& mtx) {
     if (mRagdollController == nullptr)
         return;
 
-    if (mRagdollController->sub_7101221CC4() == 0)
+    if (mRagdollController->getWorldState() == RagdollController::WorldState::AddedToWorld)
         sub_7100FBC890(mtx, false, false);
 }
 
diff --git a/src/KingSystem/Physics/physConversions.h b/src/KingSystem/Physics/physConversions.h
index 2ee002cc..e58bf8dd 100644
--- a/src/KingSystem/Physics/physConversions.h
+++ b/src/KingSystem/Physics/physConversions.h
@@ -56,6 +56,11 @@ inline void toHkQuat(hkQuaternionf* out, const sead::Quatf& quat) {
     return {quat.x, quat.y, quat.z, quat.w};
 }
 
+inline void toMtx34(sead::Matrix34f* out, const hkQsTransformf& transform) {
+    out->makeSQT(toVec3(transform.getScale()), toQuat(transform.getRotation()),
+                 toVec3(transform.getTranslation()));
+}
+
 inline void toMtx34(sead::Matrix34f* out, const hkTransformf& transform) {
     const hkRotationf& rotate = transform.getRotation();
     const hkVector4f& translate = transform.getTranslation();
diff --git a/src/KingSystem/Utils/MathUtil.h b/src/KingSystem/Utils/MathUtil.h
index 6ba5d833..9d4b0ed3 100644
--- a/src/KingSystem/Utils/MathUtil.h
+++ b/src/KingSystem/Utils/MathUtil.h
@@ -56,4 +56,10 @@ inline bool isMatrixInvalid(const sead::Matrix34f& matrix) {
     return false;
 }
 
+inline sead::Vector3f getCol(const sead::Matrix34f& mtx, int col) {
+    sead::Vector3f result;
+    mtx.getBase(result, col);
+    return result;
+}
+
 }  // namespace ksys::util