Add AssemblyListSnapshot.TryGetSimilarModuleAsync

ILSpy/AssemblyListSnapshot.cs

@@ -21,6 +21,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Linq;
 using System.Threading.Tasks;
 
 using ICSharpCode.Decompiler.Metadata;
@@ -33,7 +34,7 @@ namespace ICSharpCode.ILSpy
 		readonly ImmutableArray<LoadedAssembly> assemblies;
 		Dictionary<string, PEFile>? asmLookupByFullName;
 		Dictionary<string, PEFile>? asmLookupByShortName;
-
+		Dictionary<string, List<(PEFile module, Version version)>>? asmLookupByShortNameGrouped;
 		public ImmutableArray<LoadedAssembly> Assemblies => assemblies;
 
 		public AssemblyListSnapshot(ImmutableArray<LoadedAssembly> assemblies)
@@ -60,6 +61,20 @@ namespace ICSharpCode.ILSpy
 			return null;
 		}
 
+		public async Task<PEFile?> TryGetSimilarModuleAsync(IAssemblyReference reference)
+		{
+			var lookup = LazyInit.VolatileRead(ref asmLookupByShortNameGrouped);
+			if (lookup == null)
+			{
+				lookup = await CreateLoadedAssemblyShortNameGroupLookupAsync().ConfigureAwait(false);
+				lookup = LazyInit.GetOrSet(ref asmLookupByShortNameGrouped, lookup);
+			}
+
+			if (!lookup.TryGetValue(reference.Name, out var candidates))
+				return null;
+			return candidates.FirstOrDefault(c => c.version >= reference.Version).module ?? candidates.Last().module;
+		}
+
 		private async Task<Dictionary<string, PEFile>> CreateLoadedAssemblyLookupAsync(bool shortNames)
 		{
 			var result = new Dictionary<string, PEFile>(StringComparer.OrdinalIgnoreCase);
@@ -93,6 +108,44 @@ namespace ICSharpCode.ILSpy
 			return result;
 		}
 
+		private async Task<Dictionary<string, List<(PEFile module, Version version)>>> CreateLoadedAssemblyShortNameGroupLookupAsync()
+		{
+			var result = new Dictionary<string, List<(PEFile module, Version version)>>(StringComparer.OrdinalIgnoreCase);
+
+			foreach (LoadedAssembly loaded in assemblies)
+			{
+				try
+				{
+					var module = await loaded.GetPEFileOrNullAsync().ConfigureAwait(false);
+					var reader = module?.Metadata;
+					if (reader == null || !reader.IsAssembly)
+						continue;
+					var asmDef = reader.GetAssemblyDefinition();
+					var asmDefName = reader.GetString(asmDef.Name);
+
+					var line = (module!, version: asmDef.Version);
+
+					if (!result.TryGetValue(asmDefName, out var existing))
+					{
+						existing = new List<(PEFile module, Version version)>();
+						result.Add(asmDefName, existing);
+						existing.Add(line);
+						continue;
+					}
+
+					int index = existing.BinarySearch(line.version, l => l.version);
+					index = index < 0 ? ~index : index + 1;
+					existing.Insert(index, line);
+				}
+				catch (BadImageFormatException)
+				{
+					continue;
+				}
+			}
+
+			return result;
+		}
+
 		/// <summary>
 		/// Gets all loaded assemblies recursively, including assemblies found in bundles or packages.
 		/// </summary>

ILSpy/LoadedAssembly.cs

@@ -522,7 +522,7 @@ namespace ICSharpCode.ILSpy
 				string tfm = await tfmTask.ConfigureAwait(false);
 
 				// 1) try to find exact match by tfm + full asm name in loaded assemblies
-				module = await alreadyLoadedAssemblies.TryGetModuleAsync(reference, tfm);
+				module = await alreadyLoadedAssemblies.TryGetModuleAsync(reference, tfm).ConfigureAwait(false);
 				if (module != null)
 				{
 					referenceLoadInfo.AddMessageOnce(reference.FullName, MessageKind.Info, "Success - Found in Assembly List");
@@ -545,7 +545,7 @@ namespace ICSharpCode.ILSpy
 					}
 					if (asm != null)
 					{
-						referenceLoadInfo.AddMessage(reference.ToString(), MessageKind.Info, "Success - Loading from: " + file);
+						referenceLoadInfo.AddMessage(reference.FullName, MessageKind.Info, "Success - Loading from: " + file);
 						return await asm.GetPEFileOrNullAsync().ConfigureAwait(false);
 					}
 					return null;
@@ -553,33 +553,16 @@ namespace ICSharpCode.ILSpy
 				else
 				{
 					// Assembly not found; try to find a similar-enough already-loaded assembly
-					var candidates = new List<(LoadedAssembly assembly, Version version)>();
-
-					foreach (LoadedAssembly loaded in alreadyLoadedAssemblies.Assemblies)
+					module = await alreadyLoadedAssemblies.TryGetSimilarModuleAsync(reference).ConfigureAwait(false);
+					if (module == null)
 					{
-						module = await loaded.GetPEFileOrNullAsync().ConfigureAwait(false);
-						var reader = module?.Metadata;
-						if (reader == null || !reader.IsAssembly)
-							continue;
-						var asmDef = reader.GetAssemblyDefinition();
-						var asmDefName = reader.GetString(asmDef.Name);
-						if (reference.Name.Equals(asmDefName, StringComparison.OrdinalIgnoreCase))
-						{
-							candidates.Add((loaded, asmDef.Version));
-						}
+						referenceLoadInfo.AddMessageOnce(reference.FullName, MessageKind.Error, "Could not find reference: " + reference.FullName);
 					}
-
-					if (candidates.Count == 0)
+					else
 					{
-						referenceLoadInfo.AddMessageOnce(reference.ToString(), MessageKind.Error, "Could not find reference: " + reference);
-						return null;
+						referenceLoadInfo.AddMessageOnce(reference.FullName, MessageKind.Info, "Success - Found in Assembly List with different TFM or version: " + module.FileName);
 					}
-
-					candidates.SortBy(c => c.version);
-
-					var bestCandidate = candidates.FirstOrDefault(c => c.version >= reference.Version).assembly ?? candidates.Last().assembly;
-					referenceLoadInfo.AddMessageOnce(reference.ToString(), MessageKind.Info, "Success - Found in Assembly List with different TFM or version: " + bestCandidate.fileName);
-					return await bestCandidate.GetPEFileOrNullAsync().ConfigureAwait(false);
+					return module;
 				}
 			}