Support /FOLD in managed ilasm

src/tools/ilasm/KNOWN-ISSUES.md

@@ -21,10 +21,6 @@ managed ilasm.
 
 Currently managed ilasm does not do any optimizations of the IL written by the user.
 
-## -FOLD is a no-op
-
-Currently managed ilasm does not fold identical IL bodies from different methods into the same blob.
-
 ## ARM32 target images are not supported
 
 Managed ilasm does not support generating ARM32 (AArch32) machine images. The native `/ARM`

src/tools/ilasm/src/ILAssembler/Actions/GrammarActions.BuildImage.cs

@@ -66,7 +66,7 @@ internal sealed partial class GrammarActions
             }
 
             BlobBuilder ilStream = new();
-            Blob mvidFixup = _entityRegistry.WriteContentTo(_metadataBuilder, ilStream, _mappedFieldDataNames, _options.Deterministic);
+            Blob mvidFixup = _entityRegistry.WriteContentTo(_metadataBuilder, ilStream, _mappedFieldDataNames, _options.Deterministic, _options.Fold);
             // MetadataRootBuilder only supports module-wide validation suppression, which is
             // required because wrapped GenericParam numbers intentionally violate table ordering.
             bool suppressMetadataValidation =

src/tools/ilasm/src/ILAssembler/EntityRegistry.cs

@@ -55,6 +55,19 @@ public int GetHashCode(BlobBuilder obj)
             }
         }
 
+        private sealed class MethodBodyContentEqualityComparer : IEqualityComparer<byte[]>
+        {
+            public bool Equals(byte[]? x, byte[]? y) =>
+                x is null ? y is null : y is not null && x.AsSpan().SequenceEqual(y);
+
+            public int GetHashCode(byte[] obj)
+            {
+                HashCode hash = default;
+                hash.AddBytes(obj);
+                return hash.ToHashCode();
+            }
+        }
+
         private sealed class MethodSpecEqualityComparer : IEqualityComparer<(EntityBase, BlobBuilder)>
         {
             public bool Equals((EntityBase, BlobBuilder) x, (EntityBase, BlobBuilder) y)
@@ -92,7 +105,7 @@ public IReadOnlyList<EntityBase> GetSeenEntities(TableIndex table)
             return Array.Empty<EntityBase>();
         }
 
