From c4b3a7f5169d0e07d9bcce68a6dde049fb855b96 Mon Sep 17 00:00:00 2001
From: Nicolas Gnyra <nicolas@gnyra.com>
Date: Wed, 5 Jul 2023 18:17:01 -0400
Subject: [PATCH] Re-enable unasserter and make it work on
 IDisposable/ILateDisposable

---
 SiraUtil/Plugin.cs                            |  6 ++-
 ...erter.cs => ZenjectLifecycleUnasserter.cs} | 38 +++++++++++--------
 2 files changed, 27 insertions(+), 17 deletions(-)
 rename SiraUtil/Tweaks/{InitializableUnasserter.cs => ZenjectLifecycleUnasserter.cs} (55%)

diff --git a/SiraUtil/Plugin.cs b/SiraUtil/Plugin.cs
index 21a0e12..258f702 100644
--- a/SiraUtil/Plugin.cs
+++ b/SiraUtil/Plugin.cs
@@ -2,6 +2,7 @@
 using IPA;
 using IPA.Config.Stores;
 using IPA.Loader;
+using IPA.Utilities.Async;
 #if DEBUG
 using SiraUtil.Affinity.Harmony.Generator;
 #endif
@@ -10,6 +11,7 @@
 using SiraUtil.Tools.FPFC;
 using SiraUtil.Zenject;
 using System.Reflection;
+using System.Threading.Tasks;
 using Zenject;
 using Conf = IPA.Config.Config;
 using IPALogger = IPA.Logging.Logger;
@@ -70,7 +72,9 @@ public void OnEnable()
         public void OnDisable()
         {
             _zenjectManager.Disable();
-            _harmony.UnpatchSelf();
+
+            // delay so DisposableManager is able to run before we unpatch on shutdown
+            UnityMainThreadTaskScheduler.Factory.StartNew(() => _harmony.UnpatchSelf()).ContinueWith((task) => Log.Error($"Failed to unpatch\n{task.Exception}"), TaskContinuationOptions.OnlyOnFaulted);
 
 #if DEBUG
             DynamicHarmonyPatchGenerator.Save();
diff --git a/SiraUtil/Tweaks/InitializableUnasserter.cs b/SiraUtil/Tweaks/ZenjectLifecycleUnasserter.cs
similarity index 55%
rename from SiraUtil/Tweaks/InitializableUnasserter.cs
rename to SiraUtil/Tweaks/ZenjectLifecycleUnasserter.cs
index 6433dd6..8c7adaf 100644
--- a/SiraUtil/Tweaks/InitializableUnasserter.cs
+++ b/SiraUtil/Tweaks/ZenjectLifecycleUnasserter.cs
@@ -9,32 +9,38 @@
 
 namespace SiraUtil.Tweaks
 {
-    class InitializableUnasserter
+    [HarmonyPatch]
+    class ZenjectLifecycleUnasserter
     {
         private static readonly MethodInfo _newFail = SymbolExtensions.GetMethodInfo(() => NewErrorBehavior(null!, null!, null!));
         private static readonly MethodInfo _rootMethod = typeof(ModestTree.Assert).GetMethod(nameof(ModestTree.Assert.CreateException), new Type[] { typeof(Exception), typeof(string), typeof(object[]) });
 
-        public static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> instructions)
+        public static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> instructions, ILGenerator generator)
         {
-            var codes = instructions.ToList();
-            for (int i = 0; i < codes.Count; i++)
-            {
-                var code = codes[i];
+            CodeMatcher codeMatcher = new(instructions, generator);
+            codeMatcher
+                .MatchForward(false, new CodeMatch(OpCodes.Call, _rootMethod), new CodeMatch(OpCodes.Throw))
+                .ThrowIfInvalid($"Call to {nameof(ModestTree.Assert.CreateException)} & throw not found");
 
-                if (code.operand != null && code.Is(OpCodes.Call, _rootMethod))
-                {
-                    codes[i] = new CodeInstruction(OpCodes.Callvirt, _newFail);
-                    codes.RemoveAt(i + 1);
-                    break;
-                }
-            }
-            return codes.AsEnumerable();
+            List<ExceptionBlock> blocks = codeMatcher.InstructionAt(1).blocks;
+            codeMatcher
+                .RemoveInstructions(2)
+                .Insert(new CodeInstruction(OpCodes.Call, _newFail) { blocks = blocks });
+
+            return codeMatcher.InstructionEnumeration();
+        }
+
+        public static IEnumerable<MethodInfo> TargetMethods()
+        {
+            yield return AccessTools.Method(typeof(InitializableManager), nameof(InitializableManager.Initialize));
+            yield return AccessTools.Method(typeof(DisposableManager), nameof(DisposableManager.Dispose));
+            yield return AccessTools.Method(typeof(DisposableManager), nameof(DisposableManager.LateDispose));
         }
 
-        private static void NewErrorBehavior(Exception exception, string _, object[] parameters)
+        private static void NewErrorBehavior(Exception exception, string message, params object[] parameters)
         {
             var failedType = (Type)parameters[0];
-            string failText = $"Error occurred while initializing {nameof(IInitializable)} with type '{failedType.FullName}'";
+            string failText = string.Format(message, failedType.FullName);
             if (failedType.Name != failedType.FullName)
             {
                 Plugin.Log.Critical(failText);