From cc87b00109f05450e36ce0995b86ebf094ebfdc8 Mon Sep 17 00:00:00 2001
From: Andrew Au <cshung@gmail.com>
Date: Sun, 2 Apr 2023 09:29:52 -0700
Subject: [PATCH] Support disassembling ReadyToRun binaries compiled using
 composite mode

---
 ILSpy.ReadyToRun/ILSpy.ReadyToRun.csproj   |  2 +-
 ILSpy.ReadyToRun/ReadyToRunDisassembler.cs |  9 +--
 ILSpy.ReadyToRun/ReadyToRunLanguage.cs     | 84 ++++++++++++++++++----
 NuGet.config                               |  2 +-
 4 files changed, 74 insertions(+), 23 deletions(-)

diff --git a/ILSpy.ReadyToRun/ILSpy.ReadyToRun.csproj b/ILSpy.ReadyToRun/ILSpy.ReadyToRun.csproj
index eccbd00e53..804c7c8e60 100644
--- a/ILSpy.ReadyToRun/ILSpy.ReadyToRun.csproj
+++ b/ILSpy.ReadyToRun/ILSpy.ReadyToRun.csproj
@@ -38,7 +38,7 @@
 
   <ItemGroup>
     <PackageReference Include="Iced" Version="1.18.0" />
-    <PackageReference Include="ILCompiler.Reflection.ReadyToRun.Experimental" Version="7.0.4-servicing.23115.8" />
+    <PackageReference Include="ILCompiler.Reflection.ReadyToRun.Experimental" Version="8.0.0-preview.4.23202.2" />
     <!-- ILCompiler.Reflection.ReadyToRun has dependencies on System.Reflection.Metadata and
          System.Runtime.CompilerServices.Unsafe. Because the AddIn compiles into ILSpy's output
          directory, we're at risk of overwriting our dependencies with different versions.
diff --git a/ILSpy.ReadyToRun/ReadyToRunDisassembler.cs b/ILSpy.ReadyToRun/ReadyToRunDisassembler.cs
index 6cf7b06488..6996f211db 100644
--- a/ILSpy.ReadyToRun/ReadyToRunDisassembler.cs
+++ b/ILSpy.ReadyToRun/ReadyToRunDisassembler.cs
@@ -137,8 +137,8 @@ public void Disassemble(PEFile currentFile, int bitness, ulong address, bool sho
 				output.Write(tempOutput.ToStringAndReset());
 				DecorateUnwindInfo(unwindInfo, baseInstrIP, instr);
 				DecorateDebugInfo(instr, debugInfo, baseInstrIP);
-
 				DecorateCallSite(currentFile, showMetadataTokens, showMetadataTokensInBase10, instr);
+				output.WriteLine();
 			}
 			output.WriteLine();
 		}
@@ -451,16 +451,11 @@ private void DecorateCallSite(PEFile currentFile, bool showMetadataTokens, bool
 							methodRefToken.WriteTo(currentFile, output, default);
 							break;
 						default:
-							output.WriteLine(reader.ImportSignatures[importCellAddress].ToString(new SignatureFormattingOptions()));
+							output.Write(reader.ImportSignatures[importCellAddress].ToString(new SignatureFormattingOptions()));
 							break;
 					}
-					output.WriteLine();
 				}
 			}
-			else
-			{
-				output.WriteLine();
-			}
 		}
 	}
 }
\ No newline at end of file
diff --git a/ILSpy.ReadyToRun/ReadyToRunLanguage.cs b/ILSpy.ReadyToRun/ReadyToRunLanguage.cs
index ec3ced0bc2..95b0f3c765 100644
--- a/ILSpy.ReadyToRun/ReadyToRunLanguage.cs
+++ b/ILSpy.ReadyToRun/ReadyToRunLanguage.cs
@@ -22,6 +22,7 @@
 using System.Collections.Generic;
 using System.ComponentModel.Composition;
 using System.Diagnostics;
