Speed up publicization by patching metadata in place

src/Publicizer.Tests/InPlaceWriterTests.cs

@@ -0,0 +1,105 @@
+using dnlib.DotNet;
+using dnlib.DotNet.Writer;
+using NUnit.Framework;
+
+namespace Publicizer.Tests;
+
+/// <summary>
+/// <see cref="InPlaceWriter"/> replaces a full dnlib metadata rebuild with a byte patch, so it has to be
+/// indistinguishable from the dnlib writer in the only respect that matters — the accessibility of every
+/// member — while leaving the rest of the file untouched.
+/// </summary>
+internal static class InPlaceWriterTests
+{
+    private static PublicizerAssemblyContext WholeAssembly() => new("Fixture") { ExplicitlyPublicizeAssembly = true };
+
+    private static PublicizerAssemblyContext TargetedMember()
+    {
+        var context = new PublicizerAssemblyContext("Fixture");
+        context.PublicizeMemberPatterns.Add("Fixture.Shapes.PrivateField");
+        return context;
+    }
+
+    private static string WriteInPlace(TemporaryFolder folder, PublicizerAssemblyContext context)
+    {
+        string destination = Path.Combine(folder.Path, "inplace.dll");
+        using var module = ModuleDefMD.Load(Fixtures.ShapesPath());
+        PublicizeAssemblies.PublicizeAssembly(module, context, NullTaskLogger.Instance);
+
+        bool written = InPlaceWriter.TryWrite(module, Fixtures.ShapesPath(), destination, NullTaskLogger.Instance);
+
+        Assert.That(written, Is.True, "the fixture has an ordinary compressed-metadata layout, so the fast path must apply");
+        return destination;
+    }
+
+    private static string WriteWithDnlib(TemporaryFolder folder, PublicizerAssemblyContext context)
+    {
+        string destination = Path.Combine(folder.Path, "dnlib.dll");
+        using var module = ModuleDefMD.Load(Fixtures.ShapesPath());
+        PublicizeAssemblies.PublicizeAssembly(module, context, NullTaskLogger.Instance);
+
+        using var stream = new FileStream(destination, FileMode.Create, FileAccess.ReadWrite, FileShare.Read);
+        module.Write(stream, new ModuleWriterOptions(module)
+        {
+            MetadataOptions = new MetadataOptions(MetadataFlags.KeepOldMaxStack),
+            Logger = DummyLogger.NoThrowInstance,
+        });
+        return destination;
+    }
+
+    private static string ManifestOf(string assemblyPath)
+    {
+        using var module = ModuleDefMD.Load(assemblyPath);
+        return AccessibilityManifest.Of(module);
+    }
+
+    [Test]
+    public static void WholeAssembly_ProducesSameAccessibilityAsDnlibWriter()
+    {
+        using var folder = new TemporaryFolder();
+
+        string patched = WriteInPlace(folder, WholeAssembly());
+        string reference = WriteWithDnlib(folder, WholeAssembly());
+
+        Assert.That(ManifestOf(patched), Is.EqualTo(ManifestOf(reference)));
+    }
+
+    [Test]
+    public static void TargetedMember_ProducesSameAccessibilityAsDnlibWriter()
+    {
+        using var folder = new TemporaryFolder();
+
+        string patched = WriteInPlace(folder, TargetedMember());
+        string reference = WriteWithDnlib(folder, TargetedMember());
+
+        Assert.That(ManifestOf(patched), Is.EqualTo(ManifestOf(reference)));
+    }
+
+    [Test]
+    public static void Output_DiffersFromInputOnlyInFlagBytes()
+    {
+        using var folder = new TemporaryFolder();
+        byte[] original = File.ReadAllBytes(Fixtures.ShapesPath());
+
+        byte[] patched = File.ReadAllBytes(WriteInPlace(folder, WholeAssembly()));
+
+        Assert.That(patched, Has.Length.EqualTo(original.Length));
+
+        int differing = original.Where((b, i) => b != patched[i]).Count();
+        Assert.That(differing, Is.GreaterThan(0), "publicizing the fixture must change something");
+        // Only the Flags columns of the TypeDef/Field/Method rows may move; anything larger means the
+        // patch is straying outside the metadata tables it is supposed to touch.
+        Assert.That(differing, Is.LessThan(original.Length / 100), "patch touched far more of the file than the flag columns");
+    }
+
+    [Test]
+    public static void Output_IsStillLoadable()
+    {
+        using var folder = new TemporaryFolder();
+
+        string patched = WriteInPlace(folder, WholeAssembly());
+
+        using var module = ModuleDefMD.Load(patched);
+        Assert.That(module.Find("Fixture.Shapes", isReflectionName: true).Fields.Single(f => f.Name == "PrivateField").IsPublic, Is.True);
+    }
+}

