From 7814fd527b259b5a7c499f4e78e956f789f65fdf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A9o=20Lam?= <leo@leolam.fr>
Date: Thu, 21 Jan 2021 23:36:01 +0100
Subject: [PATCH] uking/ui: Implement shield sorting for inventory

---
 data/uking_functions.csv           |  4 +-
 src/Game/UI/uiPauseMenuDataMgr.cpp | 81 ++++++++++++++++++++++++++++++
 src/Game/UI/uiPauseMenuDataMgr.h   | 10 +++-
 3 files changed, 92 insertions(+), 3 deletions(-)

diff --git a/data/uking_functions.csv b/data/uking_functions.csv
index 05467a26..62d25076 100644
--- a/data/uking_functions.csv
+++ b/data/uking_functions.csv
@@ -56381,13 +56381,13 @@
 0x0000007100977700,sub_7100977700,248,_ZN5uking2ui22pouchItemSortPredicateEPKNS0_9PouchItemES3_
 0x00000071009777f8,sub_71009777F8,272,
 0x0000007100977908,sub_7100977908,288,
-0x0000007100977a28,sub_7100977A28,272,
+0x0000007100977a28,sub_7100977A28,272,_ZN5uking2ui13compareShieldEPKNS0_9PouchItemES3_PN4ksys3act8InfoDataE
 0x0000007100977b38,sub_7100977B38,336,_ZN5uking2ui12compareArmorEPKNS0_9PouchItemES3_PN4ksys3act8InfoDataE
 0x0000007100977c88,sub_7100977C88,308,
 0x0000007100977dbc,sub_7100977DBC,540,
 0x0000007100977fd8,sub_7100977FD8,116,_ZN5uking2ui14compareKeyItemEPKNS0_9PouchItemES3_PN4ksys3act8InfoDataE
 0x000000710097804c,sub_710097804C,676,
-0x00000071009782f0,sub_71009782F0,664,
+0x00000071009782f0,sub_71009782F0,664,_ZN5uking2uiL15doCompareShieldEPKNS0_9PouchItemES3_PN4ksys3act8InfoDataE
 0x0000007100978588,sub_7100978588,480,_ZN5uking2ui16getCookItemOrderEPKNS0_9PouchItemEPN4ksys3act8InfoDataE
 0x0000007100978768,sub_7100978768,304,_ZN5uking2ui30pouchItemSortPredicateForArrowEPKNS0_9PouchItemES3_
 0x0000007100978898,sub_7100978898,440,
diff --git a/src/Game/UI/uiPauseMenuDataMgr.cpp b/src/Game/UI/uiPauseMenuDataMgr.cpp
index d49d159e..74c3f8dc 100644
--- a/src/Game/UI/uiPauseMenuDataMgr.cpp
+++ b/src/Game/UI/uiPauseMenuDataMgr.cpp
@@ -133,6 +133,37 @@ void getSameGroupActorName(sead::SafeString* group, const sead::SafeString& item
     }
 }
 
+int getWeaponModifierSortKey(sead::TypedBitFlag<act::WeaponModifier> flags) {
+    if (flags.isOn(act::WeaponModifier::AddAtk))
+        return flags.isOn(act::WeaponModifier::IsYellow) ? 0 : 1;
+
+    if (flags.isOn(act::WeaponModifier::AddLife))
+        return flags.isOn(act::WeaponModifier::IsYellow) ? 4 : 5;
+
+    if (flags.isOn(act::WeaponModifier::AddThrow))
+        return 6;
+
+    if (flags.isOn(act::WeaponModifier::AddCrit))
+        return 7;
+
+    if (flags.isOn(act::WeaponModifier::AddGuard))
+        return flags.isOn(act::WeaponModifier::IsYellow) ? 2 : 3;
+
+    if (flags.isOn(act::WeaponModifier::AddSurfMaster))
+        return 8;
+
+    if (flags.isOn(act::WeaponModifier::AddSpreadFire))
+        return 10;
+
+    if (flags.isOn(act::WeaponModifier::AddZoomRapid))
+        return 11;
+
+    if (flags.isOn(act::WeaponModifier::AddRapidFire))
+        return 9;
+
+    return 12;
+}
+
 }  // namespace
 
 int pouchItemSortPredicate(const PouchItem* lhs, const PouchItem* rhs);