-        public Blob WriteContentTo(MetadataBuilder builder, BlobBuilder ilStream, IReadOnlyDictionary<string, int> mappedFieldDataNames, bool deterministic)
+        public Blob WriteContentTo(MetadataBuilder builder, BlobBuilder ilStream, IReadOnlyDictionary<string, int> mappedFieldDataNames, bool deterministic, bool fold)
         {
             // Set the assembly handle early since DeclarativeSecurityAttribute needs it
             // The assembly definition handle is always row 1 (there's only ever one assembly per module)
@@ -342,6 +355,7 @@ public Blob WriteContentTo(MetadataBuilder builder, BlobBuilder ilStream, IReadO
             }
 
             var bodyStreamEncoder = new MethodBodyStreamEncoder(ilStream);
+            Dictionary<byte[], int>? foldedBodies = fold ? new(new MethodBodyContentEqualityComparer()) : null;
 
             for (int i = 0; i < GetSeenEntities(TableIndex.MethodDef).Count; i++)
             {
@@ -350,6 +364,9 @@ public Blob WriteContentTo(MetadataBuilder builder, BlobBuilder ilStream, IReadO
                 int bodyOffset = -1;
                 if (methodDef.MethodBody.CodeBuilder.Count != 0)
                 {
+                    BlobBuilder? serializedBody = fold ? new BlobBuilder() : null;
+                    MethodBodyStreamEncoder encoder = fold ? new(serializedBody!) : bodyStreamEncoder;
+
                     // Add deferred exception regions now that TypeRef-to-TypeDef resolution is complete.
                     // Catch clause type handles are read here, after resolution has set the real handle.
                     foreach (var region in methodDef.ExceptionRegions)
@@ -394,7 +411,7 @@ public Blob WriteContentTo(MetadataBuilder builder, BlobBuilder ilStream, IReadO
 
                     try
                     {
-                        bodyOffset = bodyStreamEncoder.AddMethodBody(
+                        bodyOffset = encoder.AddMethodBody(
                             methodDef.MethodBody,
                             methodDef.MaxStack,
                             localsSigHandle,
@@ -406,7 +423,7 @@ public Blob WriteContentTo(MetadataBuilder builder, BlobBuilder ilStream, IReadO
                         // Method has unresolved labels or other body errors.
                         // Emit a minimal valid method body containing the raw IL bytes so
                         // the PE can still be emitted (error diagnostics are already recorded).
-                        var fallbackBody = bodyStreamEncoder.AddMethodBody(
+                        var fallbackBody = encoder.AddMethodBody(
                             methodDef.MethodBody.CodeBuilder.Count,
                             methodDef.MaxStack,
                             exceptionRegionCount: 0,
@@ -424,7 +441,7 @@ public Blob WriteContentTo(MetadataBuilder builder, BlobBuilder ilStream, IReadO
                         // errors that produced malformed control flow). Emit the IL in a
                         // minimal valid method body and omit exception regions in fallback.
                         // TODO-COMPAT: Emit the invalid exception regions manually
-                        var fallbackBody = bodyStreamEncoder.AddMethodBody(
+                        var fallbackBody = encoder.AddMethodBody(
                             methodDef.MethodBody.CodeBuilder.Count,
                             methodDef.MaxStack,
                             exceptionRegionCount: 0,
@@ -436,6 +453,31 @@ public Blob WriteContentTo(MetadataBuilder builder, BlobBuilder ilStream, IReadO
                         var writer2 = new BlobWriter(fallbackBody.Instructions);
                         methodDef.MethodBody.CodeBuilder.WriteContentTo(ref writer2);
                     }
+
+                    if (fold)
+                    {
+                        byte[] content = serializedBody!.ToArray();
+                        // The encoder may have written a partial body before falling back.
+                        if (bodyOffset != 0)
+                        {
+                            content = content.AsSpan(bodyOffset).ToArray();
+                        }
+                        if (foldedBodies!.TryGetValue(content, out int existingOffset))
+                        {
+                            bodyOffset = existingOffset;
+                        }
+                        else
+                        {
+                            // Fat method headers must be aligned relative to the IL stream.
+                            if ((content[0] & 0x3) == 0x3)
+                            {
+                                ilStream.Align(4);
+                            }
+                            bodyOffset = ilStream.Count;
+                            ilStream.WriteBytes(content);
+                            foldedBodies.Add(content, bodyOffset);
+                        }
+                    }
                 }
 
                 var methodAttributes = methodDef.MethodAttributes;

src/tools/ilasm/src/ILAssembler/Options.cs

@@ -135,7 +135,6 @@ public sealed class Options
         /// <summary>
         /// Fold identical method bodies into one.
         /// </summary>
-        /// <remarks>Not yet implemented; accepted for CLI compatibility with native ilasm.</remarks>
         public bool Fold { get; set; }
 
         /// <summary>

src/tools/ilasm/tests/ILAssembler.Tests/CompilerOptionsTests.cs

@@ -194,6 +194,90 @@ .method public static void Main() cil managed
             Assert.NotEqual(0, pe.PEHeaders.PEHeader!.SizeOfImage);
         }
 
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public void FoldOption_SharesOnlyIdenticalMethodBodies(bool fold)
+        {
+            string source = """
+                .assembly test { }
+                .class public auto ansi Test
+                {
+                    .method public static void First() cil managed { ret }
+                    .method public static void Second() cil managed { ret }
+                    .method public static void DifferentCode() cil managed { nop ret }
+                    .method public static void DifferentStack() cil managed
+                    {
+                        .maxstack 1
+                        ret
+                    }
+                    .method public static void FirstLocals() cil managed
+                    {
+                        .locals init (int32 V_0)
+                        ret
+                    }
+                    .method public static void SecondLocals() cil managed
+                    {
+                        .locals init (int32 V_0)
+                        ret
+                    }
+                    .method public static void DifferentLocals() cil managed
+                    {
+                        .locals init (int64 V_0)
+                        ret
+                    }
+                    .method public static void FirstBranch() cil managed
+                    {
+                        br.s Done
+                        Done: ret
+                    }
+                    .method public static void SecondBranch() cil managed
+                    {
+                        br.s Done
+                        Done: ret
+                    }
+                    .method public static void FirstFinally() cil managed
+                    {
+                        .try { leave.s Done }
+                        finally { endfinally }
+                        Done: ret
+                    }
+                    .method public static void SecondFinally() cil managed
+                    {
+                        .try { leave.s Done }
+                        finally { endfinally }
+                        Done: ret
+                    }
+                    .method public static void DifferentHandler() cil managed
+                    {
+                        .try { leave.s Done }
+                        fault { endfinally }
+                        Done: ret
+                    }
+                }
+                """;
+
+            using var pe = DocumentCompilerTestHelpers.CompileAndGetReader(source, new Options { Fold = fold });
+            var reader = pe.GetMetadataReader();
+            var methods = reader.MethodDefinitions
+                .Select(reader.GetMethodDefinition)
+                .ToDictionary(method => reader.GetString(method.Name), method => method.RelativeVirtualAddress);
+
+            Assert.NotEqual(0, methods["First"]);
+            Assert.Equal(fold, methods["First"] == methods["Second"]);
+            Assert.Equal(fold, methods["FirstLocals"] == methods["SecondLocals"]);
+            Assert.Equal(fold, methods["FirstBranch"] == methods["SecondBranch"]);
+            Assert.Equal(fold, methods["FirstFinally"] == methods["SecondFinally"]);
+            Assert.NotEqual(methods["First"], methods["DifferentCode"]);
+            Assert.NotEqual(methods["First"], methods["DifferentStack"]);
+            Assert.NotEqual(methods["First"], methods["FirstLocals"]);
+            Assert.NotEqual(methods["FirstLocals"], methods["DifferentLocals"]);
+            Assert.NotEqual(methods["FirstFinally"], methods["DifferentHandler"]);
+            Assert.Equal(new byte[] { 0x2a }, pe.GetMethodBody(methods["Second"]).GetILBytes());
+            Assert.False(pe.GetMethodBody(methods["SecondLocals"]).LocalSignature.IsNil);
+            Assert.Single(pe.GetMethodBody(methods["SecondFinally"]).ExceptionRegions);
+        }
+
         [Fact]
         public void PdbOption_EmitsEmbeddedPortablePdbWithoutLineDirectives()
         {

src/tools/ilasm/tests/ILAssembler.Tests/InstructionTests.cs

@@ -46,8 +46,10 @@ br UndefinedLabel
             Assert.Equal(DiagnosticSeverity.Error, error.Severity);
         }
 
-        [Fact]
-        public void UndefinedBranchTarget_WithErrorTolerantOption_PreservesNativeFatHeaderBehavior()
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public void UndefinedBranchTarget_WithErrorTolerantOption_PreservesNativeFatHeaderBehavior(bool fold)
         {
             string source = """
                 .assembly extern mscorlib { }
@@ -61,6 +63,12 @@ .maxstack 3
                         br UndefinedLabel
                         ret
                     }
+                    .method public static void OtherMethod() cil managed
+                    {
+                        .maxstack 3
+                        br UndefinedLabel
+                        ret
+                    }
                 }
                 """;
 
@@ -69,7 +77,7 @@ br UndefinedLabel
                 new SourceText(source, "test.il"),
                 _ => throw new InvalidOperationException("Unexpected include"),
                 _ => throw new InvalidOperationException("Unexpected resource"),
-                new Options { ErrorTolerant = true });
+                new Options { ErrorTolerant = true, Fold = fold });
 
             Assert.Contains(diagnostics, diagnostic => diagnostic.Id == DiagnosticIds.LabelNotFound);
             Assert.NotNull(result);
@@ -82,9 +90,14 @@ br UndefinedLabel
                 .Select(reader.GetMethodDefinition)
                 .Single(method => reader.GetString(method.Name) == "TestMethod");
             MethodBodyBlock body = pe.GetMethodBody(method.RelativeVirtualAddress);
+            MethodDefinition otherMethod = reader.MethodDefinitions
+                .Select(reader.GetMethodDefinition)
+                .Single(method => reader.GetString(method.Name) == "OtherMethod");
 
             Assert.Equal(3, body.MaxStack);
             Assert.True(body.LocalVariablesInitialized);
+            Assert.Equal(fold, method.RelativeVirtualAddress == otherMethod.RelativeVirtualAddress);
+            Assert.Equal(body.GetILBytes(), pe.GetMethodBody(otherMethod.RelativeVirtualAddress).GetILBytes());
         }
 
         [Fact]