Fix assembly rewriting causing VS to error/crash when debugging

3 files changed, 79 insertions, 14 deletions

diff --git a/src/SMAPI/Framework/ModLoading/AssemblyLoader.cs b/src/SMAPI/Framework/ModLoading/AssemblyLoader.cs
index 86b43990..98154b53 100644
--- a/src/SMAPI/Framework/ModLoading/AssemblyLoader.cs
+++ b/src/SMAPI/Framework/ModLoading/AssemblyLoader.cs
@@ -34,6 +34,9 @@ namespace StardewModdingAPI.Framework.ModLoading
         /// <summary>A minimal assembly definition resolver which resolves references to known loaded assemblies.</summary>
         private readonly AssemblyDefinitionResolver AssemblyDefinitionResolver;
 
+        private readonly SymbolReaderProvider SymbolReaderProvider;
+        private readonly SymbolWriterProvider SymbolWriterProvider;
+
         /// <summary>The objects to dispose as part of this instance.</summary>
         private readonly HashSet<IDisposable> Disposables = new HashSet<IDisposable>();
 
@@ -61,6 +64,9 @@ namespace StardewModdingAPI.Framework.ModLoading
             this.AssemblyDefinitionResolver = this.TrackForDisposal(new AssemblyDefinitionResolver());
             Constants.ConfigureAssemblyResolver(this.AssemblyDefinitionResolver);
 
+            this.SymbolReaderProvider = new SymbolReaderProvider();
+            this.SymbolWriterProvider = new SymbolWriterProvider();
+
             // generate type => assembly lookup for types which should be rewritten
             this.TypeAssemblies = new Dictionary<string, Assembly>();
             foreach (Assembly assembly in this.AssemblyMap.Targets)
@@ -114,7 +120,7 @@ namespace StardewModdingAPI.Framework.ModLoading
 
                 // rewrite assembly
                 bool changed = this.RewriteAssembly(mod, assembly.Definition, loggedMessages, logPrefix: "      ");
