Merge pull request #47274 from benaadams/Devirtualize-TokenMap

Devirtualize TokenMap

src/Compilers/Core/Portable/CodeGen/TokenMap.cs

@@ -8,8 +8,8 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
-using Microsoft.CodeAnalysis.PooledObjects;
-using Roslyn.Utilities;
+using System.Threading;
+using Microsoft.Cci;
 
 namespace Microsoft.CodeAnalysis.CodeGen
 {
@@ -24,74 +24,97 @@ namespace Microsoft.CodeAnalysis.CodeGen
     /// </summary>
     internal sealed class TokenMap
     {
-        private readonly ConcurrentDictionary<object, uint> _itemIdentityToToken = new ConcurrentDictionary<object, uint>(ReferenceEqualityComparer.Instance);
+        private readonly ConcurrentDictionary<IReferenceOrISignature, uint> _itemIdentityToToken = new();
+        private readonly Dictionary<IReferenceOrISignatureEquivalent, uint> _itemEquivalentToToken = new();
+        private object[] _items = Array.Empty<object>();
+        private int _count = 0;
 
-        private readonly Dictionary<object, uint> _itemToToken;
-        private readonly ArrayBuilder<object> _items = new ArrayBuilder<object>();
+        internal TokenMap() { }
 
-        internal TokenMap(IEqualityComparer<object> comparer)
+        public uint GetOrAddTokenFor(IReference item, out bool referenceAdded)
         {
-            _itemToToken = new Dictionary<object, uint>(comparer);
+            if (_itemIdentityToToken.TryGetValue(new IReferenceOrISignature(item), out uint token))
+            {
+                referenceAdded = false;
+                return token;
+            }
+
+            return AddItem(new IReferenceOrISignatureEquivalent(item), out referenceAdded);
         }
 
-        public uint GetOrAddTokenFor(object item, out bool referenceAdded)
+        public uint GetOrAddTokenFor(ISignature item, out bool referenceAdded)
         {
-            uint tmp;
-            if (_itemIdentityToToken.TryGetValue(item, out tmp))
+            if (_itemIdentityToToken.TryGetValue(new IReferenceOrISignature(item), out uint token))
             {
                 referenceAdded = false;
-                return (uint)tmp;
+                return token;
             }
 
-            return AddItem(item, out referenceAdded);
+            return AddItem(new IReferenceOrISignatureEquivalent(item), out referenceAdded);
         }
 
-        private uint AddItem(object item, out bool referenceAdded)
+        private uint AddItem(IReferenceOrISignatureEquivalent item, out bool referenceAdded)
         {
-            Debug.Assert(item is Cci.ISignature || item is Cci.IReference);
             uint token;
-
             // NOTE: cannot use GetOrAdd here since items and itemToToken must be in sync
             // so if we do need to add we have to take a lock and modify both collections.
-            lock (_items)
+            lock (_itemEquivalentToToken)
             {
-                if (!_itemToToken.TryGetValue(item, out token))
+                // Check if there is an equivalent type that has a token
+                if (!_itemEquivalentToToken.TryGetValue(item, out token))
                 {
-                    token = (uint)_items.Count;
-                    _items.Add(item);
-                    _itemToToken.Add(item, token);
+                    token = (uint)_count;
+                    // No equivalent, add the token for this type
+                    _itemEquivalentToToken.Add(item, token);
+
+                    var count = (int)token + 1;
+                    var items = _items;
+                    if (items.Length > count)
+                    {
+                        items[(int)token] = item.AsObject();
+                    }
+                    else
+                    {
+                        // Not enough room, we need to resize the array
+                        Array.Resize(ref items, Math.Max(8, count * 2));
+                        items[(int)token] = item.AsObject();
+
+                        // Update the updated array reference before updating _count
+                        Volatile.Write(ref _items, items);
+                    }
+
+                    Volatile.Write(ref _count, count);
                 }
             }
 
+            // Use the provided token to update the reference dictionary
             referenceAdded = _itemIdentityToToken.TryAdd(item, token);
             return token;
         }
 
         public object GetItem(uint token)
         {
-            lock (_items)
-            {
-                return _items[(int)token];
-            }
-        }
+            // If a token has been handed out, then it should be always within _count of the
+            // current array and a lock is not required.
+            Debug.Assert(token < (uint)_count && _count <= _items.Length);
 
-        public IEnumerable<object> GetAllItems()
-        {
-            lock (_items)
-            {
-                return _items.ToArray();
-            }
+            return _items[(int)token];
         }
 
         //TODO: why is this is called twice during emit?
         //      should probably return ROA instead of IE and cache that in Module. (and no need to return count)
-        public IEnumerable<object> GetAllItemsAndCount(out int count)
+        public ReadOnlySpan<object> GetAllItems()
         {
-            lock (_items)
-            {
-                count = _items.Count;
-                return _items.ToArray();
-            }
+            // Read _count before _items reference, to match inverse of the writes in AddItem.
+            // So _items is guaranteed to have at least count items; and a lock is not required.
+
+            // Read the count prior to getting the array
+            int count = Volatile.Read(ref _count);
+            // Read the array reference
+            object[] items = Volatile.Read(ref _items);
+
+            // Return a right sized view of the array based on read count and reference.
+            return new ReadOnlySpan<object>(items, 0, count);
         }
     }
 }

src/Compilers/Core/Portable/Emit/CommonPEModuleBuilder.cs

@@ -33,9 +33,9 @@ internal abstract class CommonPEModuleBuilder : Cci.IUnit, Cci.IModuleReference
         internal Cci.IMethodReference DebugEntryPoint;
 
         private readonly ConcurrentDictionary<IMethodSymbolInternal, Cci.IMethodBody> _methodBodyMap;
