Devirtualize TokenMap

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

@@ -7,9 +7,8 @@
 using System;
 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 +23,91 @@ 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];
-            }
-        }
-
-        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 object[] GetAllItems()
         {
-            lock (_items)
-            {
-                count = _items.Count;
-                return _items.ToArray();
-            }
+            // Read the count prior to getting the array
+            int count = Volatile.Read(ref _count);
+            object[] items = Volatile.Read(ref _items);
+
+            // Return a right sized copy of the array
+            return (new ReadOnlySpan<object>(items, 0, count)).ToArray();
         }
     }
 }

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 object[] ReferencesInIL()
         {
-            return _referencesInILMap.GetAllItemsAndCount(out count);
+            return _referencesInILMap.GetAllItems();
         }
 
         /// <summary>

src/Compilers/Core/Portable/IReferenceOrISignature.cs

@@ -0,0 +1,133 @@
+// 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.Collections.Generic;
+using Microsoft.Cci;
+using Microsoft.CodeAnalysis.Symbols;
+
+namespace Microsoft.CodeAnalysis
+{
+    /// <summary>
+    /// Used to devirtualize Dictionary/HashSet for EqualityComparer{T}.Default
+    /// </summary>
+    internal struct IReferenceOrISignatureEquivalent : IEquatable<IReferenceOrISignatureEquivalent>
+    {
+        private readonly object? _item;
+
+        public IReferenceOrISignatureEquivalent(IReference item)
+        {
+            _item = item;
+        }
+
+        public IReferenceOrISignatureEquivalent(ISignature item)
+        {
+            _item = item;
+        }
+
+        public IReferenceOrISignatureEquivalent(IMethodReference item)
+        {
+            _item = item;
+        }
+
+        public bool Equals(IReferenceOrISignatureEquivalent other)
+        {
+            object? x = _item;
+            object? y = other._item;
+            if (x is null)
+            {
+                return y is null;
+            }
+            else if (ReferenceEquals(x, y))
+            {
+                return true;
+            }
+
+            return Equals(x, y);
+        }
+
+        private new static bool Equals(object? x, object? y)
+        {
+            if (x is null)
+            {
+                return y is null;
+            }
+            else if (ReferenceEquals(x, y))
+            {
+                return true;
+            }
+            else 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)
+        {
+            return obj is IReferenceOrISignatureEquivalent refOrSig ?
+                Equals(refOrSig) :
+                Equals(_item, obj);
+        }
+
+        public override int GetHashCode() => _item?.GetHashCode() ?? 0;
+
+        public override string ToString() => _item?.ToString() ?? "null";
+
+        internal object AsObject() => _item!;
+    }
+
+    /// <summary>
+    /// Used to devirtualize ConcurrentDictionary for EqualityComparer{T}.Default and ReferenceEquals
+    /// </summary>
+    internal struct IReferenceOrISignature : IEquatable<IReferenceOrISignature>
+    {
+        private readonly object? _item;
+
+        public static implicit operator IReferenceOrISignature(IReferenceOrISignatureEquivalent d) => new IReferenceOrISignature(d.AsObject());
+
+        private IReferenceOrISignature(object? item)
+        {
+            _item = item;
+        }
+
+        public IReferenceOrISignature(IReference item)
+        {
+            _item = item;
+        }
+
+        public IReferenceOrISignature(ISignature item)
+        {
+            _item = item;
+        }
+
+        public bool Equals(IReferenceOrISignature other)
+        {
+            object? x = _item;
+            object? y = other._item;
+
+            return ReferenceEquals(x, y);
+        }
+
+        public override bool Equals(object? obj)
+        {
+            return obj is IReferenceOrISignature refOrSig ?
+                Equals(refOrSig) :
+                ReferenceEquals(_item, obj);
+        }
+
+        public override int GetHashCode() => _item?.GetHashCode() ?? 0;
+
+        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;
         }
 
         private void CreateIndices()