ksys/phys: Finish ModelBoneAccessor

5 files changed, 196 insertions, 32 deletions

diff --git a/src/KingSystem/Physics/Rig/physModelBoneAccessor.cpp b/src/KingSystem/Physics/Rig/physModelBoneAccessor.cpp
index 5165f5ee..1daaeda4 100644
--- a/src/KingSystem/Physics/Rig/physModelBoneAccessor.cpp
+++ b/src/KingSystem/Physics/Rig/physModelBoneAccessor.cpp
@@ -1,4 +1,5 @@
 #include "KingSystem/Physics/Rig/physModelBoneAccessor.h"
+#include <Havok/Animation/Animation/Rig/hkaPose.h>
 #include <Havok/Animation/Animation/Rig/hkaSkeleton.h>
 #include <Havok/Common/Base/Memory/Allocator/Malloc/hkMallocAllocator.h>
 #include <gsys/gsysModel.h>
@@ -8,6 +9,8 @@
 #include <nn/g3d/ResSkeleton.h>
 #include <nn/g3d/SkeletonObj.h>
 #include <type_traits>
+#include "KingSystem/Physics/physConversions.h"
+#include "KingSystem/Utils/MathUtil.h"
 #include "KingSystem/Utils/SafeDelete.h"
 #include "KingSystem/Utils/Types.h"
 
@@ -160,14 +163,16 @@ bool ModelSkeleton::constructFromModel(ModelBoneAccessor::ModelBoneFilter* bone_
 
 }  // namespace detail
 
+static int sModelBoneAccessorUnkMode;
+static bool sModelBoneAccessorUnkFlag;
+
 ModelBoneAccessor::ModelBoneAccessor() = default;
 
 ModelBoneAccessor::~ModelBoneAccessor() {
     ModelBoneAccessor::finalize();
 }
 
