Avoid allocations when adding existing items to the blob heap. (#81059)

* Add `BlobDictionary`.

* Use `BlobDictionary` in `MetadataReader` and optimize the simple cases.

* Optimize adding blobs from UTF-16 strings and from blob builders.

* Delete `ByteSequenceComparer`.

* Fix compiler errors.

* Clean-up `BlobDictionary`; move it, make it non-generic and streamline its API.

* Don't try to copy the blob builder's contents to a stack-allocated buffer.
The blob builder's default starting size is 256 bytes; the same as the stack-alllocated buffer, and making it bigger does not have any benefit since such large blobs are rare.

Author: teo-tsirpanis

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

@@ -45,7 +45,6 @@ The System.Reflection.Metadata library is built-in as part of the shared framewo
     <Compile Include="System\Reflection\Metadata\BlobWriterImpl.cs" />
     <Compile Include="System\Reflection\Metadata\BlobBuilder.cs" />
     <Compile Include="System\Reflection\Metadata\BlobBuilder.Enumerators.cs" />
-    <Compile Include="System\Reflection\Internal\Utilities\ByteSequenceComparer.cs" />
     <Compile Include="System\Reflection\Internal\Utilities\DecimalUtilities.cs" />
     <Compile Include="System\Reflection\Internal\Utilities\EnumerableExtensions.cs" />
     <Compile Include="System\Reflection\Metadata\Ecma335\CustomAttributeDecoder.cs" />
@@ -60,6 +59,7 @@ The System.Reflection.Metadata library is built-in as part of the shared framewo
     <Compile Include="System\Reflection\Metadata\Ecma335\Encoding\InstructionEncoder.cs" />
     <Compile Include="System\Reflection\Metadata\Ecma335\Encoding\LabelHandle.cs" />
     <Compile Include="System\Reflection\Metadata\IL\ILOpCode.cs" />
+    <Compile Include="System\Reflection\Metadata\Ecma335\BlobDictionary.cs" />
     <Compile Include="System\Reflection\Metadata\Ecma335\CodedIndex.cs" />
     <Compile Include="System\Reflection\Metadata\Ecma335\PortablePdbBuilder.cs" />
     <Compile Include="System\Reflection\Metadata\Ecma335\MetadataBuilder.Tables.cs" />

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

@@ -50,6 +50,7 @@ public partial class BlobBuilder
         private bool IsHead => (_length & IsFrozenMask) == 0;
         private int Length => (int)(_length & ~IsFrozenMask);
         private uint FrozenLength => _length | IsFrozenMask;
+        private Span<byte> Span => _buffer.AsSpan(0, Length);
 
         public BlobBuilder(int capacity = DefaultChunkSize)
         {
@@ -234,7 +235,7 @@ public bool ContentEquals(BlobBuilder other)
                 var right = rightEnumerator.Current;
 
                 int minLength = Math.Min(left.Length - leftStart, right.Length - rightStart);
-                if (!ByteSequenceComparer.Equals(left._buffer, leftStart, right._buffer, rightStart, minLength))
+                if (!left._buffer.AsSpan(leftStart, minLength).SequenceEqual(right._buffer.AsSpan(rightStart, minLength)))
                 {
                     return false;
                 }
@@ -318,6 +319,19 @@ public ImmutableArray<byte> ToImmutableArray(int start, int byteCount)
             return ImmutableByteArrayInterop.DangerousCreateFromUnderlyingArray(ref array);
         }
 
+        internal bool TryGetSpan(out ReadOnlySpan<byte> buffer)
+        {
+            if (_nextOrPrevious == this)
+            {
+                // If the blob builder has one chunk, we can just return it and avoid copies.
+                buffer = Span;
+                return true;
+            }
+
+            buffer = default;
+            return false;
+        }
+
         /// <exception cref="ArgumentNullException"><paramref name="destination"/> is null.</exception>
         /// <exception cref="InvalidOperationException">Content is not available, the builder has been linked with another one.</exception>
         public void WriteContentTo(Stream destination)
@@ -344,7 +358,7 @@ public void WriteContentTo(ref BlobWriter destination)
 
             foreach (var chunk in GetChunks())
             {
-                destination.WriteBytes(chunk._buffer.AsSpan(0, chunk.Length));
+                destination.WriteBytes(chunk.Span);
             }
         }
 
@@ -359,7 +373,7 @@ public void WriteContentTo(BlobBuilder destination)
 
             foreach (var chunk in GetChunks())
             {
-                destination.WriteBytes(chunk._buffer.AsSpan(0, chunk.Length));
+                destination.WriteBytes(chunk.Span);
             }
         }
 

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

