More enumeration perf fixes

src/System.Collections.Immutable/src/System/Collections/Immutable/AllocFreeConcurrentStack.cs

@@ -9,31 +9,37 @@
 
 namespace System.Collections.Immutable
 {
-    [DebuggerDisplay("Count = {stack.Count}")]
-    internal class AllocFreeConcurrentStack<T>
+    [DebuggerDisplay("Count = {stack != null ? stack.Count : 0}")]
+    internal static class AllocFreeConcurrentStack<T>
     {
-        /// <summary>
-        /// We use locks to protect this rather than ThreadLocal{T} because in perf tests
-        /// uncontested locks are much faster than looking up thread-local storage.
-        /// </summary>
-        private readonly Stack<RefAsValueType<T>> stack = new Stack<RefAsValueType<T>>();
+        private const int MaxSize = 35;
 
-        public void TryAdd(T item)
+        [ThreadStatic]
+        private static Stack<RefAsValueType<T>> stack;
+
+        public static void TryAdd(T item)
         {
-            lock (this.stack)
+            if (stack == null)
+            {
+                stack = new Stack<RefAsValueType<T>>(MaxSize);
+            }
+
+            // Just in case we're in a scenario where an object is continually requested on one thread
+            // and returned on another, avoid unbounded growth of the stack.
+            if (stack.Count < MaxSize)
             {
-                this.stack.Push(new RefAsValueType<T>(item));
+                stack.Push(new RefAsValueType<T>(item));
             }
         }
 
-        public bool TryTake(out T item)
+        public static bool TryTake(out T item)
         {
-            lock (this.stack)
+            if (stack != null)
             {
-                int count = this.stack.Count;
+                int count = stack.Count;
                 if (count > 0)
                 {
-                    item = this.stack.Pop().Value;
+                    item = stack.Pop().Value;
                     return true;
                 }
             }

src/System.Collections.Immutable/src/System/Collections/Immutable/ImmutableList`1.cs

@@ -1505,7 +1505,7 @@ public struct Enumerator : IEnumerator<T>, ISecurePooledObjectUser
             /// A unique ID for this instance of this enumerator.
             /// Used to protect pooled objects from use after they are recycled.
             /// </summary>
-            private readonly Guid poolUserId;
+            private readonly int poolUserId;
 
             /// <summary>
             /// The starting index of the collection at which to begin enumeration.
@@ -1571,25 +1571,21 @@ internal Enumerator(IBinaryTree<T> root, Builder builder = null, int startIndex
                 this.remainingCount = this.count;
                 this.reversed = reversed;
                 this.enumeratingBuilderVersion = builder != null ? builder.Version : -1;
-                this.poolUserId = Guid.NewGuid();
+                this.poolUserId = SecureObjectPool.NewId();
+                this.stack = null;
                 if (this.count > 0)
                 {
-                    this.stack = null;
                     if (!EnumeratingStacks.TryTake(this, out this.stack))
                     {
                         this.stack = EnumeratingStacks.PrepNew(this, new Stack<RefAsValueType<IBinaryTree<T>>>(root.Height));
                     }
                 }
-                else
-                {
-                    this.stack = null;
-                }
 
                 this.Reset();
             }
 
             /// <inheritdoc/>
-            Guid ISecurePooledObjectUser.PoolUserId
+            int ISecurePooledObjectUser.PoolUserId
             {
                 get { return this.poolUserId; }
             }

src/System.Collections.Immutable/src/System/Collections/Immutable/ImmutableSortedDictionary`2.cs

@@ -957,7 +957,7 @@ public struct Enumerator : IEnumerator<KeyValuePair<TKey, TValue>>, ISecurePoole
             /// A unique ID for this instance of this enumerator.
             /// Used to protect pooled objects from use after they are recycled.
             /// </summary>
-            private readonly Guid poolUserId;
+            private readonly int poolUserId;
 
             /// <summary>
             /// The set being enumerated.
@@ -992,11 +992,14 @@ internal Enumerator(IBinaryTree<KeyValuePair<TKey, TValue>> root, Builder builde
                 this.builder = builder;
                 this.current = null;
                 this.enumeratingBuilderVersion = builder != null ? builder.Version : -1;
-                this.poolUserId = Guid.NewGuid();
+                this.poolUserId = SecureObjectPool.NewId();
                 this.stack = null;
-                if (!enumeratingStacks.TryTake(this, out this.stack))
+                if (!this.root.IsEmpty)
                 {
-                    this.stack = enumeratingStacks.PrepNew(this, new Stack<RefAsValueType<IBinaryTree<KeyValuePair<TKey, TValue>>>>(root.Height));
+                    if (!enumeratingStacks.TryTake(this, out this.stack))
+                    {
+                        this.stack = enumeratingStacks.PrepNew(this, new Stack<RefAsValueType<IBinaryTree<KeyValuePair<TKey, TValue>>>>(root.Height));
+                    }
                 }
 
                 this.Reset();
@@ -1020,7 +1023,7 @@ public KeyValuePair<TKey, TValue> Current
             }
 
             /// <inheritdoc/>
-            Guid ISecurePooledObjectUser.PoolUserId
+            int ISecurePooledObjectUser.PoolUserId
             {
                 get { return this.poolUserId; }
             }
@@ -1062,21 +1065,22 @@ public bool MoveNext()
                 this.ThrowIfDisposed();
                 this.ThrowIfChanged();
 
-                using (var stack = this.stack.Use(this))
+                if (this.stack != null)
                 {
-                    if (stack.Value.Count > 0)
-                    {
-                        IBinaryTree<KeyValuePair<TKey, TValue>> n = stack.Value.Pop().Value;
-                        this.current = n;
-                        this.PushLeft(n.Right);
-                        return true;
-                    }
-                    else
+                    using (var stack = this.stack.Use(this))
                     {
-                        this.current = null;
-                        return false;
+                        if (stack.Value.Count > 0)
+                        {
+                            IBinaryTree<KeyValuePair<TKey, TValue>> n = stack.Value.Pop().Value;
+                            this.current = n;
+                            this.PushLeft(n.Right);
+                            return true;
+                        }
                     }
                 }
+
+                this.current = null;
+                return false;
             }
 
             /// <summary>
@@ -1088,12 +1092,15 @@ public void Reset()
 
                 this.enumeratingBuilderVersion = builder != null ? builder.Version : -1;
                 this.current = null;
-                using (var stack = this.stack.Use(this))
+                if (this.stack != null)
                 {
-                    stack.Value.Clear();
-                }
+                    using (var stack = this.stack.Use(this))
+                    {
+                        stack.Value.Clear();
+                    }
 
-                this.PushLeft(this.root);
+                    this.PushLeft(this.root);
+                }
             }
 
             /// <summary>