-
+                
                 // detect broken assembly reference
                 foreach (AssemblyNameReference reference in assembly.Definition.MainModule.AssemblyReferences)
                 {
@@ -134,20 +140,12 @@ namespace StardewModdingAPI.Framework.ModLoading
                     if (!oneAssembly)
                         this.Monitor.Log($"      Loading {assembly.File.Name} (rewritten)...", LogLevel.Trace);
 
-                    // load PDB file if present
-                    byte[] symbols;
-                    {
-                        string symbolsPath = Path.Combine(Path.GetDirectoryName(assemblyPath), Path.GetFileNameWithoutExtension(assemblyPath)) + ".pdb";
-                        symbols = File.Exists(symbolsPath)
-                            ? File.ReadAllBytes(symbolsPath)
-                            : null;
-                    }
-
                     // load assembly
                     using MemoryStream outStream = new MemoryStream();
-                    assembly.Definition.Write(outStream);
+                    using MemoryStream outSymbolStream = new MemoryStream();
+                    assembly.Definition.Write(outStream, new WriterParameters() { WriteSymbols = true, SymbolStream = outSymbolStream, SymbolWriterProvider = this.SymbolWriterProvider } );
                     byte[] bytes = outStream.ToArray();
-                    lastAssembly = Assembly.Load(bytes, symbols);
+                    lastAssembly = Assembly.Load(bytes, outSymbolStream.ToArray());
                 }
                 else
                 {
@@ -234,10 +232,15 @@ namespace StardewModdingAPI.Framework.ModLoading
             if (!file.Exists)
                 yield break; // not a local assembly
 
-            // read assembly
+            // read assembly and PDB (if present)
             byte[] assemblyBytes = File.ReadAllBytes(file.FullName);
             Stream readStream = this.TrackForDisposal(new MemoryStream(assemblyBytes));
-            AssemblyDefinition assembly = this.TrackForDisposal(AssemblyDefinition.ReadAssembly(readStream, new ReaderParameters(ReadingMode.Immediate) { AssemblyResolver = assemblyResolver, InMemory = true }));
+            {
+                string symbolsPath = Path.Combine(Path.GetDirectoryName(file.FullName), Path.GetFileNameWithoutExtension(file.FullName)) + ".pdb";
+                if ( File.Exists( symbolsPath ) )
+                    this.SymbolReaderProvider.AddSymbolMapping( Path.GetFileName( file.FullName ), this.TrackForDisposal( new MemoryStream( File.ReadAllBytes( symbolsPath ) ) ) );
+            }
+            AssemblyDefinition assembly = this.TrackForDisposal(AssemblyDefinition.ReadAssembly(readStream, new ReaderParameters(ReadingMode.Immediate) { AssemblyResolver = assemblyResolver, InMemory = true, ReadSymbols = true, SymbolReaderProvider = this.SymbolReaderProvider }));
 
             // skip if already visited
             if (visitedAssemblyNames.Contains(assembly.Name.Name))
diff --git a/src/SMAPI/Framework/ModLoading/SymbolReaderProvider.cs b/src/SMAPI/Framework/ModLoading/SymbolReaderProvider.cs
new file mode 100644
index 00000000..a651c167
--- /dev/null
+++ b/src/SMAPI/Framework/ModLoading/SymbolReaderProvider.cs
@@ -0,0 +1,37 @@
+using System;
+using System.Collections.Generic;
+using System.IO;
+using Mono.Cecil;
+using Mono.Cecil.Cil;
+using Mono.Cecil.Pdb;
+
+namespace StardewModdingAPI.Framework.ModLoading
+{
+    internal class SymbolReaderProvider : ISymbolReaderProvider
+    {
+        private readonly ISymbolReaderProvider BaseProvider = new DefaultSymbolReaderProvider();
+
+        private readonly Dictionary<string, Stream> SymbolMapping = new Dictionary<string, Stream>();
+
+        public void AddSymbolMapping( string dllName, Stream symbolStream )
+        {
+            this.SymbolMapping.Add( dllName, symbolStream );
+        }
+
+        public ISymbolReader GetSymbolReader( ModuleDefinition module, string fileName )
+        {
+            if ( this.SymbolMapping.ContainsKey( module.Name ) )
+                return new NativePdbReaderProvider().GetSymbolReader( module, this.SymbolMapping[ module.Name ] );
+            
+            return this.BaseProvider.GetSymbolReader( module, fileName );
+        }
+
+        public ISymbolReader GetSymbolReader( ModuleDefinition module, Stream symbolStream )
+        {
+            if ( this.SymbolMapping.ContainsKey( module.Name ) )
+                return new PortablePdbReaderProvider().GetSymbolReader( module, this.SymbolMapping[ module.Name ] );
+
+            return this.BaseProvider.GetSymbolReader( module, symbolStream );
+        }
+    }
+}
diff --git a/src/SMAPI/Framework/ModLoading/SymbolWriterProvider.cs b/src/SMAPI/Framework/ModLoading/SymbolWriterProvider.cs
new file mode 100644
index 00000000..116e341a
--- /dev/null
+++ b/src/SMAPI/Framework/ModLoading/SymbolWriterProvider.cs
@@ -0,0 +1,25 @@
+using System;
+using System.Collections.Generic;
+using System.IO;
+using Mono.Cecil;
+using Mono.Cecil.Cil;
+using Mono.Cecil.Pdb;
+
+namespace StardewModdingAPI.Framework.ModLoading
+{
+    internal class SymbolWriterProvider : ISymbolWriterProvider
+    {
+        private readonly ISymbolWriterProvider BaseProvider = new DefaultSymbolWriterProvider();
+
+        public ISymbolWriter GetSymbolWriter( ModuleDefinition module, string fileName )
+        {
+            return this.BaseProvider.GetSymbolWriter( module, fileName );
+        }
+
+        public ISymbolWriter GetSymbolWriter( ModuleDefinition module, Stream symbolStream )
+        {
+            // Not implemented in default native pdb writer, so fallback to portable
+            return new PortablePdbWriterProvider().GetSymbolWriter( module, symbolStream );
+        }
+    }
+}