src/Publicizer/InPlaceWriter.cs

@@ -0,0 +1,126 @@
+using dnlib.DotNet;
+using dnlib.DotNet.MD;
+
+namespace Publicizer;
+
+/// <summary>
+/// Writes a publicized assembly by patching visibility bits directly in a copy of the original file.
+/// </summary>
+/// <remarks>
+/// Publicization only ever flips bits in the Flags column of the TypeDef, Field and Method tables
+/// (see <see cref="AssemblyEditor"/>) — nothing is added, removed, renamed or resized. Those columns are
+/// fixed-width at fixed offsets, so the result can be produced by patching bytes instead of having dnlib
+/// tear down and rebuild the entire metadata, which is two to three orders of magnitude more expensive.
+///
+/// Patching also leaves every other byte identical to the input, so output is stable across dnlib upgrades
+/// and unaffected by the writer quirk that KeepOldMaxStack works around.
+///
+/// Layouts this cannot handle are rejected by <see cref="TryWrite"/> so the caller can fall back to the
+/// dnlib writer.
+/// </remarks>
+internal static class InPlaceWriter
+{
+    // ECMA-335 II.22.37/15/26: index of the Flags column within each table's row.
+    private const int TypeDefFlagsColumnIndex = 0;
+    private const int FieldFlagsColumnIndex = 0;
+    private const int MethodFlagsColumnIndex = 2;
+
+    private const string FlagsColumnName = "Flags";
+
+    /// <summary>
+    /// Attempts to write <paramref name="module"/>'s publicized form to <paramref name="destinationPath"/> by
+    /// patching a copy of <paramref name="sourcePath"/>. Returns false when the assembly's metadata layout is
+    /// not patchable, in which case nothing has been written and the caller should use the dnlib writer.
+    /// </summary>
+    internal static bool TryWrite(ModuleDefMD module, string sourcePath, string destinationPath, ITaskLogger logger)
+    {
+        Metadata metadata = module.Metadata;
+
+        // ENC/uncompressed metadata (#- heap) allows deleted rows and non-sequential rids, so row offsets
+        // cannot be computed from rid alone.
+        if (!metadata.IsCompressed)
+        {
+            logger.Info("Metadata is not compressed (#- heap); falling back to the dnlib writer");
+            return false;
+        }
+
+        TablesStream tables = metadata.TablesStream;
+        byte[] buffer = File.ReadAllBytes(sourcePath);
+
+        if (!TryPatchTable(buffer, tables.TypeDefTable, TypeDefFlagsColumnIndex, module.ResolveTypeDefFlags, logger) ||
+            !TryPatchTable(buffer, tables.FieldTable, FieldFlagsColumnIndex, module.ResolveFieldFlags, logger) ||
+            !TryPatchTable(buffer, tables.MethodTable, MethodFlagsColumnIndex, module.ResolveMethodFlags, logger))
+        {
+            return false;
+        }
+
+        File.WriteAllBytes(destinationPath, buffer);
+        return true;
+    }
+
+    private static bool TryPatchTable(byte[] buffer, MDTable table, int columnIndex, Func<uint, uint?> getFlags, ITaskLogger logger)
+    {
+        if (table is null || table.Rows == 0)
+        {
+            return true;
+        }
+
+        if (columnIndex >= table.Columns.Count)
+        {
+            logger.Info($"Table {table.Name} has no column {columnIndex}; falling back to the dnlib writer");
+            return false;
+        }
+
+        ColumnInfo column = table.Columns[columnIndex];
+
+        // Guards against dnlib ever reordering or resizing the column out from under these constants.
+        if (!string.Equals(column.Name, FlagsColumnName, StringComparison.Ordinal))
+        {
+            logger.Info($"Table {table.Name} column {columnIndex} is '{column.Name}', not '{FlagsColumnName}'; falling back to the dnlib writer");
+            return false;
+        }
+
+        if (column.Size is not (2 or 4))
+        {
+            logger.Info($"Table {table.Name} flags column is {column.Size} bytes; falling back to the dnlib writer");
+            return false;
+        }
+
+        long tableStart = (long)table.StartOffset;
+        long rowSize = table.RowSize;
+        long lastByte = tableStart + ((table.Rows - 1) * rowSize) + column.Offset + column.Size;
+
+        if (tableStart < 0 || lastByte > buffer.Length)
+        {
+            logger.Info($"Table {table.Name} extends past the end of the file; falling back to the dnlib writer");
+            return false;
+        }
+
+        for (uint rid = 1; rid <= table.Rows; rid++)
+        {
+            uint? flags = getFlags(rid);
+            if (flags is null)
+            {
+                logger.Info($"Table {table.Name} row {rid} did not resolve; falling back to the dnlib writer");
+                return false;
+            }
+
+            long offset = tableStart + ((rid - 1) * rowSize) + column.Offset;
+            uint value = flags.Value;
+
+            buffer[offset] = (byte)value;
+            buffer[offset + 1] = (byte)(value >> 8);
+            if (column.Size == 4)
+            {
+                buffer[offset + 2] = (byte)(value >> 16);
+                buffer[offset + 3] = (byte)(value >> 24);
+            }
+        }
+
+        return true;
+    }
+
+    private static uint? ResolveTypeDefFlags(this ModuleDefMD module, uint rid) => module.ResolveTypeDef(rid) is TypeDef type ? (uint)type.Attributes : null;
+    private static uint? ResolveFieldFlags(this ModuleDefMD module, uint rid) => module.ResolveField(rid) is FieldDef field ? (uint)field.Attributes : null;
+    private static uint? ResolveMethodFlags(this ModuleDefMD module, uint rid) => module.ResolveMethod(rid) is MethodDef method ? (uint)method.Attributes : null;
+}