@@ -54,7 +54,7 @@ public BlobWriter(byte[] buffer, int start, int count)
         /// </summary>
         public bool ContentEquals(BlobWriter other)
         {
-            return Length == other.Length && ByteSequenceComparer.Equals(_buffer, _start, other._buffer, other._start, Length);
+            return Length == other.Length && _buffer.AsSpan(_start, Length).SequenceEqual(other._buffer.AsSpan(other._start, other.Length));
         }
 
         public int Offset

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

@@ -0,0 +1,109 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// 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;
+#if NET
+using System.Runtime.InteropServices;
+#endif
+
+namespace System.Reflection.Metadata.Ecma335
+{
+    [DebuggerDisplay("Count = {Count}")]
+    internal readonly struct BlobDictionary
+    {
+        private readonly Dictionary<int, KeyValuePair<ImmutableArray<byte>, BlobHandle>> _dictionary;
+
+        // A simple LCG. Constants taken from
+        // https://github.com/imneme/pcg-c/blob/83252d9c23df9c82ecb42210afed61a7b42402d7/include/pcg_variants.h#L276-L284
+        private static int GetNextDictionaryKey(int dictionaryKey) =>
+            (int)((uint)dictionaryKey * 747796405 + 2891336453);
+
+#if NET
+        private unsafe ref KeyValuePair<ImmutableArray<byte>, BlobHandle> GetValueRefOrAddDefault(ReadOnlySpan<byte> key, out bool exists)
+        {
+            int dictionaryKey = Hash.GetFNVHashCode(key);
+            while (true)
+            {
+                ref var entry = ref CollectionsMarshal.GetValueRefOrAddDefault(_dictionary, dictionaryKey, out exists);
+                if (!exists || entry.Key.AsSpan().SequenceEqual(key))
+                {
+#pragma warning disable CS9082 // Local is returned by reference but was initialized to a value that cannot be returned by reference
+                    // In .NET 6 the assembly of GetValueRefOrAddDefault was compiled with earlier ref safety rules
+                    // and caused an error, which was turned into a warning because of unsafe and was suppressed.
+                    return ref entry;
+#pragma warning restore CS9082
+                }
+                dictionaryKey = GetNextDictionaryKey(dictionaryKey);
+            }
+        }
+
+        public BlobHandle GetOrAdd(ReadOnlySpan<byte> key, ImmutableArray<byte> immutableKey, BlobHandle value, out bool exists)
+        {
+            ref var entry = ref GetValueRefOrAddDefault(key, out exists);
+            if (exists)
+            {
+                return entry.Value;
+            }
+
+            // If we are given an immutable array, do not allocate a new one.
+            if (immutableKey.IsDefault)
+            {
+                immutableKey = key.ToImmutableArray();
+            }
+            else
+            {
+                Debug.Assert(immutableKey.AsSpan().SequenceEqual(key));
+            }
+
+            entry = new(immutableKey, value);
+            return value;
+        }
+#else
+        public BlobHandle GetOrAdd(ReadOnlySpan<byte> key, ImmutableArray<byte> immutableKey, BlobHandle value, out bool exists)
+        {
+            int dictionarykey = Hash.GetFNVHashCode(key);
+            KeyValuePair<ImmutableArray<byte>, BlobHandle> entry;
+            while (true)
+            {
+                if (!(exists = _dictionary.TryGetValue(dictionarykey, out entry))
+                    || entry.Key.AsSpan().SequenceEqual(key))
+                {
+                    break;
+                }
+                dictionarykey = GetNextDictionaryKey(dictionarykey);
+            }
+
+            if (exists)
+            {
+                return entry.Value;
+            }
+
+            // If we are given an immutable array, do not allocate a new one.
+            if (immutableKey.IsDefault)
+            {
+                immutableKey = key.ToImmutableArray();
+            }
+            else
+            {
+                Debug.Assert(immutableKey.AsSpan().SequenceEqual(key));
+            }
+
+            _dictionary.Add(dictionarykey, new(immutableKey, value));
+            return value;
+        }
+#endif
+
+        public BlobDictionary(int capacity = 0)
+        {
+            _dictionary = new(capacity);
+        }
+
+        public int Count => _dictionary.Count;
+
+        public Dictionary<int, KeyValuePair<ImmutableArray<byte>, BlobHandle>>.Enumerator GetEnumerator() =>
+            _dictionary.GetEnumerator();
+    }
+}

src/libraries/System.Reflection.Metadata/src/System/Reflection/Metadata/Ecma335/BlobDictionary.cs

@@ -3,10 +3,8 @@
 
 using System.Collections.Generic;
 using System.Collections.Immutable;
-using System.Diagnostics;
 using System.Reflection.Internal;