-bool ModelBoneAccessor::init(const hkaSkeleton* skeleton, const gsys::Model* model,
-                             sead::Heap* heap) {
+bool ModelBoneAccessor::init(const hkaSkeleton* skeleton, gsys::Model* model, sead::Heap* heap) {
     if (!skeleton || !model)
         return false;
 
@@ -194,7 +199,7 @@ bool ModelBoneAccessor::init(const hkaSkeleton* skeleton, const gsys::Model* mod
     return true;
 }
 
-bool ModelBoneAccessor::init(const gsys::Model* model, int model_unit_index, sead::Heap* heap,
+bool ModelBoneAccessor::init(gsys::Model* model, int model_unit_index, sead::Heap* heap,
                              ModelBoneAccessor::ModelBoneFilter* bone_filter) {
     if (!model)
         return false;
@@ -228,4 +233,135 @@ void ModelBoneAccessor::finalize() {
         util::safeDelete(mModelSkeleton);
 }
 
+int ModelBoneAccessor::findBoneIndex(const gsys::BoneAccessKey& key) const {
+    for (int i = 0, n = mBoneAccessKeys.size(); i < n; ++i) {
+        if (key == mBoneAccessKeys[i].key.getKey())
+            return i;
+    }
+    return -1;
+}
+
+const char* ModelBoneAccessor::getBoneName(int index) const {
+    if (index < 0 || index >= getSkeleton()->m_bones.getSize())
+        return nullptr;
+    return getSkeleton()->m_bones[index].m_name.cString();
+}
+
+// NON_MATCHING: two loop induction variables (bone_idx) swapped
+void ModelBoneAccessor::copyModelPoseToHavok(EnableScale enable_scale) const {
+    auto& hk_bones = getPose()->accessUnsyncedPoseLocalSpace();
+
+    switch (getUnkMode()) {
+    case 2:
+        enable_scale = EnableScale::Yes;
+        break;
+    case 1:
+        enable_scale = EnableScale::No;
+        break;
+    }
+
+    for (int bone_idx = 0, n = mBoneAccessKeys.getSize(); bone_idx < n; ++bone_idx) {
+        if (!mBoneAccessKeys[bone_idx].key.isValid())
+            continue;
+
+        if (!mBoneAccessKeys[bone_idx]._38)
+            continue;
+
+        const auto key = mBoneAccessKeys[bone_idx].key.getKey();
+        auto* unit = getModelUnit(bone_idx);
+
+        sead::Vector3f scale;
+        sead::Matrix34f transform;
+        unit->getBoneLocalMatrix(&transform, &scale, key.bone_index);
+
+        if (util::isMatrixInvalid(transform))
+            continue;
+
+        if ((!bool(enable_scale) || getUnkFlag()) &&
+            !scale.equals(sead::Vector3f::ones, sead::Mathf::epsilon())) {
+            if (getUnkFlag()) {
+                // leftover debug code
+                static_cast<void>(getModelUnit(bone_idx)->getName().include("Link"));
+            }
+
+            if (!bool(enable_scale))
+                scale = {1, 1, 1};
+        }
+
+        toHkQsTransform(&hk_bones[bone_idx], transform,
+                        bone_idx == 0 ? sead::Vector3f::ones : scale);
+
+        if (mModel->getScale().x != 1)
+            hk_bones[bone_idx].m_translation.mul(mModel->getScale().x);
+    }
+}
+
+void ModelBoneAccessor::copyHavokPoseToModel(EnableScale enable_scale) const {
+    const auto& hk_bones = getPoseInternal()->getSyncedPoseLocalSpace();
+
+    bool not_first_bone = false;
+
+    for (int bone_idx = 0, n = mBoneAccessKeys.getSize(); bone_idx < n; ++bone_idx) {
+        if (!mBoneAccessKeys[bone_idx]._39)
+            continue;
+
+        if (!mBoneAccessKeys[bone_idx].key.isValid())
+            continue;
+
+        sead::Vector3f translate;
+        translate = sead::Vector3f::zero;
+        if (!not_first_bone)
+            translate += mTranslate;
+
+        hkTransform hk_transform;
+        const_cast<hkQsTransformf&>(hk_bones[bone_idx]).fastRenormalize();
+        hk_bones[bone_idx].copyToTransform(hk_transform);
+
+        sead::Vector3f scale = toVec3(hk_bones[bone_idx].getScale());
+        sead::Matrix34f transform;
+        for (int i = 0; i < 3; ++i) {
+            for (int j = 0; j < 3; ++j) {
+                transform(i, j) = hk_transform.getRotation().getColumn(j)(i);
+            }
+            transform(i, 3) = translate.e[i] + hk_transform.getTranslation()(i);
+        }
+
+        const float model_scale = mModel->getScale().x;
+        if (model_scale != 1 && model_scale > 0) {
+            for (int i = 0; i < 3; ++i)
+                transform(i, 3) *= 1 / model_scale;
+        }
+
+        if (!util::isMatrixInvalid(transform)) {
+            if ((!bool(enable_scale) || getUnkFlag()) &&
+                !scale.equals(sead::Vector3f::ones, sead::Mathf::epsilon())) {
+                if (getUnkFlag()) {
+                    // leftover debug code
+                    static_cast<void>(getModelUnit(bone_idx)->getName().include("Link"));
+                }
+
+                if (!bool(enable_scale))
+                    scale = {1, 1, 1};
+            }
+
+            mModel->setBoneLocalMatrix(mBoneAccessKeys[bone_idx].key.getKey(), transform, scale);
+        }
+
+        not_first_bone = true;
+    }
+}
+
+int& ModelBoneAccessor::getUnkMode() {
+    return sModelBoneAccessorUnkMode;
+}
+
+bool& ModelBoneAccessor::getUnkFlag() {
+    return sModelBoneAccessorUnkFlag;
+}
+
+gsys::ModelUnit* ModelBoneAccessor::getModelUnit(int bone_idx) const {
+    const auto key = mBoneAccessKeys[bone_idx].key.getKey();
+    return mModel->getUnits()[key.model_unit_index]->mModelUnit;
+}
+
 }  // namespace ksys::phys
diff --git a/src/KingSystem/Physics/Rig/physModelBoneAccessor.h b/src/KingSystem/Physics/Rig/physModelBoneAccessor.h
index 1480bf3d..602983f0 100644
--- a/src/KingSystem/Physics/Rig/physModelBoneAccessor.h
+++ b/src/KingSystem/Physics/Rig/physModelBoneAccessor.h
@@ -19,6 +19,8 @@ class ModelSkeleton;
 
 class ModelBoneAccessor : public BoneAccessor {
 public:
+    enum class EnableScale : bool { Yes = true, No = false };
+
     class ModelBoneFilter {
     public:
         using BoneBitSet = sead::LongBitFlag<1024>;
@@ -30,13 +32,22 @@ public:
     ModelBoneAccessor();
     ~ModelBoneAccessor() override;
 
-    bool init(const hkaSkeleton* skeleton, const gsys::Model* model, sead::Heap* heap);
+    bool init(const hkaSkeleton* skeleton, gsys::Model* model, sead::Heap* heap);
 
-    bool init(const gsys::Model* model, int model_unit_index, sead::Heap* heap,
+    bool init(gsys::Model* model, int model_unit_index, sead::Heap* heap,
               ModelBoneFilter* bone_filter);
 
     void finalize() override;
 
+    int findBoneIndex(const gsys::BoneAccessKey& key) const;
+    const char* getBoneName(int index) const;
+
+    void copyModelPoseToHavok(EnableScale enable_scale) const;
+    void copyHavokPoseToModel(EnableScale enable_scale) const;
+
+    static int& getUnkMode();
+    static bool& getUnkFlag();
+
 protected:
     struct BoneAccessKey {
         gsys::BoneAccessKeyEx key;
@@ -44,9 +55,11 @@ protected:
         bool _39;
     };
 
-    const gsys::Model* mModel{};
+    gsys::ModelUnit* getModelUnit(int bone_idx) const;
+
+    gsys::Model* mModel{};
     sead::Buffer<BoneAccessKey> mBoneAccessKeys;
-    sead::Vector3f _38 = sead::Vector3f::zero;
+    sead::Vector3f mTranslate = sead::Vector3f::zero;
     detail::ModelSkeleton* mModelSkeleton{};
 };
 
diff --git a/src/KingSystem/Physics/RigidBody/physRigidBody.cpp b/src/KingSystem/Physics/RigidBody/physRigidBody.cpp
index f0710bca..d11aabfc 100644
--- a/src/KingSystem/Physics/RigidBody/physRigidBody.cpp
+++ b/src/KingSystem/Physics/RigidBody/physRigidBody.cpp
@@ -23,27 +23,12 @@
 #include "KingSystem/Physics/System/physSystem.h"
 #include "KingSystem/Physics/System/physUserTag.h"
 #include "KingSystem/Physics/physConversions.h"
+#include "KingSystem/Utils/MathUtil.h"
 
 namespace ksys::phys {
 
 constexpr float MinInertia = 0.001;
 
-static bool isVectorInvalid(const sead::Vector3f& vec) {
-    for (int i = 0; i < 3; ++i) {
-        if (std::isnan(vec.e[i]))
-            return true;
-    }
-    return false;
-}
-
-static bool isMatrixInvalid(const sead::Matrix34f& matrix) {
-    for (float x : matrix.a) {
-        if (std::isnan(x))
-            return true;
-    }
-    return false;
-}
-
 RigidBody::RigidBody(Type type, ContactLayerType layer_type, hkpRigidBody* hk_body,
                      const sead::SafeString& name, sead::Heap* heap, bool set_flag_10)
     : mCS(heap), mHkBody(hk_body), mRigidBodyAccessor(hk_body), mType(type) {
@@ -851,7 +836,7 @@ void RigidBody::setColor(const sead::Color4f& color, const void* a, bool b) {
 }
 
 void RigidBody::setPosition(const sead::Vector3f& position, bool propagate_to_linked_motions) {
-    if (isVectorInvalid(position)) {
+    if (util::isVectorInvalid(position)) {
         onInvalidParameter();
         return;
     }
@@ -904,7 +889,7 @@ sead::Matrix34f RigidBody::getTransform() const {
 }
 
 void RigidBody::setTransform(const sead::Matrix34f& mtx, bool propagate_to_linked_motions) {
-    if (isMatrixInvalid(mtx)) {
+    if (util::isMatrixInvalid(mtx)) {
         onInvalidParameter();
         return;
     }
@@ -1051,7 +1036,7 @@ void RigidBody::triggerScheduledMotionTypeChange() {
 }
 
 bool RigidBody::setLinearVelocity(const sead::Vector3f& velocity, float epsilon) {
-    if (isVectorInvalid(velocity)) {
+    if (util::isVectorInvalid(velocity)) {
         onInvalidParameter();
         return false;
     }
@@ -1078,7 +1063,7 @@ sead::Vector3f RigidBody::getLinearVelocity() const {
 }
 
 bool RigidBody::setAngularVelocity(const sead::Vector3f& velocity, float epsilon) {
-    if (isVectorInvalid(velocity)) {
+    if (util::isVectorInvalid(velocity)) {
         onInvalidParameter();
         return false;
     }
@@ -1347,7 +1332,7 @@ void RigidBody::applyLinearImpulse(const sead::Vector3f& impulse) {
     if (hasFlag(Flag::_400) || hasFlag(Flag::_40))
         return;
 
-    if (isVectorInvalid(impulse)) {
+    if (util::isVectorInvalid(impulse)) {
         onInvalidParameter();
         return;
     }
@@ -1363,7 +1348,7 @@ void RigidBody::applyAngularImpulse(const sead::Vector3f& impulse) {
     if (hasFlag(Flag::_400) || hasFlag(Flag::_40))
         return;
 
-    if (isVectorInvalid(impulse)) {
+    if (util::isVectorInvalid(impulse)) {
         onInvalidParameter();
         return;
     }
@@ -1379,12 +1364,12 @@ void RigidBody::applyPointImpulse(const sead::Vector3f& impulse, const sead::Vec
     if (hasFlag(Flag::_400) || hasFlag(Flag::_40))
         return;
 
-    if (isVectorInvalid(impulse)) {
+    if (util::isVectorInvalid(impulse)) {
         onInvalidParameter();
         return;
     }
 
-    if (isVectorInvalid(point)) {
+    if (util::isVectorInvalid(point)) {
         onInvalidParameter();
         return;
     }
diff --git a/src/KingSystem/Physics/physConversions.h b/src/KingSystem/Physics/physConversions.h
index 7647a3f5..2ee002cc 100644
--- a/src/KingSystem/Physics/physConversions.h
+++ b/src/KingSystem/Physics/physConversions.h
@@ -74,7 +74,7 @@ inline void toMtx34(sead::Matrix34f* out, const hkTransformf& transform) {
     mtx[2].store<4>(out->m[2]);
 }
 
-inline void toHkTransform(hkTransformf* out, const sead::Matrix34f& mtx) {
+inline void toHkTransform(hkTransform* out, const sead::Matrix34f& mtx) {
     sead::Quatf rotate;
     mtx.toQuat(rotate);
     rotate.normalize();
@@ -85,6 +85,18 @@ inline void toHkTransform(hkTransformf* out, const sead::Matrix34f& mtx) {
     out->set(toHkQuat(rotate), toHkVec4(translate));
 }
 
+inline void toHkQsTransform(hkQsTransform* out, const sead::Matrix34f& mtx,
+                            const sead::Vector3f& scale) {
+    sead::Quatf rotate;
+    mtx.toQuat(rotate);
+    rotate.normalize();
+
+    sead::Vector3f translate;
+    mtx.getTranslation(translate);
+
+    out->set(toHkVec4(translate), toHkQuat(rotate), toHkVec4(scale));
+}
+
 // Consider using toMtx34 if you have an hkTransform and wish to set both rotation and translation.
 inline void setMtxRotation(sead::Matrix34f* mtx, const hkRotationf& rotation) {
     mtx->setBase(0, toVec3(rotation.getColumn(0)));
diff --git a/src/KingSystem/Utils/MathUtil.h b/src/KingSystem/Utils/MathUtil.h
index e2db881f..6ba5d833 100644
--- a/src/KingSystem/Utils/MathUtil.h
+++ b/src/KingSystem/Utils/MathUtil.h
@@ -1,6 +1,8 @@
 #pragma once
 
+#include <cmath>
 #include <math/seadMathCalcCommon.h>
+#include <math/seadMatrix.h>
 #include <math/seadVector.h>
 
 namespace ksys::util {
@@ -38,4 +40,20 @@ inline sead::Vector3f lerp(const sead::Vector3f& a, const sead::Vector3f& b, flo
     return result;
 }
 
+inline bool isVectorInvalid(const sead::Vector3f& vec) {
+    for (int i = 0; i < 3; ++i) {
+        if (std::isnan(vec.e[i]))
+            return true;
+    }
+    return false;
+}
+
+inline bool isMatrixInvalid(const sead::Matrix34f& matrix) {
+    for (float x : matrix.a) {
+        if (std::isnan(x))
+            return true;
+    }
+    return false;
+}
+
 }  // namespace ksys::util