@@ -1857,6 +1888,56 @@ static s32 compareSortKeys(const PouchItem* lhs, const PouchItem* rhs, ksys::act
     return 0;
 }
 
+static int getShieldGuardPower(const PouchItem* item, ksys::act::InfoData* data) {
+    int power = ksys::act::getWeaponCommonGuardPower(data, item->getName().cstr());
+    if (item->getWeaponAddFlags().isOn(act::WeaponModifier::AddGuard))
+        power += item->getWeaponAddValue();
+    return power;
+}
+
+static int doCompareShield(const PouchItem* lhs, const PouchItem* rhs, ksys::act::InfoData* data) {
+    const int gp1 = getShieldGuardPower(lhs, data);
+    const int gp2 = getShieldGuardPower(rhs, data);
+    // Higher is better
+    if (gp1 > gp2)
+        return -1;
+    if (gp1 < gp2)
+        return 1;
+
+    const int mod1 = getWeaponModifierSortKey(lhs->getWeaponAddFlags());
+    const int mod2 = getWeaponModifierSortKey(rhs->getWeaponAddFlags());
+    // Lower is better
+    if (mod1 < mod2)
+        return -1;
+    if (mod1 > mod2)
+        return 1;
+
+    const int val1 = lhs->getValue();
+    const int val2 = rhs->getValue();
+    // Higher is better
+    if (val1 > val2)
+        return -1;
+    if (val1 < val2)
+        return 1;
+
+    return 0;
+}
+
+int compareShield(const PouchItem* lhs, const PouchItem* rhs, ksys::act::InfoData* data) {
+    if (ksys::gdt::getFlag_SortTypeShieldPouch()) {
+        if (auto cmp = compareSortKeys(lhs, rhs, data))
+            return cmp;
+
+        return doCompareShield(lhs, rhs, data);
+
+    } else {
+        if (auto cmp = doCompareShield(lhs, rhs, data))
+            return cmp;
+
+        return compareSortKeys(lhs, rhs, data);
+    }
+}
+
 int compareArmor(const PouchItem* lhs, const PouchItem* rhs, ksys::act::InfoData* data) {
     if (ksys::gdt::getFlag_SortTypeArmorPouch()) {
         if (auto cmp = compareSortKeys(lhs, rhs, data))
diff --git a/src/Game/UI/uiPauseMenuDataMgr.h b/src/Game/UI/uiPauseMenuDataMgr.h
index ec36f7ed..4348abbe 100644
--- a/src/Game/UI/uiPauseMenuDataMgr.h
+++ b/src/Game/UI/uiPauseMenuDataMgr.h
@@ -8,6 +8,7 @@
 #include <heap/seadDisposer.h>
 #include <math/seadVector.h>
 #include <prim/seadSafeString.h>
+#include <prim/seadTypedBitFlag.h>
 #include <thread/seadCriticalSection.h>
 #include "KingSystem/Utils/Types.h"
 
@@ -21,8 +22,9 @@ class InfoData;
 }  // namespace ksys::act
 
 namespace uking::act {
+enum class WeaponModifier : u32;
 struct WeaponModifierInfo;
-}
+}  // namespace uking::act
 
 namespace ksys {
 struct CookItem;
@@ -192,6 +194,12 @@ public:
     WeaponData& getWeaponData() { return mData.weapon; }
     const WeaponData& getWeaponData() const { return mData.weapon; }
 
+    sead::TypedBitFlag<act::WeaponModifier> getWeaponAddFlags() const {
+        if (!isWeapon())
+            return {};
+        return sead::TypedBitFlag<act::WeaponModifier>{mData.weapon.mAddType};
+    }
+
     u32 getWeaponAddValue() const {
         if (!isWeapon())
             return 0;