[2018-06] Implement IEnumerable for ConditionalWeakTable (#10657)

* Implement IEnumerable for ConditionalWeakTable

* Update ConditionalWeakTable.cs

2 files changed, 116 insertions, 1 deletions

diff --git a/mcs/class/corlib/System.Runtime.CompilerServices/ConditionalWeakTable.cs b/mcs/class/corlib/System.Runtime.CompilerServices/ConditionalWeakTable.cs
index 17474df4f97..aa37ecc4249 100644
--- a/mcs/class/corlib/System.Runtime.CompilerServices/ConditionalWeakTable.cs
+++ b/mcs/class/corlib/System.Runtime.CompilerServices/ConditionalWeakTable.cs
@@ -31,6 +31,8 @@
 using System;
 using System.Collections;
 using System.Collections.Generic;
+using System.Diagnostics;
+using System.Threading;
 
 namespace System.Runtime.CompilerServices
 {
@@ -375,11 +377,109 @@ namespace System.Runtime.CompilerServices
 			}
 		}
 
+		// IEnumerable implementation was copied from CoreCLR
 		IEnumerator<KeyValuePair<TKey, TValue>> IEnumerable<KeyValuePair<TKey, TValue>>.GetEnumerator ()
 		{
-			throw new NotImplementedException ();
+			lock (_lock)
+			{
+				return size == 0 ?
+					((IEnumerable<KeyValuePair<TKey, TValue>>)Array.Empty<KeyValuePair<TKey, TValue>>()).GetEnumerator() :
+					new Enumerator(this);
+			}
 		}
 
 		IEnumerator IEnumerable.GetEnumerator () => ((IEnumerable<KeyValuePair<TKey, TValue>>)this).GetEnumerator ();
+		
+		/// <summary>Provides an enumerator for the table.</summary>
+		private sealed class Enumerator : IEnumerator<KeyValuePair<TKey, TValue>>
+		{
+			// The enumerator would ideally hold a reference to the Container and the end index within that
+			// container.  However, the safety of the CWT depends on the only reference to the Container being
+			// from the CWT itself; the Container then employs a two-phase finalization scheme, where the first
+			// phase nulls out that parent CWT's reference, guaranteeing that the second time it's finalized there
+			// can be no other existing references to it in use that would allow for concurrent usage of the
+			// native handles with finalization.  We would break that if we allowed this Enumerator to hold a
+			// reference to the Container.  Instead, the Enumerator holds a reference to the CWT rather than to
+			// the Container, and it maintains the CWT._activeEnumeratorRefCount field to track whether there
+			// are outstanding enumerators that have yet to be disposed/finalized.  If there aren't any, the CWT
+			// behaves as it normally does.  If there are, certain operations are affected, in particular resizes.
+			// Normally when the CWT is resized, it enumerates the contents of the table looking for indices that
+			// contain entries which have been collected or removed, and it frees those up, effectively moving
+			// down all subsequent entries in the container (not in the existing container, but in a replacement).
+			// This, however, would cause the enumerator's understanding of indices to break.  So, as long as
+			// there is any outstanding enumerator, no compaction is performed.
+
+			private ConditionalWeakTable<TKey, TValue> _table; // parent table, set to null when disposed
+			private int _currentIndex = -1;                    // the current index into the container
+			private KeyValuePair<TKey, TValue> _current;       // the current entry set by MoveNext and returned from Current
+
+			public Enumerator(ConditionalWeakTable<TKey, TValue> table)
+			{
+				// Store a reference to the parent table and increase its active enumerator count.
+				_table = table;
+				_currentIndex = -1;
+			}
+
+			~Enumerator() { Dispose(); }
+
+			public void Dispose()
+			{
+				// Use an interlocked operation to ensure that only one thread can get access to
+				// the _table for disposal and thus only decrement the ref count once.
+				ConditionalWeakTable<TKey, TValue> table = Interlocked.Exchange(ref _table, null);
+				if (table != null)
+				{
+					// Ensure we don't keep the last current alive unnecessarily
+					_current = default;
+
+					// Finalization is purely to decrement the ref count.  We can suppress it now.
+					GC.SuppressFinalize(this);
+				}
+			}
+
+			public bool MoveNext()
+			{
+				// Start by getting the current table.  If it's already been disposed, it will be null.
+				ConditionalWeakTable<TKey, TValue> table = _table;
+				if (table != null)
+				{
+					// Once have the table, we need to lock to synchronize with other operations on
+					// the table, like adding.
+					lock (table._lock)
+					{
+						var tombstone = GC.EPHEMERON_TOMBSTONE;
+						while (_currentIndex < table.data.Length - 1)
+						{
+							_currentIndex++;
+							var currentDataItem = table.data[_currentIndex];
+							if (currentDataItem.key != null && currentDataItem.key != tombstone)
+							{
+								_current = new KeyValuePair<TKey, TValue>((TKey)currentDataItem.key, (TValue)currentDataItem.value);
+								return true;
+							}
+						}
+					}
+				}
+
+				// Nothing more to enumerate.
+				return false;
+			}
+
+			public KeyValuePair<TKey, TValue> Current
+			{
+				get
+				{
+					if (_currentIndex < 0)
+					{
+						ThrowHelper.ThrowInvalidOperationException_InvalidOperation_EnumOpCantHappen();
+					}
+					return _current;
+				}
+			}
+
+			object IEnumerator.Current => Current;
+
+			public void Reset() { }
+		}
 	}
 }
diff --git a/mcs/class/corlib/Test/System.Runtime.CompilerServices/ConditionalWeakTableTest.cs b/mcs/class/corlib/Test/System.Runtime.CompilerServices/ConditionalWeakTableTest.cs
index 1db16372c72..a016888ca22 100644
--- a/mcs/class/corlib/Test/System.Runtime.CompilerServices/ConditionalWeakTableTest.cs
+++ b/mcs/class/corlib/Test/System.Runtime.CompilerServices/ConditionalWeakTableTest.cs
@@ -29,6 +29,7 @@
 
 using NUnit.Framework;
 using System;
+using System.Linq;
 using System.Reflection;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
@@ -486,6 +487,20 @@ namespace MonoTests.System.Runtime.CompilerServices {
 				Assert.IsTrue (table.Remove (key), "#2-" + i + "-k-" + key);
 		}
 	}
+
+	[Test]
+	public void ConditionalWeakTableEnumerable()
+	{
+		var cwt = new ConditionalWeakTable<string, string>();
+		Assert.AreEqual(0, cwt.ToArray().Length);
+		cwt.Add("test1", "foo1");
+		cwt.Add("test2", "foo2");
+		Assert.AreEqual(2, cwt.ToArray().Length);
+		cwt.Remove("test1");
+		Assert.AreEqual(1, cwt.ToArray().Length);
+		cwt.Remove("test2");
+		Assert.AreEqual(0, cwt.ToArray().Length);
+	}
 	}
 }