Implement native wheel node matrix retrieval and enhance topology validation

runtime/include/vr/native_wheel_topology.h

@@ -2,8 +2,10 @@
 #pragma once
 
 #include <array>
+#include <cmath>
 #include <cstddef>
 #include <cstdint>
+#include <cstring>
 #include <numeric>
 #include <vector>
 
@@ -171,4 +173,46 @@ inline bool ReadNativeWheelTopology(const uint8_t *mdl, size_t size, uint32_t ar
     return found;
 }
 
+// The bone GX draws node `node` through: its 0x70 matrix is that bone's
+// model-space transform, parents included. Found by node id, not dictionary
+// position: the Flame Flyer and Cheep Charger bodies open with an nw4r_root
+// bone on node 2 or 3, and the node-0 bone their wheel belongs to comes third.
+inline bool ReadNativeWheelNodeMatrix(const uint8_t *mdl, size_t size, uint32_t node, float out[12]) {
+    if (!mdl || size < 0x40)
+        return false;
+    const auto read32 = [&](size_t at) {
+        return (uint32_t(mdl[at]) << 24) | (uint32_t(mdl[at + 1]) << 16) | (uint32_t(mdl[at + 2]) << 8) | mdl[at + 3];
+    };
+    const auto contains = [&](size_t at, size_t length) { return at <= size && length <= size - at; };
+    const auto version = read32(8);
+    if (read32(0) != 0x4d444c30 || version < 8 || version > 11 || read32(4) != size)
+        return false;
+    const size_t dictionary = read32(0x14);
+    if (!dictionary || !contains(dictionary, 8))
+        return false;
+    const auto count = read32(dictionary + 4);
+    if (!count || count > 4096 || !contains(dictionary + 8, size_t(count + 1) * 16))
+        return false;
+    for (uint32_t entry = 1; entry <= count; ++entry) {
+        const size_t offset = read32(dictionary + 8 + entry * 16 + 12);
+        if (offset > size - dictionary)
+            return false;
+        const size_t bone = dictionary + offset;
+        if (!contains(bone, 0xa0))
+            return false;
+        if (read32(bone + 0x10) != node)
+            continue;
+        for (unsigned i = 0; i < 12; ++i) {
+            const uint32_t bits = read32(bone + 0x70 + i * 4);
+            float value;
+            std::memcpy(&value, &bits, sizeof value);
+            if (!std::isfinite(value))
+                return false;
+            out[i] = value;
+        }
+        return true;
+    }
+    return false;
+}
+
 } // namespace mkw::vr

runtime/src/vr/mkw_vr_first_person.cpp

