From 02584dfcfb9f76533238d6cdc4e10eed4a1178ff Mon Sep 17 00:00:00 2001
From: pas2704 <pas2704@outlook.com>
Date: Sun, 3 Nov 2024 18:34:34 -0500
Subject: [PATCH] Allow plugins to find non-harmony types in other plugins

---
 .../Patches/IntrospectionPatches.cs           | 13 ++++++
 .../Loader/PluginAssemblyLoadContext.cs       | 44 ++++++++++++++++++-
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/CringeLauncher/Patches/IntrospectionPatches.cs b/CringeLauncher/Patches/IntrospectionPatches.cs
index f647dc4..6aaed9b 100644
--- a/CringeLauncher/Patches/IntrospectionPatches.cs
+++ b/CringeLauncher/Patches/IntrospectionPatches.cs
@@ -78,6 +78,19 @@ public static class IntrospectionPatches
         return false;
     }
 
+    [HarmonyPrefix, HarmonyPatch(typeof(AccessTools), nameof(AccessTools.TypeByName))]
+    private static bool TypeByNameHarmonyPrefix(ref Type __result, string name)
+    {
+        if (!PluginAssemblyLoadContext.TypeToAssembly.TryGetValue(name, out var assembly))
+            return true;
+
+        var module = assembly.GetMainModule();
+
+        __result = module.GetType(name, true, false)!;
+
+        return false;
+    }
+
     [HarmonyPrefix, HarmonyPatch(typeof(AccessTools), nameof(AccessTools.AllAssemblies))]
     private static bool AllAssembliesHarmonyPrefix(ref IEnumerable<Assembly> __result)
     {
diff --git a/CringePlugins/Loader/PluginAssemblyLoadContext.cs b/CringePlugins/Loader/PluginAssemblyLoadContext.cs
index c587f84..dc65715 100644
--- a/CringePlugins/Loader/PluginAssemblyLoadContext.cs
+++ b/CringePlugins/Loader/PluginAssemblyLoadContext.cs
@@ -1,23 +1,51 @@
-using System.Reflection;
+using System.Collections.Concurrent;
+using System.Reflection;
 using System.Runtime.Loader;
 using CringeBootstrap.Abstractions;
+using CringePlugins.Utils;
 using SharedCringe.Loader;
 
 namespace CringePlugins.Loader;
 
 internal class PluginAssemblyLoadContext : DerivedAssemblyLoadContext
 {
+    //todo: refactor?
+    public static readonly ConcurrentDictionary<string, Assembly> TypeToAssembly = [];
+
     private readonly string _entrypointPath;
     private readonly AssemblyDependencyResolver _dependencyResolver;
+    private readonly HashSet<string> _loadedTypes = [];
     private Assembly? _assembly;
 
     internal PluginAssemblyLoadContext(ICoreLoadContext parentContext, string entrypointPath) : base(parentContext, $"Plugin Context {Path.GetFileNameWithoutExtension(entrypointPath)}")
     {
         _entrypointPath = entrypointPath;
         _dependencyResolver = new(entrypointPath);
+
+        Unloading += OnUnload;
     }
 
-    public Assembly LoadEntrypoint() => _assembly ??= LoadFromAssemblyPath(_entrypointPath);
+    public Assembly LoadEntrypoint()
+    {
+        if (_assembly is not null)
+            return _assembly;
+
+        _assembly = LoadFromAssemblyPath(_entrypointPath);
+
+        var module = _assembly.GetMainModule();
+
+        foreach (var type in module.GetTypes())
+        {
+            var name = type.FullName?.Replace('/', '+');
+
+            if (string.IsNullOrEmpty(name) || !_loadedTypes.Add(name))
+                continue;
+
+            TypeToAssembly[name] = _assembly;
+        }
+
+        return _assembly;
+    }
 
     protected override Assembly? Load(AssemblyName assemblyName)
     {
@@ -34,4 +62,16 @@ internal class PluginAssemblyLoadContext : DerivedAssemblyLoadContext
         
         return base.LoadUnmanagedDll(unmanagedDllName);
     }
+
+    private static void OnUnload(AssemblyLoadContext context)
+    {
+        if (context is not PluginAssemblyLoadContext pluginContext)
+            return;
+
+        foreach (var typeStr in pluginContext._loadedTypes)
+        {
+            TypeToAssembly.Remove(typeStr);
+        }
+        pluginContext._loadedTypes.Clear();
+    }
 }
\ No newline at end of file