Fix heap corruption issue in PriorityQueue.Remove (#107603)

Co-authored-by: Eirik Tsarpalis <eirik.tsarpalis@gmail.com>
dotnet · Sep 12, 2024 · 496a16c · 496a16c
1 parent e9633ab
commit 496a16c
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 3 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
@@ -532,16 +532,30 @@ public bool Remove(
             if (index < newSize)
             {
                 // We're removing an element from the middle of the heap.
-                // Pop the last element in the collection and sift downward from the removed index.
+                // Pop the last element in the collection and sift from the removed index.
                 (TElement Element, TPriority Priority) lastNode = nodes[newSize];
 
                 if (_comparer == null)
                 {
-                    MoveDownDefaultComparer(lastNode, index);
+                    if (Comparer<TPriority>.Default.Compare(lastNode.Priority, priority) < 0)
+                    {
+                        MoveUpDefaultComparer(lastNode, index);
+                    }
+                    else
+                    {
+                        MoveDownDefaultComparer(lastNode, index);
+                    }
                 }
                 else
                 {
-                    MoveDownCustomComparer(lastNode, index);
+                    if (_comparer.Compare(lastNode.Priority, priority) < 0)
+                    {
+                        MoveUpCustomComparer(lastNode, index);
+                    }
+                    else
+                    {
+                        MoveDownCustomComparer(lastNode, index);
+                    }
                 }
             }
 

diff --git a/src/libraries/System.Collections/tests/Generic/PriorityQueue/PriorityQueue.Tests.cs b/src/libraries/System.Collections/tests/Generic/PriorityQueue/PriorityQueue.Tests.cs
@@ -266,6 +266,33 @@ public void PriorityQueue_EmptyCollection_Remove_ShouldReturnFalse()
             Assert.Null(removedPriority);
         }
 
+        [Fact]
+        public void PriorityQueue_LargeCollection_Remove_ShouldPreserveHeapInvariant()
+        {
+            // Regression test for https://github.com/dotnet/runtime/issues/107292
+
+            PriorityQueue<int, int> queue = new();
+            for (int i = 19; i >= 0; i--)
+            {
+                queue.Enqueue(i, i);
+            }
+
+            queue.Remove(10, out int _, out int _);
+
+            List<int> sortedValues = queue.UnorderedItems
+                .OrderBy(e => e.Priority)
+                .Select(e => e.Element)
+                .ToList();
+
+            List<int> dequeuedValues = new();
+            while (queue.Count > 0)
+            {
+                dequeuedValues.Add(queue.Dequeue());
+            }
+
+            Assert.Equal(sortedValues, dequeuedValues);
+        }
+
         #region EnsureCapacity, TrimExcess
 
         [Fact]