ksys/phys: Add ContactLayerCollisionInfoGroup

8 files changed, 258 insertions, 15 deletions

diff --git a/src/KingSystem/Physics/CMakeLists.txt b/src/KingSystem/Physics/CMakeLists.txt
index 010eeae4..e6cbf839 100644
--- a/src/KingSystem/Physics/CMakeLists.txt
+++ b/src/KingSystem/Physics/CMakeLists.txt
@@ -115,6 +115,8 @@ target_sources(uking PRIVATE
   System/physContactInfoParam.h
   System/physContactLayerCollisionInfo.cpp
   System/physContactLayerCollisionInfo.h
+  System/physContactLayerCollisionInfoGroup.cpp
+  System/physContactLayerCollisionInfoGroup.h
   System/physContactListener.cpp
   System/physContactListener.h
   System/physContactMgr.cpp
diff --git a/src/KingSystem/Physics/RigidBody/physRigidBody.cpp b/src/KingSystem/Physics/RigidBody/physRigidBody.cpp
index c76bb15c..46d6a9d2 100644
--- a/src/KingSystem/Physics/RigidBody/physRigidBody.cpp
+++ b/src/KingSystem/Physics/RigidBody/physRigidBody.cpp
@@ -477,7 +477,7 @@ void RigidBody::x_10() {
 void RigidBody::setCollisionInfo(CollisionInfo* info) {
     if (mCollisionInfo != info) {
         if (mCollisionInfo && isFlag8Set())
-            System::instance()->registerRigidBodyForContactSystem(this);
+            System::instance()->removeRigidBodyFromContactSystem(this);
         mCollisionInfo = info;
     }
 
@@ -717,7 +717,7 @@ void RigidBody::setCollisionFilterInfo(u32 info) {
     if (getCollisionFilterInfo() != info) {
         if (isFlag8Set()) {
             if (int(current_layer) != getContactLayer(EntityCollisionMask(info)))
-                System::instance()->registerRigidBodyForContactSystem(this);
+                System::instance()->removeRigidBodyFromContactSystem(this);
         }
 
         mHkBody->setCollisionFilterInfo(info);
diff --git a/src/KingSystem/Physics/System/physCollisionInfo.h b/src/KingSystem/Physics/System/physCollisionInfo.h
index eec2263f..5f6e0694 100644
--- a/src/KingSystem/Physics/System/physCollisionInfo.h
+++ b/src/KingSystem/Physics/System/physCollisionInfo.h
@@ -13,8 +13,7 @@ namespace ksys::phys {
 class RigidBody;
 
 struct CollidingBodies {
-    RigidBody* body_a;
-    RigidBody* body_b;
+    RigidBody* bodies[2];
     sead::ListNode list_node;
 
     static constexpr size_t getListNodeOffset() { return offsetof(CollidingBodies, list_node); }
diff --git a/src/KingSystem/Physics/System/physContactLayerCollisionInfo.h b/src/KingSystem/Physics/System/physContactLayerCollisionInfo.h
index 798dd5f5..068b6fab 100644
--- a/src/KingSystem/Physics/System/physContactLayerCollisionInfo.h
+++ b/src/KingSystem/Physics/System/physContactLayerCollisionInfo.h
@@ -6,7 +6,7 @@
 
 namespace ksys::phys {
 
-/// Tracks contact points for a contact layer.
+/// Tracks colliding bodies that belong to a specific contact layer.
 class ContactLayerCollisionInfo : public CollisionInfoBase {
 public:
     explicit ContactLayerCollisionInfo(ContactLayer layer);
diff --git a/src/KingSystem/Physics/System/physContactLayerCollisionInfoGroup.cpp b/src/KingSystem/Physics/System/physContactLayerCollisionInfoGroup.cpp
new file mode 100644
index 00000000..26cfbbe6
--- /dev/null
+++ b/src/KingSystem/Physics/System/physContactLayerCollisionInfoGroup.cpp
@@ -0,0 +1,117 @@
+#include "KingSystem/Physics/System/physContactLayerCollisionInfoGroup.h"
+#include "KingSystem/Physics/System/physContactLayerCollisionInfo.h"
+#include "KingSystem/Physics/System/physSystem.h"
+
+namespace ksys::phys {
+
+ContactLayerCollisionInfoGroup::ContactLayerCollisionInfoGroup(ContactLayer layer,
+                                                               const sead::SafeString& name)
+    : sead::INamable(name), mLayer(layer) {}
+
+ContactLayerCollisionInfoGroup::~ContactLayerCollisionInfoGroup() = default;
+
+ContactLayerCollisionInfoGroup* ContactLayerCollisionInfoGroup::make(sead::Heap* heap,
+                                                                     ContactLayer layer,
+                                                                     int capacity,
+                                                                     const sead::SafeString& name) {
+    return System::instance()->makeContactLayerCollisionInfoGroup(heap, layer, capacity, name);
+}
+
+void ContactLayerCollisionInfoGroup::free(ContactLayerCollisionInfoGroup* group) {
+    System::instance()->freeContactLayerCollisionInfoGroup(group);
+}
+
+void ContactLayerCollisionInfoGroup::init(sead::Heap* heap, int capacity) {
+    mCollisionInfoInstances.allocBuffer(capacity, heap);
+    mLayers.allocBufferAssert(capacity, heap);
+}
+
+void ContactLayerCollisionInfoGroup::finalize() {
+    mCollisionInfoInstances.freeBuffer();
+    mLayers.freeBuffer();
+}
+
+// NON_MATCHING: trivial reordering
+void ContactLayerCollisionInfoGroup::addLayer(ContactLayer layer) {
+    auto& info = mLayers[mCollisionInfoInstances.size()];
+    info.layer = layer;
+    info.layer_gt = mLayer > layer;
+    info.layer_le = !info.layer_gt;
+
+    auto* collision_info = System::instance()->trackLayerPair(mLayer, layer);
+    mCollisionInfoInstances.pushBack(collision_info);
+}
+
+void ContactLayerCollisionInfoGroup::ensureLayersAreTracked() {
+    for (int i = 0; i < mCollisionInfoInstances.size(); ++i) {
+        System::instance()->trackLayerPair(mLayer, mLayers[i].layer);
+    }
+}
+
+ContactLayerCollisionInfoGroup::CollidingBodiesIterator::CollidingBodiesIterator(
+    const ContactLayerCollisionInfoGroup* group, int index, IsStart start)
+    : mGroup(group), mInfoIndex(index) {
+    if (!bool(start))
+        return;
+
+    initIterator(group);
+
+    // If there is no colliding body, turn this iterator into an end iterator.
+    if (!mCollidingBodiesEntry) {
+        mInfo = nullptr;
+        mInfoIndex = group->mCollisionInfoInstances.size();
+    }
+}
+
+ContactLayerCollisionInfoGroup::CollidingBodiesIterator::~CollidingBodiesIterator() {
+    if (mInfo)
+        mInfo->unlock();
+}
+
+void ContactLayerCollisionInfoGroup::CollidingBodiesIterator::initIterator(
+    const ContactLayerCollisionInfoGroup* group) {
+    for (mInfoIndex = 0; mInfoIndex < group->mCollisionInfoInstances.size(); ++mInfoIndex) {
+        mInfo = group->mCollisionInfoInstances[mInfoIndex];
+        if (!mInfo)
+            continue;
+
+        mInfo->lock();
+        mCollidingBodiesEntry = mInfo->getCollidingBodies().front();
+        if (mCollidingBodiesEntry) {
+            // Keep the ContactLayerCollisionInfo locked.
+            break;
+        }
+        // Otherwise, unlock the current ContactLayerCollisionInfo and try the next one.
+        mInfo->unlock();
+    }
+}
+
+ContactLayerCollisionInfoGroup::CollidingBodiesIterator&
+ContactLayerCollisionInfoGroup::CollidingBodiesIterator::operator++() {
+    auto* group = mGroup;
+
+    mCollidingBodiesEntry = mInfo->getCollidingBodies().next(mCollidingBodiesEntry);
+    if (mCollidingBodiesEntry)
+        return *this;
+
+    // If we reached the last entry in the current ContactLayerCollisionInfo,
+    // move on to the next one.
+
+    ++mInfoIndex;
+    for (; mInfoIndex < group->mCollisionInfoInstances.size(); ++mInfoIndex) {
+        auto* next_info = group->mCollisionInfoInstances[mInfoIndex];
+        if (!next_info)
+            continue;
+
+        mInfo->unlock();
+        mInfo = next_info;
+        mInfo->lock();
+        mCollidingBodiesEntry = mInfo->getCollidingBodies().front();
+        if (mCollidingBodiesEntry)
+            break;
+    }
+
+    return *this;
+}
+
+}  // namespace ksys::phys
diff --git a/src/KingSystem/Physics/System/physContactLayerCollisionInfoGroup.h b/src/KingSystem/Physics/System/physContactLayerCollisionInfoGroup.h
new file mode 100644
index 00000000..59cff64e
--- /dev/null
+++ b/src/KingSystem/Physics/System/physContactLayerCollisionInfoGroup.h
@@ -0,0 +1,114 @@
+#pragma once
+
+#include <container/seadBuffer.h>
+#include <container/seadListImpl.h>
+#include <container/seadPtrArray.h>
+#include <prim/seadNamable.h>
+#include "KingSystem/Physics/physDefines.h"
+
+namespace ksys::phys {
+
+struct CollidingBodies;
+class ContactLayerCollisionInfo;
+
+/// Container for ContactLayerCollisionInfo instances that pertain to a contact layer
+/// paired with other layers.
+class ContactLayerCollisionInfoGroup : public sead::INamable {
+public:
+    class CollidingBodiesIterator;
+    struct CollidingBodiesRange;
+
+    struct LayerInfo {
+        ContactLayer layer;
+        /// Is `layer` greater than the layer of the group?
+        bool layer_gt;
+        /// Is `layer` lower or equal to the layer of the group?
+        bool layer_le;
+    };
+
+    static ContactLayerCollisionInfoGroup* make(sead::Heap* heap, ContactLayer layer, int capacity,
+                                                const sead::SafeString& name);
+    static void free(ContactLayerCollisionInfoGroup* group);
+
+    ContactLayerCollisionInfoGroup(ContactLayer layer, const sead::SafeString& name);
+    virtual ~ContactLayerCollisionInfoGroup();
+
+    /// @param capacity The maximum number of layers that can be added to this group.
+    void init(sead::Heap* heap, int capacity);
+    void finalize();
+
+    /// Add (mLayer, layer) to the list of layer pairs in this group.
+    void addLayer(ContactLayer layer);
+
+    /// Call this to ensure that all layer pairs in this group are being tracked by
+    /// the relevant contact listener. This may be necessary if the listener has been reset.
+    void ensureLayersAreTracked();
+
+    CollidingBodiesIterator collidingBodiesBegin() const;
+    CollidingBodiesIterator collidingBodiesEnd() const;
+    CollidingBodiesRange getCollidingBodies() const;
+
+    static constexpr size_t getListNodeOffset() {
+        return offsetof(ContactLayerCollisionInfoGroup, mListNode);
+    }
+
+private:
+    ContactLayer mLayer;
+    sead::PtrArray<ContactLayerCollisionInfo> mCollisionInfoInstances;
+    sead::Buffer<LayerInfo> mLayers;
+    sead::ListNode mListNode;
+};
+
+class ContactLayerCollisionInfoGroup::CollidingBodiesIterator {
+public:
+    enum class IsStart : bool { Yes = true, No = false };
+
+    CollidingBodiesIterator(const ContactLayerCollisionInfoGroup* group, int index, IsStart start);
+    ~CollidingBodiesIterator();
+
+    const LayerInfo& getLayerInfo() const { return mGroup->mLayers[mInfoIndex]; }
+
+    const CollidingBodies& operator*() const { return *mCollidingBodiesEntry; }
+    const CollidingBodies* operator->() const { return mCollidingBodiesEntry; }
+    CollidingBodiesIterator& operator++();
+
+    bool operator==(const CollidingBodiesIterator& other) const {
+        return mInfoIndex == other.mInfoIndex &&
+               mCollidingBodiesEntry == other.mCollidingBodiesEntry;
+    }
+
+    bool operator!=(const CollidingBodiesIterator& other) const { return !operator==(other); }
+
+private:
+    void initIterator(const ContactLayerCollisionInfoGroup* group);
+
+    const ContactLayerCollisionInfoGroup* mGroup{};
+    ContactLayerCollisionInfo* mInfo{};
+    /// The current CollidingBodies pair within the current ContactLayerCollisionInfo (mInfo).
+    const CollidingBodies* mCollidingBodiesEntry{};
+    int mInfoIndex{};
+};
+
+struct ContactLayerCollisionInfoGroup::CollidingBodiesRange {
+    auto begin() const { return group->collidingBodiesBegin(); }
+    auto end() const { return group->collidingBodiesEnd(); }
+
+    const ContactLayerCollisionInfoGroup* group;
+};
+
+inline ContactLayerCollisionInfoGroup::CollidingBodiesIterator
+ContactLayerCollisionInfoGroup::collidingBodiesBegin() const {
+    return {this, 0, CollidingBodiesIterator::IsStart::Yes};
+}
+
+inline ContactLayerCollisionInfoGroup::CollidingBodiesIterator
+ContactLayerCollisionInfoGroup::collidingBodiesEnd() const {
+    return {this, mCollisionInfoInstances.size(), CollidingBodiesIterator::IsStart::No};
+}
+
+inline ContactLayerCollisionInfoGroup::CollidingBodiesRange
+ContactLayerCollisionInfoGroup::getCollidingBodies() const {
+    return {this};
+}
+
+}  // namespace ksys::phys
diff --git a/src/KingSystem/Physics/System/physContactMgr.cpp b/src/KingSystem/Physics/System/physContactMgr.cpp
index 9359e04e..a467736c 100644
--- a/src/KingSystem/Physics/System/physContactMgr.cpp
+++ b/src/KingSystem/Physics/System/physContactMgr.cpp
@@ -214,7 +214,7 @@ void ContactMgr::removeCollisionEntriesWithBody(RigidBody* body) {
             continue;
 
         for (auto& colliding_pair : collision_info.getCollidingBodies().robustRange()) {
-            if (colliding_pair.body_b != body)
+            if (colliding_pair.bodies[1] != body)
                 continue;
 
             collision_info.getCollidingBodies().erase(&colliding_pair);
@@ -363,8 +363,8 @@ void ContactMgr::registerCollision(CollisionInfo* info, RigidBody* body_a, Rigid
 
     auto* entry = mCollidingBodiesFreeList.popFront();
     if (entry) {
-        entry->body_a = body_a;
-        entry->body_b = body_b;
+        entry->bodies[0] = body_a;
+        entry->bodies[1] = body_b;
         info->getCollidingBodies().pushBack(entry);
     }
 }
@@ -375,8 +375,8 @@ void ContactMgr::registerCollision(ContactLayerCollisionInfo* info, RigidBody* b
 
     auto* entry = mCollidingBodiesFreeList.popFront();
     if (entry) {
-        entry->body_a = body_a;
-        entry->body_b = body_b;
+        entry->bodies[0] = body_a;
+        entry->bodies[1] = body_b;
         info->getCollidingBodies().pushBack(entry);
     }
 }
@@ -386,7 +386,7 @@ void ContactMgr::unregisterCollision(CollisionInfo* info, RigidBody* body_a, Rig
     auto info_lock = sead::makeScopedLock(*info);
 
     for (auto& entry : info->getCollidingBodies()) {
-        if (entry.body_a == body_a && entry.body_b == body_b) {
+        if (entry.bodies[0] == body_a && entry.bodies[1] == body_b) {
             info->getCollidingBodies().erase(&entry);
             mCollidingBodiesFreeList.pushBack(&entry);
             break;
@@ -400,7 +400,7 @@ void ContactMgr::unregisterCollision(ContactLayerCollisionInfo* info, RigidBody*
     auto info_lock = sead::makeScopedLock(*info);
 
     for (auto& entry : info->getCollidingBodies()) {
-        if (entry.body_a == body_a && entry.body_b == body_b) {
+        if (entry.bodies[0] == body_a && entry.bodies[1] == body_b) {
             info->getCollidingBodies().erase(&entry);
             freeCollidingBodiesEntry(&entry);
             break;
@@ -415,8 +415,8 @@ void ContactMgr::unregisterCollisionWithBody(ContactLayerCollisionInfo* info, Ri
     auto body_layer = body->getContactLayer();
 
     for (auto& entry : info->getCollidingBodies().robustRange()) {
-        auto* body_a = entry.body_a;
-        auto* body_b = int(body_a->getContactLayer()) == body_layer ? body_a : entry.body_b;
+        auto* body_a = entry.bodies[0];
+        auto* body_b = int(body_a->getContactLayer()) == body_layer ? body_a : entry.bodies[1];
         if (body_b == body) {
             info->getCollidingBodies().erase(&entry);
             freeCollidingBodiesEntry(&entry);
diff --git a/src/KingSystem/Physics/System/physSystem.h b/src/KingSystem/Physics/System/physSystem.h
index 3a1ebba8..68fd7bac 100644
--- a/src/KingSystem/Physics/System/physSystem.h
+++ b/src/KingSystem/Physics/System/physSystem.h
@@ -10,6 +10,8 @@
 namespace ksys::phys {
 
 class CollisionInfo;
+class ContactLayerCollisionInfo;
+class ContactLayerCollisionInfoGroup;
 class ContactMgr;
 class GroupFilter;
 class MaterialTable;
@@ -69,8 +71,17 @@ public:
     // 0x00000071012169ac
     void freeCollisionInfo(CollisionInfo* info) const;
 
+    // 0x00000071012169b4
+    ContactLayerCollisionInfoGroup*
+    makeContactLayerCollisionInfoGroup(sead::Heap* heap, ContactLayer layer, int capacity,
+                                       const sead::SafeString& name);
+    // 0x00000071012169c0
+    void freeContactLayerCollisionInfoGroup(ContactLayerCollisionInfoGroup* group);
+    // 0x00000071012169c8
+    ContactLayerCollisionInfo* trackLayerPair(ContactLayer layer_a, ContactLayer layer_b);
+
     // 0x0000007101216a20
-    void registerRigidBodyForContactSystem(RigidBody* body);
+    void removeRigidBodyFromContactSystem(RigidBody* body);
 
     void removeSystemGroupHandler(SystemGroupHandler* handler);