From ab324da11bb0833cf967edd03736c217ce265b29 Mon Sep 17 00:00:00 2001
From: slxdy <devslidy@gmail.com>
Date: Sat, 14 Dec 2024 01:57:05 +0100
Subject: [PATCH] Fix generator memory leak

---
 Il2CppInterop.Generator/Contexts/AssemblyRewriteContext.cs | 5 +----
 Il2CppInterop.Generator/Contexts/RewriteGlobalContext.cs   | 2 ++
 Il2CppInterop.Generator/Passes/Pass16ScanMethodRefs.cs     | 4 ++--
 .../Passes/Pass81FillUnstrippedMethodBodies.cs             | 3 +++
 .../Runners/InteropAssemblyGenerator.cs                    | 7 +++++++
 5 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/Il2CppInterop.Generator/Contexts/AssemblyRewriteContext.cs b/Il2CppInterop.Generator/Contexts/AssemblyRewriteContext.cs
index 01a714cf..057a2ac3 100644
--- a/Il2CppInterop.Generator/Contexts/AssemblyRewriteContext.cs
+++ b/Il2CppInterop.Generator/Contexts/AssemblyRewriteContext.cs
@@ -9,9 +9,6 @@ namespace Il2CppInterop.Generator.Contexts;
 [DebuggerDisplay($"{{{nameof(GetDebuggerDisplay)}(),nq}}")]
 public class AssemblyRewriteContext
 {
-    // TODO: Dispose
-    private static readonly Dictionary<ModuleDefinition, RuntimeAssemblyReferences> ImportsMap = new();
-
     public readonly RewriteGlobalContext GlobalContext;
 
     public readonly RuntimeAssemblyReferences Imports;
@@ -30,7 +27,7 @@ public AssemblyRewriteContext(RewriteGlobalContext globalContext, AssemblyDefini
         NewAssembly = newAssembly;
         GlobalContext = globalContext;
 
-        Imports = ImportsMap.GetOrCreate(newAssembly.ManifestModule!,
+        Imports = globalContext.ImportsMap.GetOrCreate(newAssembly.ManifestModule!,
             mod => new RuntimeAssemblyReferences(mod, globalContext));
     }
 
diff --git a/Il2CppInterop.Generator/Contexts/RewriteGlobalContext.cs b/Il2CppInterop.Generator/Contexts/RewriteGlobalContext.cs
index a9f390f1..454a926f 100644
--- a/Il2CppInterop.Generator/Contexts/RewriteGlobalContext.cs
+++ b/Il2CppInterop.Generator/Contexts/RewriteGlobalContext.cs
@@ -19,6 +19,8 @@ public class RewriteGlobalContext : IDisposable
 
     internal readonly Dictionary<(object?, string, int), List<TypeDefinition>> RenameGroups = new();
 
+    internal readonly Dictionary<ModuleDefinition, RuntimeAssemblyReferences> ImportsMap = new();
+
     public RewriteGlobalContext(GeneratorOptions options, IIl2CppMetadataAccess gameAssemblies,
         IMetadataAccess unityAssemblies)
     {
diff --git a/Il2CppInterop.Generator/Passes/Pass16ScanMethodRefs.cs b/Il2CppInterop.Generator/Passes/Pass16ScanMethodRefs.cs
index 9dea948c..f16a3541 100644
--- a/Il2CppInterop.Generator/Passes/Pass16ScanMethodRefs.cs
+++ b/Il2CppInterop.Generator/Passes/Pass16ScanMethodRefs.cs
@@ -10,8 +10,8 @@ namespace Il2CppInterop.Generator.Passes;
 
 public static class Pass16ScanMethodRefs
 {
-    public static readonly HashSet<long> NonDeadMethods = new();
-    public static IDictionary<long, List<XrefInstance>> MapOfCallers = new Dictionary<long, List<XrefInstance>>();
+    internal static HashSet<long> NonDeadMethods = new();
+    internal static IDictionary<long, List<XrefInstance>> MapOfCallers = new Dictionary<long, List<XrefInstance>>();
 
     public static void DoPass(RewriteGlobalContext context, GeneratorOptions options)
     {
diff --git a/Il2CppInterop.Generator/Passes/Pass81FillUnstrippedMethodBodies.cs b/Il2CppInterop.Generator/Passes/Pass81FillUnstrippedMethodBodies.cs
index ffcef3a1..705e419a 100644
--- a/Il2CppInterop.Generator/Passes/Pass81FillUnstrippedMethodBodies.cs
+++ b/Il2CppInterop.Generator/Passes/Pass81FillUnstrippedMethodBodies.cs
@@ -32,6 +32,9 @@ public static void DoPass(RewriteGlobalContext context)
             }
         }
 
+        StuffToProcess.Clear();
+        StuffToProcess.Capacity = 0;
+
         Logger.Instance.LogInformation("IL unstrip statistics: {MethodsSucceeded} successful, {MethodsFailed} failed", methodsSucceeded,
             methodsFailed);
     }
diff --git a/Il2CppInterop.Generator/Runners/InteropAssemblyGenerator.cs b/Il2CppInterop.Generator/Runners/InteropAssemblyGenerator.cs
index b1c66b7c..f3f3efe7 100644
--- a/Il2CppInterop.Generator/Runners/InteropAssemblyGenerator.cs
+++ b/Il2CppInterop.Generator/Runners/InteropAssemblyGenerator.cs
@@ -1,4 +1,5 @@
 using Il2CppInterop.Common;
+using Il2CppInterop.Common.XrefScans;
 using Il2CppInterop.Generator.Contexts;
 using Il2CppInterop.Generator.MetadataAccess;
 using Il2CppInterop.Generator.Passes;
@@ -201,6 +202,12 @@ public void Run(GeneratorOptions options)
             Pass91GenerateMethodPointerMap.DoPass(rewriteContext, options);
         }
 
+        using (new TimingCookie("Clearing static data"))
+        {
+            Pass16ScanMethodRefs.MapOfCallers = new Dictionary<long, List<XrefInstance>>();
+            Pass16ScanMethodRefs.NonDeadMethods = [];
+        }
+
         Logger.Instance.LogInformation("Done!");
 
         rewriteContext.Dispose();