src/Publicizer/PublicizeAssemblies.cs

@@ -107,7 +107,7 @@ public override bool Execute()
             }
             else
             {
-                using ModuleDef module = ModuleDefMD.Load(assemblyPath);
+                using var module = ModuleDefMD.Load(assemblyPath);
                 scopedLogger.Info("Publicizing members...");
                 bool isAssemblyModified = PublicizeAssembly(module, assemblyContext, scopedLogger);
                 if (!isAssemblyModified)
@@ -116,17 +116,21 @@ public override bool Execute()
                     continue;
                 }
 
-                using var fileStream = new FileStream(outputAssemblyPath, FileMode.Create, FileAccess.ReadWrite, FileShare.Read);
+                scopedLogger.Info($"Saving publicized assembly to {outputAssemblyPath}");
 
-                var writerOptions = new ModuleWriterOptions(module)
+                if (!InPlaceWriter.TryWrite(module, assemblyPath, outputAssemblyPath, scopedLogger))
                 {
-                    // Writing the module sometime fails without this flag due to how it was originally compiled.
-                    // https://github.com/krafs/Publicizer/issues/42
-                    MetadataOptions = new MetadataOptions(MetadataFlags.KeepOldMaxStack),
-                    Logger = DummyLogger.NoThrowInstance
-                };
-                scopedLogger.Info($"Saving publicized assembly to {outputAssemblyPath}");
-                module.Write(fileStream, writerOptions);
+                    using var fileStream = new FileStream(outputAssemblyPath, FileMode.Create, FileAccess.ReadWrite, FileShare.Read);
+
+                    var writerOptions = new ModuleWriterOptions(module)
+                    {
+                        // Writing the module sometime fails without this flag due to how it was originally compiled.
+                        // https://github.com/krafs/Publicizer/issues/42
+                        MetadataOptions = new MetadataOptions(MetadataFlags.KeepOldMaxStack),
+                        Logger = DummyLogger.NoThrowInstance
+                    };
+                    module.Write(fileStream, writerOptions);
+                }
 
                 string assemblyDirectory = Path.GetDirectoryName(assemblyPath);
                 string originalDocumentationFullPath = Path.Combine(assemblyDirectory, assemblyName + ".xml");