+using System.IO;
 using System.Linq;
 using System.Reflection.Metadata;
 using System.Reflection.PortableExecutable;
@@ -121,10 +122,13 @@ public override ProjectId DecompileAssembly(LoadedAssembly assembly, ITextOutput
 			else
 			{
 				ReadyToRunReader reader = cacheEntry.readyToRunReader;
-				WriteCommentLine(output, reader.Machine.ToString());
-				WriteCommentLine(output, reader.OperatingSystem.ToString());
-				WriteCommentLine(output, reader.CompilerIdentifier);
-				WriteCommentLine(output, "TODO - display more header information");
+				WriteCommentLine(output, $"Machine                  : {reader.Machine}");
+				WriteCommentLine(output, $"OperatingSystem          : {reader.OperatingSystem}");
+				WriteCommentLine(output, $"CompilerIdentifier       : {reader.CompilerIdentifier}");
+				if (reader.OwnerCompositeExecutable != null)
+				{
+					WriteCommentLine(output, $"OwnerCompositeExecutable : {reader.OwnerCompositeExecutable}");
+				}
 			}
 
 			return base.DecompileAssembly(assembly, output, options);
@@ -153,9 +157,22 @@ public override void DecompileMethod(IMethod method, ITextOutput output, Decompi
 				}
 				if (cacheEntry.methodMap == null)
 				{
-					cacheEntry.methodMap = reader.Methods.ToList()
-						.GroupBy(m => m.MethodHandle)
-						.ToDictionary(g => g.Key, g => g.ToArray());
+					IEnumerable<ReadyToRunMethod> readyToRunMethods = null;
+					if (cacheEntry.compositeReadyToRunReader == null)
+					{
+						readyToRunMethods = reader.Methods;
+					}
+					else
+					{
+						readyToRunMethods = cacheEntry.compositeReadyToRunReader.Methods
+							.Where(m => {
+								MetadataReader mr = m.ComponentReader.MetadataReader;
+								return string.Equals(mr.GetString(mr.GetAssemblyDefinition().Name), method.ParentModule.Name, StringComparison.OrdinalIgnoreCase);
+							});
+					}
+					cacheEntry.methodMap = readyToRunMethods.ToList()
+							.GroupBy(m => m.MethodHandle)
+							.ToDictionary(g => g.Key, g => g.ToArray());
 				}
 				var displaySettings = MainWindow.Instance.CurrentDisplaySettings;
 				bool showMetadataTokens = displaySettings.ShowMetadataTokens;
@@ -174,7 +191,20 @@ public override void DecompileMethod(IMethod method, ITextOutput output, Decompi
 #endif
 						foreach (RuntimeFunction runtimeFunction in readyToRunMethod.RuntimeFunctions)
 						{
-							new ReadyToRunDisassembler(output, reader, runtimeFunction).Disassemble(method.ParentModule.PEFile, bitness, (ulong)runtimeFunction.StartAddress, showMetadataTokens, showMetadataTokensInBase10);
+							PEFile file = null;
+							ReadyToRunReader disassemblingReader = null;
+							if (cacheEntry.compositeReadyToRunReader == null)
+							{
+								disassemblingReader = reader;
+								file = method.ParentModule.PEFile;
+							}
+							else
+							{
+								disassemblingReader = cacheEntry.compositeReadyToRunReader;
+								file = ((IlSpyAssemblyMetadata)readyToRunMethod.ComponentReader).Module;
+							}
+
+							new ReadyToRunDisassembler(output, disassemblingReader, runtimeFunction).Disassemble(file, bitness, (ulong)runtimeFunction.StartAddress, showMetadataTokens, showMetadataTokensInBase10);
 						}
 					}
 				}
@@ -206,6 +236,11 @@ private ReadyToRunReaderCacheEntry GetReader(LoadedAssembly assembly, PEFile mod
 							result.failureReason = $"Architecture {result.readyToRunReader.Machine} is not currently supported.";
 							result.readyToRunReader = null;
 						}