-using System.Runtime.CompilerServices;
-using System.Text;
+using System.Runtime.InteropServices;
 
 namespace System.Reflection.Metadata.Ecma335
 {
@@ -44,7 +42,7 @@ internal void SetCapacity(int capacity)
         private int _stringHeapCapacity = 4 * 1024;
 
         // #Blob heap
-        private readonly Dictionary<ImmutableArray<byte>, BlobHandle> _blobs = new Dictionary<ImmutableArray<byte>, BlobHandle>(1024, ByteSequenceComparer.Instance);
+        private readonly BlobDictionary _blobs = new BlobDictionary(1024);
         private readonly int _blobHeapStartOffset;
         private int _blobHeapSize;
 
@@ -118,7 +116,7 @@ public MetadataBuilder(
             // beginning of the delta blob.
             _userStringBuilder.WriteByte(0);
 
-            _blobs.Add(ImmutableArray<byte>.Empty, default(BlobHandle));
+            _blobs.GetOrAdd(ReadOnlySpan<byte>.Empty, ImmutableArray<byte>.Empty, default, out _);
             _blobHeapSize = 1;
 
             // When EnC delta is applied #US, #String and #Blob heaps are appended.
@@ -193,7 +191,11 @@ public BlobHandle GetOrAddBlob(BlobBuilder value)
                 Throw.ArgumentNull(nameof(value));
             }
 
-            // TODO: avoid making a copy if the blob exists in the index
+            if (value.TryGetSpan(out ReadOnlySpan<byte> buffer))
+            {
+                return GetOrAddBlob(buffer);
+            }
+
             return GetOrAddBlob(value.ToImmutableArray());
         }
 
@@ -210,8 +212,19 @@ public BlobHandle GetOrAddBlob(byte[] value)
                 Throw.ArgumentNull(nameof(value));
             }
 
-            // TODO: avoid making a copy if the blob exists in the index
-            return GetOrAddBlob(ImmutableArray.Create(value));
+            return GetOrAddBlob(new ReadOnlySpan<byte>(value));
+        }
+
+        private BlobHandle GetOrAddBlob(ReadOnlySpan<byte> value, ImmutableArray<byte> immutableValue = default)
+        {
+            BlobHandle nextHandle = BlobHandle.FromOffset(_blobHeapStartOffset + _blobHeapSize);
+            BlobHandle handle = _blobs.GetOrAdd(value, immutableValue, nextHandle, out bool exists);
+            if (!exists)
+            {
+                _blobHeapSize += BlobWriterImpl.GetCompressedIntegerSize(value.Length) + value.Length;
+            }
+
+            return handle;
         }
 
         /// <summary>
@@ -227,16 +240,7 @@ public BlobHandle GetOrAddBlob(ImmutableArray<byte> value)
                 Throw.ArgumentNull(nameof(value));
             }
 
-            BlobHandle handle;
-            if (!_blobs.TryGetValue(value, out handle))
-            {
-                handle = BlobHandle.FromOffset(_blobHeapStartOffset + _blobHeapSize);
-                _blobs.Add(value, handle);
-
-                _blobHeapSize += BlobWriterImpl.GetCompressedIntegerSize(value.Length) + value.Length;
-            }
-
-            return handle;
+            return GetOrAddBlob(value.AsSpan(), value);
         }
 
         /// <summary>
@@ -267,6 +271,16 @@ public unsafe BlobHandle GetOrAddConstantBlob(object? value)
         /// <exception cref="ArgumentNullException"><paramref name="value"/> is null.</exception>
         public BlobHandle GetOrAddBlobUTF16(string value)
         {
+            if (value is null)
+            {
+                Throw.ArgumentNull(nameof(value));
+            }
+
+            if (BitConverter.IsLittleEndian)
+            {
+                return GetOrAddBlob(MemoryMarshal.AsBytes(value.AsSpan()));
+            }
+
             var builder = PooledBlobBuilder.GetInstance();
             builder.WriteUTF16(value);
             var handle = GetOrAddBlob(builder);
@@ -611,8 +625,8 @@ private void WriteAlignedBlobHeap(BlobBuilder builder)
             int startOffset = _blobHeapStartOffset;
             foreach (var entry in _blobs)
             {
-                int heapOffset = entry.Value.GetHeapOffset();
-                var blob = entry.Key;
+                int heapOffset = entry.Value.Value.GetHeapOffset();
+                var blob = entry.Value.Key;
 
                 writer.Offset = (heapOffset == 0) ? 0 : heapOffset - startOffset;
                 writer.WriteCompressedInteger(blob.Length);

src/libraries/System.Reflection.Metadata/src/System/Reflection/Metadata/Ecma335/MetadataBuilder.Heaps.cs

@@ -61,7 +61,7 @@ private void TestContentEquals(byte[] left, byte[] right)
             var builder2 = new BlobBuilder(0);
             builder2.WriteBytes(right);
 
-            bool expected = ByteSequenceComparer.Equals(left, right);
+            bool expected = left.AsSpan().SequenceEqual(right.AsSpan());
             Assert.Equal(expected, builder1.ContentEquals(builder2));
         }