@@ -740,22 +740,18 @@ bool PublishNativeWheelMesh(uint32_t part, const Mtx34& model_view, const Mtx34&
         if (version < 8 || version > 11) {
             return false;
         }
+        const uint32_t mdl_size = Memory::Read32(mdl + 4);
+        const uint8_t* mdl_bytes = mdl_size <= 0x1000000 ? Memory::GetPointer(mdl, mdl_size) : nullptr;
         Mtx34 body_from_vertices = kIdentityMtx34;
         Mtx34 wheel_model_view = model_view;
         if (!whole_part) {
-            // Body::mtx is the model placement, while the root bone may have
-            // its own authored rotation (notably the Baby Booster). GX draws
-            // its positions through placement * root, so both wheel selection
-            // and Aurora's local-player matrix match must include that root.
-            const uint32_t bone_dic_offset = Memory::Read32(mdl + 0x14);
-            if (!bone_dic_offset || bone_dic_offset > 0x100000) return false;
-            const uint32_t bone_dic = mdl + bone_dic_offset;
-            if (!Memory::Contains(bone_dic, 40) || !Memory::Read32(bone_dic + 4)) return false;
-            const uint32_t bone_offset = Memory::Read32(bone_dic + 36);
-            if (bone_offset > 0x100000) return false;
-            const uint32_t root_bone = bone_dic + bone_offset;
-            if (!Memory::Contains(root_bone, 0xa0) || Memory::Read32(root_bone + 0x10) != 0 ||
-                !ReadGuestMtx34(root_bone + 0x70, body_from_vertices)) return false;
+            // Body::mtx is the model placement, while the bone GX draws the
+            // body's node 0 through may have its own authored transform (the
+            // Baby Booster's rotated root). GX draws its positions through
+            // placement * bone, so both wheel selection and Aurora's
+            // local-player matrix match must include it. It is found by node
+            // id: the Flame Flyer and Cheep Charger list an nw4r_root first.
+            if (!ReadNativeWheelNodeMatrix(mdl_bytes, mdl_size, 0, body_from_vertices.data())) return false;
             wheel_model_view = ComposeMtx(model_view, body_from_vertices);
         }
         const uint32_t dic_offset = Memory::Read32(mdl + 0x18);
@@ -814,8 +810,6 @@ bool PublishNativeWheelMesh(uint32_t part, const Mtx34& model_view, const Mtx34&
                     point = detail::TransformPoint(correction, point.x, point.y, point.z);
                 }
             } else {
-                const uint32_t mdl_size = Memory::Read32(mdl + 4);
-                const uint8_t* mdl_bytes = mdl_size <= 0x1000000 ? Memory::GetPointer(mdl, mdl_size) : nullptr;
                 NativeWheelTopology topology(num);
                 if (!ReadNativeWheelTopology(mdl_bytes, mdl_size, Memory::Read32(header + 0x10), topology) ||
                     RotateNativeWheelVertices(points, topology, center, radius, angle, &correction,

runtime/tests/vr_cockpit_tests.cpp

@@ -10,7 +10,9 @@
 #include "vr/mkw_vr_first_person.h"
 #include "vr/native_wheel_mesh.h"
 
+#include <algorithm>
 #include <cmath>
+#include <cstring>
 #include <iostream>
 #include <vector>
 
@@ -337,6 +339,45 @@ void TestNativeWheelTopology() {
     Check(!ReadNativeWheelTopology(mdl.data(), mdl.size(), 1, model), "unrelated position array ignored");
     put32(shape + 0x2c, UINT32_MAX);
     Check(!ReadNativeWheelTopology(mdl.data(), mdl.size(), 0, model), "escaping primitive offset rejected");
+    // The bone the wheel's positions are drawn through is found by node id:
+    // the Flame Flyer and Cheep Charger list an nw4r_root bone (node 2)
+    // before the node-0 bone, which used to leave them on the VR wheel.
+    std::vector<uint8_t> bones(0x1c0, 0);
+    const auto putBone32 = [&](size_t at, uint32_t value) {
+        for (unsigned i = 0; i < 4; ++i)
+            bones[at + i] = uint8_t(value >> ((3 - i) * 8));
+    };
+    const auto putBoneFloat = [&](size_t at, float value) {
+        uint32_t bits;
+        std::memcpy(&bits, &value, sizeof bits);
+        putBone32(at, bits);
+    };
+    putBone32(0, 0x4d444c30);
+    putBone32(4, uint32_t(bones.size()));
+    putBone32(8, 11);
+    putBone32(0x14, 0x40);
+    putBone32(0x44, 2);
+    putBone32(0x40 + 8 + 16 + 12, 0x40);
+    putBone32(0x40 + 8 + 32 + 12, 0xe0);
+    putBone32(0x80 + 0x10, 2);
+    putBone32(0x120 + 0x10, 0);
+    const float authored[12]{0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0};
+    for (unsigned i = 0; i < 12; ++i) {
+        putBoneFloat(0x80 + 0x70 + i * 4, i % 5 == 0 ? 1.0f : 0.0f);
+        putBoneFloat(0x120 + 0x70 + i * 4, authored[i]);
+    }
+    float bone[12]{};
+    Check(ReadNativeWheelNodeMatrix(bones.data(), bones.size(), 0, bone) &&
+              std::equal(std::begin(authored), std::end(authored), bone),
+          "node-0 bone found behind an nw4r_root bone");
+    Check(ReadNativeWheelNodeMatrix(bones.data(), bones.size(), 2, bone) && bone[0] == 1.0f && bone[2] == 0.0f,
+          "bones are told apart by node id");
+    Check(!ReadNativeWheelNodeMatrix(bones.data(), bones.size(), 5, bone), "missing node rejected");
+    Check(!ReadNativeWheelNodeMatrix(bones.data(), bones.size() - 1, 0, bone), "MDL0 size mismatch rejected");
+    putBoneFloat(0x120 + 0x70, NAN);
+    Check(!ReadNativeWheelNodeMatrix(bones.data(), bones.size(), 0, bone), "non-finite bone matrix rejected");
+    putBone32(0x40 + 8 + 32 + 12, uint32_t(bones.size()));
+    Check(!ReadNativeWheelNodeMatrix(bones.data(), bones.size(), 0, bone), "escaping bone offset rejected");
 }
 
 } // namespace