Address feedback.

dotnet · Oct 26, 2023 · f21a683 · stephentoub · Oct 26, 2023 · f21a683
1 parent c52e196
commit f21a683
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 5 deletions.
diff --git a/src/libraries/System.Collections/src/System/Collections/Generic/PriorityQueue.cs b/src/libraries/System.Collections/src/System/Collections/Generic/PriorityQueue.cs
@@ -513,6 +513,8 @@ public void EnqueueRange(IEnumerable<TElement> elements, TPriority priority)
         /// <remarks>
         /// The method performs a linear-time scan of every element in the heap, removing the first value found to match the <paramref name="element"/> parameter.
         /// In case of duplicate entries, what entry does get removed is non-deterministic and does not take priority into account.
+        ///
+        /// If no <paramref name="equalityComparer"/> is specified, <see cref="EqualityComparer{TElement}.Default"/> will be used instead.
         /// </remarks>
         public bool Remove(
             TElement element,
@@ -549,6 +551,7 @@ public bool Remove(
             }
 
             nodes[newSize] = default;
+            _version++;
             return true;
         }
 
@@ -865,15 +868,14 @@ private void MoveDownCustomComparer((TElement Element, TPriority Priority) node,
         private int FindIndex(TElement element, IEqualityComparer<TElement>? equalityComparer)
         {
             equalityComparer ??= EqualityComparer<TElement>.Default;
-            (TElement Element, TPriority Priority)[] nodes = _nodes;
-            int size = _size;
+            Span<(TElement Element, TPriority Priority)> nodes = _nodes.AsSpan(0, _size);
 
             // Currently the JIT doesn't optimize direct EqualityComparer<T>.Default.Equals
             // calls for reference types, so we want to cache the comparer instance instead.
             // TODO https://github.com/dotnet/runtime/issues/10050: Update if this changes in the future.
             if (typeof(TElement).IsValueType && equalityComparer == EqualityComparer<TElement>.Default)
             {
-                for (int i = 0; i < size; i++)
+                for (int i = 0; i < nodes.Length; i++)
                 {
                     if (EqualityComparer<TElement>.Default.Equals(element, nodes[i].Element))
                     {
@@ -883,7 +885,7 @@ private int FindIndex(TElement element, IEqualityComparer<TElement>? equalityCom
             }
             else
             {
-                for (int i = 0; i < size; i++)
+                for (int i = 0; i < nodes.Length; i++)
                 {
                     if (equalityComparer.Equals(element, nodes[i].Element))
                     {

diff --git a/src/libraries/System.Collections/tests/Generic/PriorityQueue/PriorityQueue.Tests.cs b/src/libraries/System.Collections/tests/Generic/PriorityQueue/PriorityQueue.Tests.cs
@@ -173,10 +173,10 @@ public void PriorityQueue_Generic_Remove_MatchingElement()
             PriorityQueue<string, int> queue = new PriorityQueue<string, int>();
             queue.EnqueueRange([("value0", 0), ("value1", 1), ("value2", 2)]);
 
-
             Assert.True(queue.Remove("value1", out string removedElement, out int removedPriority));
             Assert.Equal("value1", removedElement);
             Assert.Equal(1, removedPriority);
+            Assert.Equal(2, queue.Count);
         }
 
         [Fact]
@@ -188,6 +188,19 @@ public void PriorityQueue_Generic_Remove_MismatchElement()
             Assert.False(queue.Remove("value4", out string removedElement, out int removedPriority));
             Assert.Null(removedElement);
             Assert.Equal(0, removedPriority);
+            Assert.Equal(3, queue.Count);
+        }
+
+        [Fact]
+        public void PriorityQueue_Generic_Remove_DuplicateElement()
+        {
+            PriorityQueue<string, int> queue = new PriorityQueue<string, int>();
+            queue.EnqueueRange([("value0", 0), ("value1", 1), ("value0", 2)]);
+
+            Assert.True(queue.Remove("value0", out string removedElement, out int removedPriority));
+            Assert.Equal("value0", removedElement);
+            Assert.True(removedPriority is 0 or 2);
+            Assert.Equal(2, queue.Count);
         }
 
         [Fact]
@@ -200,6 +213,7 @@ public void PriorityQueue_Generic_Remove_CustomEqualityComparer()
             Assert.True(queue.Remove("someOtherValue1", out string removedElement, out int removedPriority, equalityComparer));
             Assert.Equal("value1", removedElement);
             Assert.Equal(1, removedPriority);
+            Assert.Equal(2, queue.Count);
         }
 
         [Fact]