Offload the load context caching to bepinex

BepInEx.Core/Bootstrap/BaseChainloader.cs

@@ -146,7 +146,7 @@ public abstract class BaseChainloader<TPlugin>
     ///     Value holds info of the originating provider. These contexts are used by chainloader to
     ///     load plugins. Plugins that are loaded successfully are added to <see cref="Plugins" />.
     /// </summary>
-    public Dictionary<IPluginLoadContext, PluginInfo> LoadContexts { get; } = new();
+    internal Dictionary<CachedPluginLoadContext, PluginInfo> LoadContexts { get; } = new();
     
     /// <summary>
     ///     Collection of error chainloader messages that occured during plugin loading.
@@ -350,7 +350,7 @@ public abstract class BaseChainloader<TPlugin>
             LoadPlugins(providers, provider: true);
             FinishedProviders?.Invoke();
             
-            var plugins = TypeLoader.GetPluginsFromLoadContexts(LoadContexts.Keys, ToPluginInfo<TPlugin>, HasBepinPluginType<BepInPlugin>, "chainloader");
+            var plugins = TypeLoader.GetPluginsFromLoadContexts(LoadContexts.Keys.Cast<IPluginLoadContext>(), ToPluginInfo<TPlugin>, HasBepinPluginType<BepInPlugin>, "chainloader");
             Logger.Log(LogLevel.Info, $"{plugins.Count} plugin{(plugins.Count == 1 ? "" : "s")} to load");
             LoadPlugins(plugins, provider: false);
             Finished?.Invoke();
@@ -414,23 +414,24 @@ public abstract class BaseChainloader<TPlugin>
                     var pluginLoadContexts = ((BasePluginProvider)pluginInfo.Instance).GetPlugins();
                     foreach (IPluginLoadContext context in pluginLoadContexts)
                     {
-                        LoadContexts[context] = pluginInfo;
+                        var cachedContext = new CachedPluginLoadContext(context);
+                        LoadContexts[cachedContext] = pluginInfo;
                     }
                     ProviderLoaded?.Invoke(pluginInfo);
                 }
                 else
                 {
-                    pluginInfo.Source = LoadContexts[pluginInfo.LoadContext];
+                    pluginInfo.Source = LoadContexts[(CachedPluginLoadContext)pluginInfo.LoadContext];
                     pluginInfo.Instance = LoadPlugin(pluginInfo, ass);
                     PluginLoaded?.Invoke(pluginInfo);
-                    pluginInfo.LoadContext.Dispose();
+                    foreach (CachedPluginLoadContext context in LoadContexts.Keys)
+                        context.Dispose();
                 }
             }
             catch (Exception ex)
             {
                 invalidPlugins.Add(pluginInfo.Metadata.GUID);
                 existingPluginInfos.Remove(pluginInfo.Metadata.GUID);
-                pluginInfo.LoadContext?.Dispose();
 
                 Logger.Log(LogLevel.Error,
                            $"Error loading [{pluginInfo}]: {(ex is ReflectionTypeLoadException re ? TypeLoader.TypeLoadExceptionToString(re) : ex.ToString())}");

BepInEx.Core/Bootstrap/TypeLoader.cs

@@ -10,6 +10,47 @@ using Mono.Cecil;
 
 namespace BepInEx.Bootstrap;
 
+internal class CachedPluginLoadContext : IPluginLoadContext, IDisposable
+{
+    public IPluginLoadContext PluginLoadContext { get; }
+    private byte[] assemblyData;
+    private byte[] assemblySymbolsData;
+    private readonly Dictionary<string, byte[]> files = new();
+
+    public CachedPluginLoadContext(IPluginLoadContext pluginLoadContext)
+    {
+        PluginLoadContext = pluginLoadContext;
+    }
+
+    public string AssemblyIdentifier => PluginLoadContext.AssemblyIdentifier;
+    public string AssemblyHash => PluginLoadContext.AssemblyHash;
+    public byte[] GetAssemblyData()
+    {
+        return assemblyData ??= PluginLoadContext.GetAssemblyData();
+    }
+
+    public byte[] GetAssemblySymbolsData()
+    {
+        return assemblySymbolsData ??= PluginLoadContext.GetAssemblySymbolsData();
+    }
+
+    public byte[] GetFile(string relativePath)
+    {
+        if (!files.ContainsKey(relativePath))
+        {
+            files[relativePath] = PluginLoadContext.GetFile(relativePath);
+        }
+        return files[relativePath];
+    }
+
+    public void Dispose()
+    {
+        assemblyData = null;
+        assemblySymbolsData = null;
+        files.Clear();
+    }
+}
+
 /// <summary>
 ///     A cacheable metadata item. Can be used with <see cref="TypeLoader.LoadAssemblyCache{T}" /> and
 ///     <see cref="TypeLoader.SaveAssemblyCache{T}" /> to cache plugin metadata.

BepInEx.Core/Contract/IPluginLoadContext.cs

@@ -6,7 +6,7 @@ namespace BepInEx;
 ///     An interface that allows to dynamically load and track an assembly
 ///     which can be obtained from a provider
 /// </summary>
-public interface IPluginLoadContext : IDisposable
+public interface IPluginLoadContext
 {
     /// <summary>
     ///     An identifier that uniquely identifies an assembly from a provider no matter its revision