diff --git a/Assets/Scripts/Editor/AssemblyPostProcess/AddInternalsVisibleToForAllUserAssembliesPostProcess.cs b/Assets/Scripts/Editor/AssemblyPostProcess/AddInternalsVisibleToForAllUserAssembliesPostProcess.cs index 9496b344..a2772470 100644 --- a/Assets/Scripts/Editor/AssemblyPostProcess/AddInternalsVisibleToForAllUserAssembliesPostProcess.cs +++ b/Assets/Scripts/Editor/AssemblyPostProcess/AddInternalsVisibleToForAllUserAssembliesPostProcess.cs @@ -1,4 +1,7 @@ -using System.IO; +using System; +using System.Collections.Generic; +using System.IO; +using System.Linq; using System.Reflection; using Mono.Cecil; using UnityEditor; @@ -7,7 +10,7 @@ namespace FastScriptReload.Editor.AssemblyPostProcess { [InitializeOnLoad] - public static class AddInternalsVisibleToForAllUserAssembliesPostProcess + public static class AddInternalsVisibleToForAllUserAssembliesPostProcess { public static readonly DirectoryInfo AdjustedAssemblyRoot; @@ -15,13 +18,14 @@ static AddInternalsVisibleToForAllUserAssembliesPostProcess() { AdjustedAssemblyRoot = new DirectoryInfo(Path.Combine(Application.dataPath, "..", "Temp", "Fast Script Reload", "AdjustedDlls")); } - + public static string CreateAssemblyWithInternalsContentsVisibleTo(Assembly changedAssembly, string visibleToAssemblyName) { - if(!AdjustedAssemblyRoot.Exists) - AdjustedAssemblyRoot.Create(); + if (!AdjustedAssemblyRoot.Exists) + AdjustedAssemblyRoot.Create(); - using (var assembly = AssemblyDefinition.ReadAssembly(changedAssembly.Location, new ReaderParameters { ReadWrite = false })) + var assemblyResolver = new UnityMonoEditorAssemblyResolver(changedAssembly.Location); + using (var assembly = AssemblyDefinition.ReadAssembly(changedAssembly.Location, new ReaderParameters { ReadWrite = false, AssemblyResolver = assemblyResolver })) { var mainModule = assembly.MainModule; @@ -35,7 +39,7 @@ public static string CreateAssemblyWithInternalsContentsVisibleTo(Assembly chang ); assembly.CustomAttributes.Add(attribute); - + var newAssemblyPath = new FileInfo(Path.Combine(AdjustedAssemblyRoot.FullName, assembly.Name.Name) + ".dll").FullName; assembly.Write(newAssemblyPath); @@ -43,4 +47,98 @@ public static string CreateAssemblyWithInternalsContentsVisibleTo(Assembly chang } } } + + class UnityMonoEditorAssemblyResolver : IAssemblyResolver + { + static readonly HashSet s_unityEditorSearchDirs; + static UnityMonoEditorAssemblyResolver() + { + // s_unityEditorSearchDirs will be populated on each domain reload + s_unityEditorSearchDirs = new(); + foreach (var loadedAssembly in AppDomain.CurrentDomain.GetAssemblies()) + { + if (loadedAssembly.IsDynamic) continue; + if (!File.Exists(loadedAssembly.Location)) continue; + var assemblyDir = Path.GetDirectoryName(loadedAssembly.Location); + s_unityEditorSearchDirs.Add(assemblyDir); + } + } + + private Dictionary cache; + private string majorSearchDir; + + public UnityMonoEditorAssemblyResolver(string targetAssemblyPath) + { + cache = new(); + majorSearchDir = Path.GetDirectoryName(targetAssemblyPath); + } + + private AssemblyDefinition GetAssembly(string file, ReaderParameters parameters) + { + parameters.AssemblyResolver ??= this; + return ModuleDefinition.ReadModule(file, parameters).Assembly; + } + + private bool TryResolveAssemblyInDirectory(string directory, ReaderParameters parameters, AssemblyNameReference name, out AssemblyDefinition assembly) + { + const string AssemblyExtension = ".dll"; + var filepath = Path.Combine(directory, name.Name + AssemblyExtension); + assembly = null; + if (File.Exists(filepath)) + { + try + { + assembly = GetAssembly(filepath, parameters); + return true; + } + catch { } + } + + return false; + } + + private AssemblyDefinition ResolveInternalCacheLess(AssemblyNameReference name, ReaderParameters parameters) + { + if (TryResolveAssemblyInDirectory(majorSearchDir, parameters, name, out var assembly)) + { + return assembly; + } + + foreach (var dir in s_unityEditorSearchDirs) + { + if (TryResolveAssemblyInDirectory(dir, parameters, name, out assembly)) + { + return assembly; + } + } + + throw new AssemblyResolutionException(name); + } + + public AssemblyDefinition Resolve(AssemblyNameReference name) + { + return Resolve(name, new ReaderParameters()); + } + + public AssemblyDefinition Resolve(AssemblyNameReference name, ReaderParameters parameters) + { + if (!cache.TryGetValue(name.FullName, out var value)) + { + cache[name.FullName] = value = ResolveInternalCacheLess(name, parameters); + } + + return value; + } + + public void Dispose() + { + foreach (var (path, assembly) in cache) + { + assembly.Dispose(); + } + + cache.Clear(); + } + + } } \ No newline at end of file