Add InstructionEncoder.Switch. (#76526)

* Avoid passing a redundant parameter in BranchInfo.GetBranchDistance.

The opcode is already stored in a field.
And it's no more stored as a byte; it takes the same space as an ILOpCode due to packing.

* Track only the label operand instead of the whole instruction in ControlFlowBuilder.

It simplifies the control flow builder and paves the way to support the switch instruction in the future.
As a drawback the BranchInfo struct became bigger due to fields that are necessary for error reporting.

* Replace immutable array builders with regular lists in ControlFlowBuilder.

There didn't seem to be any reason to use them; they were not converted to immutable arrays.

* Clean-up EnumerableExtensions; remove an unused method and forward another one.

* Fix the tests.

And rename instructionStartOffset to ilOffset.

* Add InstructionEncoder.Switch.

* Add tests.

* Remove a potentially unclear comment.

* Validate that an `InstructionEncoder` is not used in any other way while switch branches are being emitted.

* Use named parameters instead of just `default`.

* Remove an automatically generated comment in the resx.

* Revert "Remove a potentially unclear comment."

This reverts commit 92a1f80.

* Apply suggestions from code review

Co-authored-by: Stephen Toub <[email protected]>

* Remove `EnumerableExtensions.FirstOrDefault`.

* Rename `AssertNotInSwitch` to `ValidateNotInSwitch` and call it on `DefineLabel`.

Co-authored-by: Stephen Toub <[email protected]>

src/libraries/System.Reflection.Metadata/ref/System.Reflection.Metadata.cs

@@ -2689,6 +2689,7 @@ public void MarkLabel(System.Reflection.Metadata.Ecma335.LabelHandle label) { }
         public void OpCode(System.Reflection.Metadata.ILOpCode code) { }
         public void StoreArgument(int argumentIndex) { }
         public void StoreLocal(int slotIndex) { }
+        public System.Reflection.Metadata.Ecma335.SwitchInstructionEncoder Switch(int branchCount) { throw null; }
         public void Token(int token) { }
         public void Token(System.Reflection.Metadata.EntityHandle handle) { }
     }
@@ -3069,6 +3070,10 @@ public void UInt64() { }
         public void UIntPtr() { }
         public void VoidPointer() { }
     }
