Merge pull request #11254 from CyrusNajmabadi/metadataReading3

Improve OrderPreservingMultiDictionary.

src/Compilers/Core/Portable/Collections/OrderPreservingMultiDictionary.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
+using System;
+using System.Collections;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics;
@@ -8,14 +10,14 @@
 namespace Microsoft.CodeAnalysis.Collections
 {
     /// <summary>
-    /// A MultiDictionary that allows only adding, and 
-    /// preserves the order of values added to the dictionary.
-    /// Thread-safe for reading, but not for adding.
+    /// A MultiDictionary that allows only adding, and preserves the order of values added to the 
+    /// dictionary. Thread-safe for reading, but not for adding.
     /// </summary>
     /// <remarks>
     /// Always uses the default comparer.
     /// </remarks>
-    internal sealed class OrderPreservingMultiDictionary<K, V>
+    internal sealed class OrderPreservingMultiDictionary<K, V> :
+        IEnumerable<KeyValuePair<K, OrderPreservingMultiDictionary<K, V>.ValueSet>>
     {
         #region Pooling
 
@@ -30,6 +32,12 @@ public void Free()
         {
             if (_dictionary != null)
             {
+                // Allow our ValueSets to return their underlying ArrayBuilders to the pool.
+                foreach (var kvp in _dictionary)
+                {
+                    kvp.Value.Free();
+                }
+
                 _dictionary.Free();
                 _dictionary = null;
             }
@@ -57,159 +65,231 @@ public void Free()
 
         #endregion Pooling
 
+        // An empty dictionary we keep around to simplify certain operations (like "Keys")
+        // when we don't have an underlying dictionary of our own.
+        private static readonly Dictionary<K, ValueSet> s_emptyDictionary = new Dictionary<K, ValueSet>();
+
+        // The underlying dictionary we store our data in.  null if we are empty.
+        private PooledDictionary<K, ValueSet> _dictionary;
+
         public OrderPreservingMultiDictionary()
         {
         }
 
-        // store either a single V or an ArrayBuilder<V>
-        /// <summary>
-        /// Each value is either a single V or an <see cref="ArrayBuilder{V}"/>.
-        /// Null when the dictionary is empty.
-        /// Don't access the field directly.
-        /// </summary>
-        private PooledDictionary<K, object> _dictionary;
-
         private void EnsureDictionary()
         {
-            _dictionary = _dictionary ?? PooledDictionary<K, object>.GetInstance();
+            _dictionary = _dictionary ?? PooledDictionary<K, ValueSet>.GetInstance();
         }
 
-        public bool IsEmpty
-        {
-            get { return _dictionary == null; }
-        }
+        public bool IsEmpty => _dictionary == null;
 
         /// <summary>
         /// Add a value to the dictionary.
         /// </summary>
         public void Add(K k, V v)
         {
-            object item;
-            if (!this.IsEmpty && _dictionary.TryGetValue(k, out item))
+            ValueSet valueSet;
+            if (!this.IsEmpty && _dictionary.TryGetValue(k, out valueSet))
             {
-                var arrayBuilder = item as ArrayBuilder<V>;
-                if (arrayBuilder == null)
-                {
-                    // Promote from singleton V to ArrayBuilder<V>.
-                    Debug.Assert(item is V, "Item must be either a V or an ArrayBuilder<V>");
-                    arrayBuilder = new ArrayBuilder<V>(2);
-                    arrayBuilder.Add((V)item);
-                    arrayBuilder.Add(v);
-                    _dictionary[k] = arrayBuilder;
-                }
-                else
-                {
-                    arrayBuilder.Add(v);
-                }
+                Debug.Assert(valueSet.Count >= 1);
+                // Have to re-store the ValueSet in case we upgraded the existing ValueSet from 
+                // holding a single item to holding multiple items.
+                _dictionary[k] = valueSet.WithAddedItem(v);
             }
             else
             {
                 this.EnsureDictionary();
-                _dictionary[k] = v;
+                _dictionary[k] = new ValueSet(v);
             }
         }
 
-        /// <summary>
-        /// Add multiple values to the dictionary.
-        /// </summary>
-        public void AddRange(K k, ImmutableArray<V> values)
+        public Dictionary<K, ValueSet>.Enumerator GetEnumerator()
         {
-            if (values.IsEmpty)
-                return;
-
-            object item;
-            ArrayBuilder<V> arrayBuilder;
+            return IsEmpty ? s_emptyDictionary.GetEnumerator() : _dictionary.GetEnumerator();
+        }
 
-            if (!this.IsEmpty && _dictionary.TryGetValue(k, out item))
-            {
-                arrayBuilder = item as ArrayBuilder<V>;
-                if (arrayBuilder == null)
-                {
-                    // Promote from singleton V to ArrayBuilder<V>.
-                    Debug.Assert(item is V, "Item must be either a V or an ArrayBuilder<V>");
-                    arrayBuilder = new ArrayBuilder<V>(1 + values.Length);
-                    arrayBuilder.Add((V)item);
-                    arrayBuilder.AddRange(values);
-                    _dictionary[k] = arrayBuilder;
-                }
-                else
-                {
-                    arrayBuilder.AddRange(values);
-                }
-            }
-            else
-            {
-                this.EnsureDictionary();
+        IEnumerator<KeyValuePair<K, ValueSet>> IEnumerable<KeyValuePair<K, ValueSet>>.GetEnumerator()
+        {
+            return GetEnumerator();
+        }
 
-                if (values.Length == 1)
-                {
-                    _dictionary[k] = values[0];
-                }
-                else
-                {
-                    arrayBuilder = new ArrayBuilder<V>(values.Length);
-                    arrayBuilder.AddRange(values);
-                    _dictionary[k] = arrayBuilder;
-                }
-            }
+        IEnumerator IEnumerable.GetEnumerator()
+        {
+            return GetEnumerator();
         }
 
         /// <summary>
-        /// Get the number of values associated with a key.
+        /// Get all values associated with K, in the order they were added.
+        /// Returns empty read-only array if no values were present.
         /// </summary>
-        public int GetCountForKey(K k)
+        public ImmutableArray<V> this[K k]
         {
-            object item;
-            if (!this.IsEmpty && _dictionary.TryGetValue(k, out item))
+            get
             {
-                return (item as ArrayBuilder<V>)?.Count ?? 1;
-            }
+                ValueSet valueSet;
+                if (!this.IsEmpty && _dictionary.TryGetValue(k, out valueSet))
+                {
+                    Debug.Assert(valueSet.Count >= 1);
+                    return valueSet.Items;
+                }
 
-            return 0;
+                return ImmutableArray<V>.Empty;
+            }
         }
 
         /// <summary>
-        /// Returns true if one or more items with given key have been added.
+        /// Get a collection of all the keys.
         /// </summary>
-        public bool ContainsKey(K k)
+        public Dictionary<K, ValueSet>.KeyCollection Keys
         {
-            return !this.IsEmpty && _dictionary.ContainsKey(k);
+            get { return this.IsEmpty ? s_emptyDictionary.Keys : _dictionary.Keys; }
         }
 
-        /// <summary>
-        /// Get all values associated with K, in the order they were added.
-        /// Returns empty read-only array if no values were present.
-        /// </summary>
-        public ImmutableArray<V> this[K k]
+        public struct ValueSet : IEnumerable<V>
         {
-            get
+            /// <summary>
+            /// Each value is either a single V or an <see cref="ArrayBuilder{V}"/>.
+            /// Never null.
+            /// </summary>
+            private readonly object _value;
+
+            internal ValueSet(V value)
             {
-                object item;
-                if (!this.IsEmpty && _dictionary.TryGetValue(k, out item))
+                _value = value;
+            }
+
+            internal ValueSet(ArrayBuilder<V> values)
+            {
+                _value = values;
+            }
+
+            internal void Free()
+            {
+                var arrayBuilder = _value as ArrayBuilder<V>;
+                arrayBuilder?.Free();
+            }
+
+            internal V this[int index]
+            {
+                get
                 {
-                    var arrayBuilder = item as ArrayBuilder<V>;
+                    Debug.Assert(this.Count >= 1);
+
+                    var arrayBuilder = _value as ArrayBuilder<V>;
+                    if (arrayBuilder == null)
+                    {
+                        if (index == 0)
+                        {
+                            return (V)_value;
+                        }
+                        else
+                        {
+                            throw new IndexOutOfRangeException();
+                        }
+                    }
+                    else
+                    {
+                        return arrayBuilder[index];
+                    }
+                }
+            }
+
+            internal ImmutableArray<V> Items
+            {
+                get
+                {
+                    Debug.Assert(this.Count >= 1);
+
+                    var arrayBuilder = _value as ArrayBuilder<V>;
                     if (arrayBuilder == null)
                     {
                         // promote singleton to set
-                        Debug.Assert(item is V, "Item must be either a V or an ArrayBuilder<V>");
-                        return ImmutableArray.Create<V>((V)item);
+                        Debug.Assert(_value is V, "Item must be a a V");
+                        return ImmutableArray.Create<V>((V)_value);
                     }
                     else
                     {
                         return arrayBuilder.ToImmutable();
                     }
                 }
+            }
 
-                return ImmutableArray<V>.Empty;
+            internal int Count => (_value as ArrayBuilder<V>)?.Count ?? 1;
+
+            IEnumerator IEnumerable.GetEnumerator()
+            {
+                return GetEnumerator();
             }
-        }
 
-        /// <summary>
-        /// Get a collection of all the keys.
-        /// </summary>
-        public ICollection<K> Keys
-        {
-            get { return this.IsEmpty ? SpecializedCollections.EmptyCollection<K>() : _dictionary.Keys; }
+            IEnumerator<V> IEnumerable<V>.GetEnumerator()
+            {
+                return GetEnumerator();
+            }
+
+            public Enumerator GetEnumerator()
+            {
+                return new Enumerator(this);
+            }
+
+            internal ValueSet WithAddedItem(V item)
+            {
+                Debug.Assert(this.Count >= 1);
+
+                var arrayBuilder = _value as ArrayBuilder<V>;
+                if (arrayBuilder == null)
+                {
+                    // Promote from singleton V to ArrayBuilder<V>.
+                    Debug.Assert(_value is V, "_value must be a V");
+
+                    // By default we allocate array builders with a size of two.  That's to store
+                    // the single item already in _value, and to store the item we're adding.  
+                    // In general, we presume that the amount of values per key will be low, so this
+                    // means we have very little overhead when there are multiple keys per value.
+                    arrayBuilder = ArrayBuilder<V>.GetInstance(capacity: 2);
+                    arrayBuilder.Add((V)_value);
+                    arrayBuilder.Add(item);
+                }
+                else
+                {
+                    arrayBuilder.Add(item);
+                }
+
+                return new ValueSet(arrayBuilder);
+            }
+
+            public struct Enumerator : IEnumerator<V>
+            {
+                private readonly ValueSet _valueSet;
+                private readonly int _count;
+                private int _index;
+
+                public Enumerator(ValueSet valueSet)
+                {
+                    _valueSet = valueSet;
+                    _count = _valueSet.Count;
+                    Debug.Assert(_count >= 1);
+                    _index = -1;
+                }
+
+                public V Current => _valueSet[_index];
+
+                object IEnumerator.Current => Current;
+
+                public bool MoveNext()
+                {
+                    _index++;
+                    return _index < _count;
+                }
+
+                public void Reset()
+                {
+                    _index = -1;
+                }
+
+                public void Dispose()
+                {
+                }
+            }
         }
     }
-}
+}
\ No newline at end of file

