Fix generator memory leak (#193)

BepInEx · Dec 14, 2024 · f2b2e69 · f2b2e69
1 parent 9d4599d
commit f2b2e69
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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();