From 7deae071014237e995ec3017825f7534305ec0c4 Mon Sep 17 00:00:00 2001
From: joachimschmidt557 <joachim.schmidt557@outlook.com>
Date: Thu, 26 May 2022 18:23:07 +0200
Subject: [PATCH] std.PriorityQueue: fix missing siftUp in remove

When the replacement node is smaller than its parent, we need to sift
up instead of sifting down.
---
 lib/std/priority_queue.zig | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig
index 52d9f59951..ebc13a9974 100644
--- a/lib/std/priority_queue.zig
+++ b/lib/std/priority_queue.zig
@@ -100,7 +100,19 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
             const item = self.items[index];
             self.items[index] = last;
             self.len -= 1;
-            siftDown(self, index);
+
+            if (index == 0) {
+                siftDown(self, index);
+            } else {
+                const parent_index = ((index - 1) >> 1);
+                const parent = self.items[parent_index];
+                if (compareFn(self.context, last, parent) == .gt) {
+                    siftDown(self, index);
+                } else {
+                    siftUp(self, index);
+                }
+            }
+
             return item;
         }
 
@@ -576,6 +588,20 @@ test "std.PriorityQueue: update same max heap" {
     try expectEqual(@as(u32, 1), queue.remove());
 }
 
+test "std.PriorityQueue: siftUp in remove" {
+    var queue = PQlt.init(testing.allocator, {});
+    defer queue.deinit();
+
+    try queue.addSlice(&.{ 0, 1, 100, 2, 3, 101, 102, 4, 5, 6, 7, 103, 104, 105, 106, 8 });
+
+    _ = queue.removeIndex(std.mem.indexOfScalar(u32, queue.items[0..queue.len], 102).?);
+
+    const sorted_items = [_]u32{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 100, 101, 103, 104, 105, 106 };
+    for (sorted_items) |e| {
+        try expectEqual(e, queue.remove());
+    }
+}
+
 fn contextLessThan(context: []const u32, a: usize, b: usize) Order {
     return std.math.order(context[a], context[b]);
 }