src/Workspaces/Core/Portable/FindSymbols/SymbolTree/SymbolTreeInfo_Metadata.cs

@@ -6,26 +6,13 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Collections;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.FindSymbols
 {
     internal partial class SymbolTreeInfo
     {
-        private static SimplePool<MultiDictionary<string, MetadataDefinition>> s_definitionMapPool =
-            new SimplePool<MultiDictionary<string, MetadataDefinition>>(() => new MultiDictionary<string, MetadataDefinition>());
-
-        private static MultiDictionary<string, MetadataDefinition> AllocateDefinitionMap()
-        {
-            return s_definitionMapPool.Allocate();
-        }
-
-        private static void FreeDefinitionMap(MultiDictionary<string, MetadataDefinition> definitionMap)
-        {
-            definitionMap.Clear();
-            s_definitionMapPool.Free(definitionMap);
-        }
-
         /// <summary>
         /// this gives you SymbolTreeInfo for a metadata
         /// </summary>
@@ -117,7 +104,7 @@ private static void GenerateMetadataNodes(MetadataReader metadataReader, List<No
             NamespaceDefinition globalNamespace, 
             List<Node> unsortedNodes)
         {
-            var definitionMap = AllocateDefinitionMap();
+            var definitionMap = OrderPreservingMultiDictionary<string, MetadataDefinition>.GetInstance();
             try
             {
                 LookupMetadataDefinitions(reader, globalNamespace, definitionMap);
@@ -132,7 +119,7 @@ private static void GenerateMetadataNodes(MetadataReader metadataReader, List<No
             }
             finally
             {
-                FreeDefinitionMap(definitionMap);
+                definitionMap.Free();
             }
         }
 
@@ -140,7 +127,7 @@ private static void GenerateMetadataNodes(MetadataReader metadataReader, List<No
             MetadataReader reader,
             string name,
             int parentIndex,
-            MultiDictionary<string, MetadataDefinition>.ValueSet definitionsWithSameName,
+            OrderPreservingMultiDictionary<string, MetadataDefinition>.ValueSet definitionsWithSameName,
             List<Node> unsortedNodes)
         {
             var node = new Node(name, parentIndex);
@@ -148,7 +135,7 @@ private static void GenerateMetadataNodes(MetadataReader metadataReader, List<No
             unsortedNodes.Add(node);
 
             // Add all child members
-            var definitionMap = AllocateDefinitionMap();
+            var definitionMap = OrderPreservingMultiDictionary<string, MetadataDefinition>.GetInstance();
             try
             {
                 foreach (var definition in definitionsWithSameName)
@@ -166,13 +153,13 @@ private static void GenerateMetadataNodes(MetadataReader metadataReader, List<No
             }
             finally
             {
-                FreeDefinitionMap(definitionMap);
+                definitionMap.Free();
             }
         }
 
         private static void LookupMetadataDefinitions(
             MetadataReader reader, MetadataDefinition definition,
-            MultiDictionary<string, MetadataDefinition> definitionMap)
+            OrderPreservingMultiDictionary<string, MetadataDefinition> definitionMap)
         {
             switch (definition.Kind)
             {
@@ -187,7 +174,7 @@ private static void GenerateMetadataNodes(MetadataReader metadataReader, List<No
 
         private static void LookupMetadataDefinitions(
             MetadataReader reader, TypeDefinition typeDefinition,
-            MultiDictionary<string, MetadataDefinition> definitionMap)
+            OrderPreservingMultiDictionary<string, MetadataDefinition> definitionMap)
         {
             // Only bother looking for extension methods in static types.
             if ((typeDefinition.Attributes & TypeAttributes.Abstract) != 0 &&
@@ -235,7 +222,7 @@ private static void GenerateMetadataNodes(MetadataReader metadataReader, List<No
 
         private static void LookupMetadataDefinitions(
             MetadataReader reader, NamespaceDefinition namespaceDefinition,
-            MultiDictionary<string, MetadataDefinition> definitionMap)
+            OrderPreservingMultiDictionary<string, MetadataDefinition> definitionMap)
         {
             foreach (var child in namespaceDefinition.NamespaceDefinitions)
             {

src/Workspaces/Core/Portable/Workspaces.csproj

@@ -41,6 +41,9 @@
     <Compile Include="..\..\..\Compilers\Core\Portable\Collections\ImmutableArrayExtensions.cs">
       <Link>InternalUtilities\ImmutableArrayExtensions.cs</Link>
     </Compile>
+    <Compile Include="..\..\..\Compilers\Core\Portable\Collections\OrderPreservingMultiDictionary.cs">
+      <Link>Utilities\CompilerUtilities\OrderPreservingMultiDictionary.cs</Link>
+    </Compile>
     <Compile Include="..\..\..\Compilers\Core\Portable\Collections\PooledStringBuilder.cs">
       <Link>Collections\PooledStringBuilder.cs</Link>
     </Compile>