From 82d9d8218ead1aecadbcb4d56e752c9dfb334d3d Mon Sep 17 00:00:00 2001
From: Menci <mencici@msn.com>
Date: Sun, 26 Jan 2025 17:00:03 +0800
Subject: [PATCH] fix

---
 .../Fix/Stability/SanitizeUserData.cs         | 22 ++++++------
 AquaMai.Mods/GameSystem/Unlock.cs             | 34 +++++++++++--------
 2 files changed, 30 insertions(+), 26 deletions(-)

diff --git a/AquaMai.Mods/Fix/Stability/SanitizeUserData.cs b/AquaMai.Mods/Fix/Stability/SanitizeUserData.cs
index ee3ad65..958af1b 100644
--- a/AquaMai.Mods/Fix/Stability/SanitizeUserData.cs
+++ b/AquaMai.Mods/Fix/Stability/SanitizeUserData.cs
@@ -71,12 +71,12 @@ private static Variant OnNetPacketComplete(string api, Variant request, Variant
     private static void OnUserDataResponse(ProxyObject _, ProxyObject response)
     {
         var userData = GetObjectOrSetDefault(response, "userData");
-        SanitizeItemIdField(userData, "iconId", "GetIcons");
-        SanitizeItemIdField(userData, "plateId", "GetPlates");
-        SanitizeItemIdField(userData, "titleId", "GetTitles");
-        SanitizeItemIdField(userData, "partnerId", "GetPartners");
-        SanitizeItemIdField(userData, "frameId", "GetFrames");
-        SanitizeItemIdField(userData, "selectMapId", "GetMapDatas");
+        SanitizeItemIdField(userData, "iconId", "GetIcons", false);
+        SanitizeItemIdField(userData, "plateId", "GetPlates", false);
+        SanitizeItemIdField(userData, "titleId", "GetTitles", false);
+        SanitizeItemIdField(userData, "partnerId", "GetPartners", false);
+        SanitizeItemIdField(userData, "frameId", "GetFrames", false);
+        SanitizeItemIdField(userData, "selectMapId", "GetMapDatas", false);
         var charaSlot = GetArrayOrSetDefault(userData, "charaSlot");
         for (var i = 0; i < 5; i++)
         {
@@ -98,8 +98,8 @@ private static void OnUserDataResponse(ProxyObject _, ProxyObject response)
     private static void OnUserExtendResponse(ProxyObject _, ProxyObject response)
     {
         var userExtend = GetObjectOrSetDefault(response, "userExtend");
-        SanitizeItemIdField(userExtend, "selectMusicId", "GetMusics");
-        SanitizeItemIdField(userExtend, "selectDifficultyId", "GetMusicDifficultys");
+        SanitizeItemIdField(userExtend, "selectMusicId", "GetMusics", true);
+        SanitizeItemIdField(userExtend, "selectDifficultyId", "GetMusicDifficultys", true);
         // categoryIndex?
         // musicIndex?
         SanitizeEnumFieldIfDefined(userExtend, "selectScoreType", ResolveEnumType("ConstParameter.ScoreKind"));
@@ -269,12 +269,12 @@ private static void SanitizeEnumFieldIfDefined(ProxyObject obj, string fieldName
             value => enumType == null || Enum.IsDefined(enumType, value),
             enumType == null ? 0 : (int)enumType.GetEnumValues().GetValue(0));
 
-    private static void SanitizeItemIdField(ProxyObject obj, string fieldName, string dataManagerGetDictionaryMethod) =>
+    private static void SanitizeItemIdField(ProxyObject obj, string fieldName, string dataManagerGetDictionaryMethod, bool defaultZero) =>
         SanitizeInt32Field(
             obj,
             fieldName,
-            itemId => SafelyCheckItemId(dataManagerGetDictionaryMethod, itemId),
-            SafelyGetDefaultItemId(dataManagerGetDictionaryMethod));
+            itemId => (itemId == 0 && defaultZero) || SafelyCheckItemId(dataManagerGetDictionaryMethod, itemId),
+            defaultZero ? 0 : SafelyGetDefaultItemId(dataManagerGetDictionaryMethod));
 
     // The corresponding DataManager methods may not exist in all game versions
     private static object SafelyGetDataMangerDictionary(string dataManagerGetDictionaryMethod)
diff --git a/AquaMai.Mods/GameSystem/Unlock.cs b/AquaMai.Mods/GameSystem/Unlock.cs
index 220f135..8d10164 100644
--- a/AquaMai.Mods/GameSystem/Unlock.cs
+++ b/AquaMai.Mods/GameSystem/Unlock.cs
@@ -251,14 +251,14 @@ public class CollectionHook
                         pair
                             .GetType()
                             .GetProperty("Key")
-                            .GetValue(pair) is not int id
-                                ? null
-                                : new Manager.UserDatas.UserItem
-                                  {
-                                      itemId = id,
-                                      stock = 1,
-                                      isValid = true
-                                  })
+                            .GetValue(pair))
+                    .Select(id =>
+                        new Manager.UserDatas.UserItem
+                        {
+                            itemId = (int)id,
+                            stock = 1,
+                            isValid = true
+                        })
                     .ToList();
             allUnlockedItemsCache[dataManagerMethod] = result;
             return result;
@@ -268,34 +268,38 @@ public class CollectionHook
 
         public record PropertyChangeLog(object From, object To);
 
-        public static void Prefix(out Dictionary<PropertyInfo, PropertyChangeLog> __state)
+        public static void Prefix(out Dictionary<UserData, Dictionary<PropertyInfo, PropertyChangeLog>> __state)
         {
             __state = [];
             ModifyUserData(false, ref __state);
         }
 
-        public static void Postfix(Dictionary<PropertyInfo, PropertyChangeLog> __state)
+        public static void Postfix(Dictionary<UserData, Dictionary<PropertyInfo, PropertyChangeLog>> __state)
         {
-            ModifyUserData(false, ref __state);
+            ModifyUserData(true, ref __state);
         }
 
-        private static void ModifyUserData(bool restore, ref Dictionary<PropertyInfo, PropertyChangeLog> backup)
+        private static void ModifyUserData(bool restore, ref Dictionary<UserData, Dictionary<PropertyInfo, PropertyChangeLog>> backup)
         {
             for (int i = 0; i < 2; i++)
             {
                 var userData = UserDataManager.Instance.GetUserData(i);
                 if (!userData.IsEntry) continue;
+                if (!backup.TryGetValue(userData, out var userBackup))
+                {
+                    backup[userData] = userBackup = [];
+                }
                 foreach (var (_, userDataProperty, dataManagerMethod) in collectionHooks)
                 {
                     var currentValue = userDataProperty.GetValue(userData);
                     if (restore)
                     {
-                        if (!backup.TryGetValue(userDataProperty, out var backupData))
+                        if (!userBackup.TryGetValue(userDataProperty, out var backupData))
                         {
                             MelonLogger.Error($"[Unlock.CollectionHook] Failed to restore {userDataProperty.Name} to the original value. Backup data not found.");
                             continue;
                         }
-                        else if (currentValue != backupData.From)
+                        else if (currentValue != backupData.To)
                         {
                             MelonLogger.Error($"[Unlock.CollectionHook] Failed to restore {userDataProperty.Name} to the original value. Value changed unexpectedly, incompatible mods loaded?");
                             continue;
@@ -305,7 +309,7 @@ private static void ModifyUserData(bool restore, ref Dictionary<PropertyInfo, Pr
                     else
                     {
                         var allUnlockedItems = GetAllUnlockedItemList(dataManagerMethod);
-                        backup[userDataProperty] = new(From: currentValue, To: allUnlockedItems);
+                        userBackup[userDataProperty] = new(From: currentValue, To: allUnlockedItems);
                         userDataProperty.SetValue(userData, allUnlockedItems);
                     }
                 }