-        private readonly TokenMap _referencesInILMap = new TokenMap(MetadataEntityReferenceComparer.ConsiderEverything);
-        private readonly ItemTokenMap<string> _stringsInILMap = new ItemTokenMap<string>();
-        private readonly ItemTokenMap<Cci.DebugSourceDocument> _sourceDocumentsInILMap = new ItemTokenMap<Cci.DebugSourceDocument>();
+        private readonly TokenMap _referencesInILMap = new();
+        private readonly ItemTokenMap<string> _stringsInILMap = new();
+        private readonly ItemTokenMap<Cci.DebugSourceDocument> _sourceDocumentsInILMap = new();
 
         private ImmutableArray<Cci.AssemblyReferenceAlias> _lazyAssemblyReferenceAliases;
         private ImmutableArray<Cci.ManagedResource> _lazyManagedResources;
@@ -335,9 +335,9 @@ public string GetStringFromToken(uint token)
             return _stringsInILMap.GetItem(token);
         }
 
-        public IEnumerable<object> ReferencesInIL(out int count)
+        public ReadOnlySpan<object> ReferencesInIL()
         {
-            return _referencesInILMap.GetAllItemsAndCount(out count);
+            return _referencesInILMap.GetAllItems();
         }
 
         /// <summary>

src/Compilers/Core/Portable/IReferenceOrISignature.cs

@@ -0,0 +1,97 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+#nullable enable
+
+using System;
+using System.Runtime.CompilerServices;
+using Microsoft.Cci;
+using Microsoft.CodeAnalysis.Symbols;
+
+namespace Microsoft.CodeAnalysis
+{
+    // These types are to enable fast-path devirtualization in the Jit. Dictionary<K, V>, HashTable<T>
+    // and ConcurrentDictionary<K, V> will devirtualize (and potentially inline) the IEquatable<T>.Equals
+    // method for a struct when the Comparer is unspecified in .NET Core, .NET 5; whereas specifying
+    // a Comparer will make .Equals and GetHashcode slower interface calls.
+
+    /// <summary>
+    /// Used to devirtualize Dictionary/HashSet for EqualityComparer{T}.Default
+    /// </summary>
+    internal readonly struct IReferenceOrISignatureEquivalent : IEquatable<IReferenceOrISignatureEquivalent>
+    {
+        private readonly object _item;
+
+        public IReferenceOrISignatureEquivalent(IReference item) => _item = item;
+
+        public IReferenceOrISignatureEquivalent(ISignature item) => _item = item;
+
+        // Needed to resolve ambiguity for types that implement both IReference and ISignature
+        public IReferenceOrISignatureEquivalent(IMethodReference item) => _item = item;
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public bool Equals(IReferenceOrISignatureEquivalent other)
+        {
+            // Fast inlinable ReferenceEquals
+            var x = _item;
+            var y = other._item;
+            if (ReferenceEquals(x, y))
+            {
+                return true;
+            }
+
+            return EqualsSlow(x, y);
+        }
+
+        private static bool EqualsSlow(object x, object y)
+        {
+            if (x is ISymbolInternal sx && y is ISymbolInternal sy)
+            {
+                return sx.Equals(sy, TypeCompareKind.ConsiderEverything);
+            }
+            else if (x is ISymbolCompareKindComparableInternal cx && y is ISymbolCompareKindComparableInternal cy)
+            {
+                return cx.Equals(cy, TypeCompareKind.ConsiderEverything);
+            }
+            else
+            {
+                return x.Equals(y);
+            }
+        }
+
+        public override bool Equals(object? obj) => false;
+
+        public override int GetHashCode() => _item.GetHashCode();
+
+        public override string ToString() => _item.ToString() ?? "null";
+
+        internal object AsObject() => _item;
+    }
+
+    /// <summary>
+    /// Used to devirtualize ConcurrentDictionary for EqualityComparer{T}.Default and ReferenceEquals
+    /// </summary>
+    internal readonly struct IReferenceOrISignature : IEquatable<IReferenceOrISignature>
+    {
+        private readonly object _item;
+
+        public IReferenceOrISignature(IReference item) => _item = item;
+
+        public IReferenceOrISignature(ISignature item) => _item = item;
+
+        // Used by implicit conversion
+        private IReferenceOrISignature(object item) => _item = item;
+
+        public static implicit operator IReferenceOrISignature(IReferenceOrISignatureEquivalent item)
+            => new IReferenceOrISignature(item.AsObject());
+
+        public bool Equals(IReferenceOrISignature other) => ReferenceEquals(_item, other._item);
+
+        public override bool Equals(object? obj) => false;
+
+        public override int GetHashCode() => _item.GetHashCode();
+
+        public override string ToString() => _item.ToString() ?? "null";
+    }
+}

src/Compilers/Core/Portable/PEWriter/MetadataWriter.cs

@@ -464,18 +464,10 @@ private bool IsMinimalDelta
 
         private void CreateMethodBodyReferenceIndex()
         {
-            int count;
-            var referencesInIL = module.ReferencesInIL(out count);
+            var referencesInIL = module.ReferencesInIL();
 
-            _pseudoSymbolTokenToTokenMap = new EntityHandle[count];
-            _pseudoSymbolTokenToReferenceMap = new object[count];
-
-            int cur = 0;
-            foreach (object o in referencesInIL)
-            {
-                _pseudoSymbolTokenToReferenceMap[cur] = o;
-                cur++;
-            }
+            _pseudoSymbolTokenToTokenMap = new EntityHandle[referencesInIL.Length];
+            _pseudoSymbolTokenToReferenceMap = referencesInIL.ToArray();
         }
 
         private void CreateIndices()