src/System.Collections.Immutable/src/System/Collections/Immutable/ImmutableSortedSet`1.cs

@@ -1131,7 +1131,7 @@ public struct Enumerator : IEnumerator<T>, ISecurePooledObjectUser
             /// A unique ID for this instance of this enumerator.
             /// Used to protect pooled objects from use after they are recycled.
             /// </summary>
-            private readonly Guid poolUserId;
+            private readonly int poolUserId;
 
             /// <summary>
             /// A flag indicating whether this enumerator works in reverse sort order.
@@ -1181,7 +1181,7 @@ internal Enumerator(IBinaryTree<T> root, Builder builder = null, bool reverse =
                 this.current = null;
                 this.reverse = reverse;
                 this.enumeratingBuilderVersion = builder != null ? builder.Version : -1;
-                this.poolUserId = Guid.NewGuid();
+                this.poolUserId = SecureObjectPool.NewId();
                 this.stack = null;
                 if (!enumeratingStacks.TryTake(this, out this.stack))
                 {
@@ -1192,7 +1192,7 @@ internal Enumerator(IBinaryTree<T> root, Builder builder = null, bool reverse =
             }
 
             /// <inheritdoc/>
-            Guid ISecurePooledObjectUser.PoolUserId
+            int ISecurePooledObjectUser.PoolUserId
             {
                 get { return this.poolUserId; }
             }

src/System.Collections.Immutable/src/System/Collections/Immutable/SecureObjectPool.cs

@@ -11,27 +11,53 @@
 
 namespace System.Collections.Immutable
 {
+    /// <summary>
+    /// Object pooling utilities.
+    /// </summary>
+    internal class SecureObjectPool
+    {
+        /// <summary>
+        /// The ever-incrementing (and wrap-on-overflow) integer for owner id's.
+        /// </summary>
+        private static int poolUserIdCounter;
+
+        /// <summary>
+        /// The ID reserved for unassigned objects.
+        /// </summary>
+        internal const int UnassignedId = -1;
+
+        /// <summary>
+        /// Returns a new ID.
+        /// </summary>
+        internal static int NewId()
+        {
+            int result;
+            do
+            {
+                result = Interlocked.Increment(ref poolUserIdCounter);
+            }
+            while (result == UnassignedId);
+
+            return result;
+        }
+    }
+
     internal class SecureObjectPool<T, TCaller>
         where TCaller : ISecurePooledObjectUser
     {
-        private AllocFreeConcurrentStack<SecurePooledObject<T>> pool = new AllocFreeConcurrentStack<SecurePooledObject<T>>();
-
         public void TryAdd(TCaller caller, SecurePooledObject<T> item)
         {
-            lock (item)
+            // Only allow the caller to recycle this object if it is the current owner.
+            if (caller.PoolUserId == item.Owner)
             {
-                // Only allow the caller to recycle this object if it is the current owner.
-                if (caller.PoolUserId == item.Owner)
-                {
-                    item.Owner = Guid.Empty;
-                    this.pool.TryAdd(item);
-                }
+                item.Owner = SecureObjectPool.UnassignedId;
+                AllocFreeConcurrentStack<SecurePooledObject<T>>.TryAdd(item);
             }
         }
 
         public bool TryTake(TCaller caller, out SecurePooledObject<T> item)
         {
-            if (caller.PoolUserId != Guid.Empty && this.pool.TryTake(out item))
+            if (caller.PoolUserId != SecureObjectPool.UnassignedId && AllocFreeConcurrentStack<SecurePooledObject<T>>.TryTake(out item))
             {
                 item.Owner = caller.PoolUserId;
                 return true;
@@ -54,37 +80,27 @@ public SecurePooledObject<T> PrepNew(TCaller caller, T newValue)
 
     internal interface ISecurePooledObjectUser
     {
-        Guid PoolUserId { get; }
+        int PoolUserId { get; }
     }
 
     internal class SecurePooledObject<T>
     {
         private readonly T value;
-        private Guid owner;
+        private int owner;
 
         internal SecurePooledObject(T newValue)
         {
             Requires.NotNullAllowStructs(newValue, "newValue");
             this.value = newValue;
         }
 
-        internal Guid Owner
+        /// <summary>
+        /// Gets or sets the current owner of this recyclable object.
+        /// </summary>
+        internal int Owner
         {
-            get
-            {
-                lock (this)
-                {
-                    return this.owner;
-                }
-            }
-
-            set
-            {
-                lock (this)
-                {
-                    this.owner = value;
-                }
-            }
+            get { return this.owner; }
+            set { this.owner = value; }
         }
 
         internal SecurePooledObjectUser Use<TCaller>(TCaller caller)
@@ -110,7 +126,6 @@ internal struct SecurePooledObjectUser : IDisposable
             internal SecurePooledObjectUser(SecurePooledObject<T> value)
             {
                 this.value = value;
-                Monitor.Enter(value);
             }
 
             internal T Value
@@ -120,7 +135,6 @@ internal T Value
 
             public void Dispose()
             {
-                Monitor.Exit(value);
             }
         }
     }

src/System.Collections.Immutable/tests/ImmutableDictionaryTest.cs

@@ -298,9 +298,10 @@ public void GetValueOrDefaultOfConcreteType()
         [Fact]
         public void EnumeratorRecyclingMisuse()
         {
-            var collection = ImmutableDictionary.Create<int, int>();
+            var collection = ImmutableDictionary.Create<int, int>().Add(5, 3);
             var enumerator = collection.GetEnumerator();
             var enumeratorCopy = enumerator;
+            Assert.True(enumerator.MoveNext());
             Assert.False(enumerator.MoveNext());
             enumerator.Dispose();
             Assert.Throws<ObjectDisposedException>(() => enumerator.MoveNext());
@@ -315,6 +316,7 @@ public void EnumeratorRecyclingMisuse()
             // We expect that acquiring a new enumerator will use the same underlying Stack<T> object,
             // but that it will not throw exceptions for the new enumerator.
             enumerator = collection.GetEnumerator();
+            Assert.True(enumerator.MoveNext());
             Assert.False(enumerator.MoveNext());
             Assert.Throws<InvalidOperationException>(() => enumerator.Current);
             enumerator.Dispose();

src/System.Collections.Immutable/tests/ImmutableHashSetTest.cs

@@ -73,9 +73,10 @@ public void EnumeratorWithHashCollisionsTest()
         [Fact]
         public void EnumeratorRecyclingMisuse()
         {
-            var collection = ImmutableHashSet.Create<int>();
+            var collection = ImmutableHashSet.Create<int>().Add(5);
             var enumerator = collection.GetEnumerator();
             var enumeratorCopy = enumerator;
+            Assert.True(enumerator.MoveNext());
             Assert.False(enumerator.MoveNext());
             enumerator.Dispose();
             Assert.Throws<ObjectDisposedException>(() => enumerator.MoveNext());
@@ -90,6 +91,7 @@ public void EnumeratorRecyclingMisuse()
             // We expect that acquiring a new enumerator will use the same underlying Stack<T> object,
             // but that it will not throw exceptions for the new enumerator.
             enumerator = collection.GetEnumerator();
+            Assert.True(enumerator.MoveNext());
             Assert.False(enumerator.MoveNext());
             Assert.Throws<InvalidOperationException>(() => enumerator.Current);
             enumerator.Dispose();