+    public readonly struct SwitchInstructionEncoder
+    {
+        public void Branch(System.Reflection.Metadata.Ecma335.LabelHandle label) { }
+    }
     public enum TableIndex : byte
     {
         Module = (byte)0,

src/libraries/System.Reflection.Metadata/src/Resources/Strings.resx

@@ -1,4 +1,4 @@
-<root>
+<root>
   <xsd:schema id="root" xmlns="" xmlns:xsd="http://www.w3.org/2001/XMLSchema" xmlns:msdata="urn:schemas-microsoft-com:xml-msdata">
     <xsd:import namespace="http://www.w3.org/XML/1998/namespace" />
     <xsd:element name="root" msdata:IsDataSet="true">
@@ -405,4 +405,10 @@
   <data name="UnexpectedValueUnknownType" xml:space="preserve">
     <value>Unexpected value '{0}' of unknown type.</value>
   </data>
-</root>
+  <data name="SwitchInstructionEncoderTooFewBranches" xml:space="preserve">
+    <value>The SwitchInstructionEncoder.Branch method was invoked too few times.</value>
+  </data>
+  <data name="SwitchInstructionEncoderTooManyBranches" xml:space="preserve">
+    <value>The SwitchInstructionEncoder.Branch method was invoked too many times.</value>
+  </data>
+</root>

src/libraries/System.Reflection.Metadata/src/System.Reflection.Metadata.csproj

@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
   <PropertyGroup>
     <TargetFrameworks>$(NetCoreAppCurrent);$(NetCoreAppPrevious);$(NetCoreAppMinimum);netstandard2.0;$(NetFrameworkMinimum)</TargetFrameworks>
     <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
@@ -31,6 +31,7 @@ The System.Reflection.Metadata library is built-in as part of the shared framewo
     <Compile Include="System\Reflection\Metadata\Ecma335\Encoding\FunctionPointerAttributes.cs" />
     <Compile Include="System\Reflection\Metadata\Ecma335\Encoding\MethodBodyStreamEncoder.cs" />
     <Compile Include="System\Reflection\Metadata\Ecma335\Encoding\MethodBodyAttributes.cs" />
+    <Compile Include="System\Reflection\Metadata\Ecma335\Encoding\SwitchInstructionEncoder.cs" />
     <Compile Include="System\Reflection\Metadata\Ecma335\MetadataBuilder.cs" />
     <Compile Include="System\Reflection\Metadata\Ecma335\MetadataRootBuilder.cs" />
     <Compile Include="System\Reflection\Metadata\Ecma335\SerializedMetadataHeaps.cs" />
@@ -89,7 +90,7 @@ The System.Reflection.Metadata library is built-in as part of the shared framewo
     <Compile Include="System\Reflection\PortableExecutable\ResourceSectionBuilder.cs" />
     <Compile Include="System\Reflection\PortableExecutable\SectionLocation.cs" />
     <Compile Include="System\Reflection\Internal\Utilities\BlobUtilities.cs" />
-    <None Include="README.md" Pack="true" PackagePath="\"/>
+    <None Include="README.md" Pack="true" PackagePath="\" />
   </ItemGroup>
 
   <ItemGroup>

src/libraries/System.Reflection.Metadata/src/System/Reflection/Internal/Utilities/EnumerableExtensions.cs

@@ -2,8 +2,6 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Generic;
-using System.Collections.Immutable;
-using System.Diagnostics.CodeAnalysis;
 
 namespace System.Reflection.Internal
 {
@@ -13,19 +11,6 @@ namespace System.Reflection.Internal
     /// </summary>
     internal static class EnumerableExtensions
     {
-        public static T? FirstOrDefault<T>(this ImmutableArray<T> collection, Func<T, bool> predicate)
-        {
-            foreach (var item in collection)
-            {
-                if (predicate(item))
-                {
-                    return item;
-                }
-            }
-
-            return default;
-        }
-
         // used only in debugger display so we needn't get fancy with optimizations.
         public static IEnumerable<TResult> Select<TSource, TResult>(this IEnumerable<TSource> source, Func<TSource, TResult> selector)
         {
@@ -35,11 +20,6 @@ public static IEnumerable<TResult> Select<TSource, TResult>(this IEnumerable<TSo
             }
         }
 
-        public static T Last<T>(this ImmutableArray<T>.Builder source)
-        {
-            return source[source.Count - 1];
-        }
-
         public static IEnumerable<T> OrderBy<T>(this List<T> source, Comparison<T> comparison)
         {
             // Produce an iterator that represents a stable sort of source.

src/libraries/System.Reflection.Metadata/src/System/Reflection/Metadata/Ecma335/Encoding/ControlFlowBuilder.cs

@@ -2,9 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Generic;
-using System.Collections.Immutable;
 using System.Diagnostics;
-using System.Reflection.Internal;
 
 namespace System.Reflection.Metadata.Ecma335
 {
@@ -13,37 +11,50 @@ public sealed class ControlFlowBuilder
         // internal for testing:
         internal readonly struct BranchInfo
         {
-            internal readonly int ILOffset;
+            // The offset to the label operand inside the instruction.
+            internal readonly int OperandOffset;
             internal readonly LabelHandle Label;
-            private readonly byte _opCode;
+            // Label offsets are calculated from the end of the instruction that contains them.
+            // This value contains the displacement from the start of the label operand
+            // to the end of the instruction. It is equal to one on short branches,
+            // four on long branches and bigger on the switch instruction.
+            private readonly int _instructionEndDisplacement;
+
+            // The following two fields are used for error reporting and tests.
+
+            // The offset to the start of the instruction.
+            internal readonly int ILOffset;
+            internal readonly ILOpCode OpCode;
 
-            internal ILOpCode OpCode => (ILOpCode)_opCode;
+            internal bool IsShortBranch => _instructionEndDisplacement == 1;
+            internal int OperandSize => Math.Min(_instructionEndDisplacement, 4);
 
-            internal BranchInfo(int ilOffset, LabelHandle label, ILOpCode opCode)
+            internal BranchInfo(int operandOffset, LabelHandle label, int instructionEndDisplacement, int ilOffset, ILOpCode opCode)
             {
-                ILOffset = ilOffset;
+                OperandOffset = operandOffset;
                 Label = label;
-                _opCode = (byte)opCode;
+                _instructionEndDisplacement = instructionEndDisplacement;
+                ILOffset = ilOffset;
+                OpCode = opCode;
             }
 
-            internal int GetBranchDistance(ImmutableArray<int>.Builder labels, ILOpCode branchOpCode, int branchILOffset, bool isShortBranch)
+            internal int GetBranchDistance(List<int> labels)
             {
                 int labelTargetOffset = labels[Label.Id - 1];
                 if (labelTargetOffset < 0)
                 {
                     Throw.InvalidOperation_LabelNotMarked(Label.Id);
                 }
 
-                int branchInstructionSize = 1 + (isShortBranch ? sizeof(sbyte) : sizeof(int));
-                int distance = labelTargetOffset - (ILOffset + branchInstructionSize);
+                int distance = labelTargetOffset - (OperandOffset + _instructionEndDisplacement);
 
-                if (isShortBranch && unchecked((sbyte)distance) != distance)
+                if (IsShortBranch && unchecked((sbyte)distance) != distance)
                 {
                     // We could potentially implement algorithm that automatically fixes up branch instructions to accommodate for bigger distances (short vs long),
                     // however an optimal algorithm would be rather complex (something like: calculate topological ordering of crossing branch instructions
                     // and then use fixed point to eliminate cycles). If the caller doesn't care about optimal IL size they can use long branches whenever the
                     // distance is unknown upfront. If they do they probably implement more sophisticated algorithm for IL layout optimization already.
-                    throw new InvalidOperationException(SR.Format(SR.DistanceBetweenInstructionAndLabelTooBig, branchOpCode, branchILOffset, distance));
+                    throw new InvalidOperationException(SR.Format(SR.DistanceBetweenInstructionAndLabelTooBig, OpCode, ILOffset, distance));
                 }
 
                 return distance;
@@ -75,14 +86,14 @@ public ExceptionHandlerInfo(
             }
         }
 
-        private readonly ImmutableArray<BranchInfo>.Builder _branches;
-        private readonly ImmutableArray<int>.Builder _labels;
-        private ImmutableArray<ExceptionHandlerInfo>.Builder? _lazyExceptionHandlers;
+        private readonly List<BranchInfo> _branches;
+        private readonly List<int> _labels;
+        private List<ExceptionHandlerInfo>? _lazyExceptionHandlers;
 
         public ControlFlowBuilder()
         {
-            _branches = ImmutableArray.CreateBuilder<BranchInfo>();
-            _labels = ImmutableArray.CreateBuilder<int>();
+            _branches = new List<BranchInfo>();
+            _labels = new List<int>();
         }
 
         /// <summary>
@@ -93,25 +104,42 @@ public void Clear()
             _branches.Clear();
             _labels.Clear();
             _lazyExceptionHandlers?.Clear();
+            RemainingSwitchBranches = 0;
         }
 
         internal LabelHandle AddLabel()
         {
+            ValidateNotInSwitch();
             _labels.Add(-1);
             return new LabelHandle(_labels.Count);
         }
 
-        internal void AddBranch(int ilOffset, LabelHandle label, ILOpCode opCode)
+        internal void AddBranch(int operandOffset, LabelHandle label, int instructionEndDisplacement, int ilOffset, ILOpCode opCode)
         {
-            Debug.Assert(ilOffset >= 0);
-            Debug.Assert(_branches.Count == 0 || ilOffset > _branches.Last().ILOffset);
+            Debug.Assert(operandOffset >= 0);
+            Debug.Assert(_branches.Count == 0 || operandOffset > _branches[_branches.Count - 1].OperandOffset);
             ValidateLabel(label, nameof(label));
-            _branches.Add(new BranchInfo(ilOffset, label, opCode));
+#if DEBUG
+            switch (instructionEndDisplacement)
+            {
+                case 1:
+                    Debug.Assert(opCode.GetBranchOperandSize() == 1);
+                    break;
+                case 4:
+                    Debug.Assert(opCode == ILOpCode.Switch || opCode.GetBranchOperandSize() == 4);
+                    break;
+                default:
+                    Debug.Assert(instructionEndDisplacement > 4 && instructionEndDisplacement % 4 == 0 && opCode == ILOpCode.Switch);
+                    break;
+            }
+#endif
+            _branches.Add(new BranchInfo(operandOffset, label, instructionEndDisplacement, ilOffset, opCode));
         }
 
         internal void MarkLabel(int ilOffset, LabelHandle label)
         {
             Debug.Assert(ilOffset >= 0);
+            ValidateNotInSwitch();
             ValidateLabel(label, nameof(label));
             _labels[label.Id - 1] = ilOffset;
         }
@@ -214,8 +242,9 @@ private void AddExceptionRegion(
             ValidateLabel(tryEnd, nameof(tryEnd));
             ValidateLabel(handlerStart, nameof(handlerStart));
             ValidateLabel(handlerEnd, nameof(handlerEnd));
+            ValidateNotInSwitch();
 
-            _lazyExceptionHandlers ??= ImmutableArray.CreateBuilder<ExceptionHandlerInfo>();
+            _lazyExceptionHandlers ??= new List<ExceptionHandlerInfo>();
 
             _lazyExceptionHandlers.Add(new ExceptionHandlerInfo(kind, tryStart, tryEnd, handlerStart, handlerEnd, filterStart, catchType));
         }
@@ -230,6 +259,25 @@ private void AddExceptionRegion(
 
         internal int ExceptionHandlerCount => _lazyExceptionHandlers?.Count ?? 0;
 
+        internal int RemainingSwitchBranches { get; set; }
+
+        internal void ValidateNotInSwitch()
+        {
+            if (RemainingSwitchBranches > 0)
+            {
+                Throw.InvalidOperation(SR.SwitchInstructionEncoderTooFewBranches);
+            }
+        }
+
+        internal void SwitchBranchAdded()
+        {
+            if (RemainingSwitchBranches == 0)
+            {
+                Throw.InvalidOperation(SR.SwitchInstructionEncoderTooManyBranches);
+            }
+            RemainingSwitchBranches--;
+        }
+
         /// <exception cref="InvalidOperationException" />
         internal void CopyCodeAndFixupBranches(BlobBuilder srcBuilder, BlobBuilder dstBuilder)
         {
@@ -252,7 +300,7 @@ internal void CopyCodeAndFixupBranches(BlobBuilder srcBuilder, BlobBuilder dstBu
                 while (true)
                 {
                     // copy bytes preceding the next branch, or till the end of the blob:
-                    int chunkSize = Math.Min(branch.ILOffset - srcOffset, srcBlob.Length - srcBlobOffset);
+                    int chunkSize = Math.Min(branch.OperandOffset - srcOffset, srcBlob.Length - srcBlobOffset);
                     dstBuilder.WriteBytes(srcBlob.Buffer, srcBlobOffset, chunkSize);
                     srcOffset += chunkSize;
                     srcBlobOffset += chunkSize;
@@ -264,22 +312,17 @@ internal void CopyCodeAndFixupBranches(BlobBuilder srcBuilder, BlobBuilder dstBu
                         break;
                     }
 
-                    Debug.Assert(srcBlob.Buffer[srcBlobOffset] == (byte)branch.OpCode);
-
-                    int operandSize = branch.OpCode.GetBranchOperandSize();
-                    bool isShortInstruction = operandSize == 1;
+                    int operandSize = branch.OperandSize;
+                    bool isShortInstruction = branch.IsShortBranch;
 
                     // Note: the 4B operand is contiguous since we wrote it via BlobBuilder.WriteInt32()
                     Debug.Assert(
-                        srcBlobOffset + 1 == srcBlob.Length ||
+                        srcBlobOffset == srcBlob.Length ||
                         (isShortInstruction ?
-                           srcBlob.Buffer[srcBlobOffset + 1] == 0xff :
-                           BitConverter.ToUInt32(srcBlob.Buffer, srcBlobOffset + 1) == 0xffffffff));
-
-                    // write branch opcode:
-                    dstBuilder.WriteByte(srcBlob.Buffer[srcBlobOffset]);
+                           srcBlob.Buffer[srcBlobOffset] == 0xff :
+                           BitConverter.ToUInt32(srcBlob.Buffer, srcBlobOffset) == 0xffffffff));
 
-                    int branchDistance = branch.GetBranchDistance(_labels, branch.OpCode, srcOffset, isShortInstruction);
+                    int branchDistance = branch.GetBranchDistance(_labels);
 
                     // write branch operand:
                     if (isShortInstruction)
@@ -291,13 +334,16 @@ internal void CopyCodeAndFixupBranches(BlobBuilder srcBuilder, BlobBuilder dstBu
                         dstBuilder.WriteInt32(branchDistance);
                     }
 
-                    srcOffset += sizeof(byte) + operandSize;
+                    srcOffset += operandSize;
 
                     // next branch:
                     branchIndex++;
                     if (branchIndex == _branches.Count)
                     {
-                        branch = new BranchInfo(int.MaxValue, label: default, opCode: default);
+                        // We have processed all branches. The MaxValue will cause the rest
+                        // of the IL stream to be directly copied to the destination blob.
+                        branch = new BranchInfo(operandOffset: int.MaxValue, label: default,
+                            instructionEndDisplacement: default, ilOffset: default, opCode: default);
                     }
                     else
                     {
@@ -311,8 +357,8 @@ internal void CopyCodeAndFixupBranches(BlobBuilder srcBuilder, BlobBuilder dstBu
                         break;
                     }
 
-                    // skip fake branch instruction:
-                    srcBlobOffset += sizeof(byte) + operandSize;
+                    // skip fake branch operand:
+                    srcBlobOffset += operandSize;
                 }
             }
         }