+						else if (result.readyToRunReader.OwnerCompositeExecutable != null)
+						{
+							string compositePath = Path.Combine(Path.GetDirectoryName(module.FileName), result.readyToRunReader.OwnerCompositeExecutable);
+							result.compositeReadyToRunReader = new ReadyToRunReader(new ReadyToRunAssemblyResolver(assembly), compositePath);
+						}
 					}
 					catch (BadImageFormatException e)
 					{
@@ -219,31 +254,52 @@ private ReadyToRunReaderCacheEntry GetReader(LoadedAssembly assembly, PEFile mod
 
 		private class ReadyToRunAssemblyResolver : ILCompiler.Reflection.ReadyToRun.IAssemblyResolver
 		{
+			private LoadedAssembly loadedAssembly;
 			private Decompiler.Metadata.IAssemblyResolver assemblyResolver;
 
 			public ReadyToRunAssemblyResolver(LoadedAssembly loadedAssembly)
 			{
+				this.loadedAssembly = loadedAssembly;
 				assemblyResolver = loadedAssembly.GetAssemblyResolver();
 			}
 
 			public IAssemblyMetadata FindAssembly(MetadataReader metadataReader, AssemblyReferenceHandle assemblyReferenceHandle, string parentFile)
 			{
-				PEFile module = assemblyResolver.Resolve(new Decompiler.Metadata.AssemblyReference(metadataReader, assemblyReferenceHandle));
-				PEReader reader = module?.Reader;
-				return reader == null ? null : new StandaloneAssemblyMetadata(reader);
+				return GetAssemblyMetadata(assemblyResolver.Resolve(new Decompiler.Metadata.AssemblyReference(metadataReader, assemblyReferenceHandle)));
 			}
 
 			public IAssemblyMetadata FindAssembly(string simpleName, string parentFile)
 			{
-				// This is called only for the composite R2R scenario, 
-				// So it will never be called before the feature is released.
-				throw new NotSupportedException("Composite R2R format is not currently supported");
+				return GetAssemblyMetadata(assemblyResolver.ResolveModule(loadedAssembly.GetPEFileOrNull(), simpleName + ".dll"));
+			}
+
+			private IAssemblyMetadata GetAssemblyMetadata(PEFile module)
+			{
+				if (module.Reader == null)
+				{
+					return null;
+				}
+				else
+				{
+					return new IlSpyAssemblyMetadata(module);
+				}
+			}
+		}
+
+		private class IlSpyAssemblyMetadata : StandaloneAssemblyMetadata
+		{
+			public PEFile Module { get; private set; }
+
+			public IlSpyAssemblyMetadata(PEFile module) : base(module.Reader)
+			{
+				Module = module;
 			}
 		}
 
 		private class ReadyToRunReaderCacheEntry
 		{
 			public ReadyToRunReader readyToRunReader;
+			public ReadyToRunReader compositeReadyToRunReader;
 			public string failureReason;
 			public Dictionary<EntityHandle, ReadyToRunMethod[]> methodMap;
 		}
diff --git a/NuGet.config b/NuGet.config
index e6def65d28..40b60b1652 100644
--- a/NuGet.config
+++ b/NuGet.config
@@ -3,6 +3,6 @@
   <packageSources>
     <add key="Nuget Official" value="https://api.nuget.org/v3/index.json" />
     <add key="dotnet-tools" value="https://pkgs.dev.azure.com/dnceng/public/_packaging/dotnet-tools/nuget/v3/index.json" />
-    <add key="dotnet7-transport" value="https://pkgs.dev.azure.com/dnceng/public/_packaging/dotnet7-transport/nuget/v3/index.json" />
+    <add key="dotnet8-transport" value="https://pkgs.dev.azure.com/dnceng/public/_packaging/dotnet8-transport/nuget/v3/index.json" />
   </packageSources>
 </configuration>