Add removeIndex function to PriorityQueue (#4070)

It's awkward to use, but lets me cancel events in an event queue.

Co-authored-by: Dmitry Atamanov <data-man@users.noreply.github.com>

1 files changed, 38 insertions, 3 deletions

diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig
index 77c76e4a70..6c56f469f9 100644
--- a/lib/std/priority_queue.zig
+++ b/lib/std/priority_queue.zig
@@ -1,8 +1,10 @@
 const std = @import("std.zig");
 const Allocator = std.mem.Allocator;
 const debug = std.debug;
+const assert = debug.assert;
 const expect = std.testing.expect;
 const expectEqual = std.testing.expectEqual;
+const expectError = std.testing.expectError;
 
 /// Priority queue for storing generic data. Initialize with `init`.
 pub fn PriorityQueue(comptime T: type) type {
@@ -77,13 +79,23 @@ pub fn PriorityQueue(comptime T: type) type {
             return if (self.len > 0) self.remove() else null;
         }
 
+        /// Remove and return the highest priority element from the
+        /// queue.
         pub fn remove(self: *Self) T {
-            const first = self.items[0];
+            return self.removeIndex(0);
+        }
+
+        /// Remove and return element at index. Indices are in the
+        /// same order as iterator, which is not necessarily priority
+        /// order.
+        pub fn removeIndex(self: *Self, index: usize) T {
+            assert(self.len > index);
             const last = self.items[self.len - 1];
-            self.items[0] = last;
+            const item = self.items[index];
+            self.items[index] = last;
             self.len -= 1;
             siftDown(self, 0);
-            return first;
+            return item;
         }
 
         /// Return the number of elements remaining in the priority
@@ -388,3 +400,26 @@ test "std.PriorityQueue: iterator" {
 
     expectEqual(@as(usize, 0), map.count());
 }
+
+test "std.PriorityQueue: remove at index" {
+    var queue = PQ.init(debug.global_allocator, lessThan);
+    defer queue.deinit();
+
+    try queue.add(3);
+    try queue.add(2);
+    try queue.add(1);
+
+    var it = queue.iterator();
+    var elem = it.next();
+    var idx: usize = 0;
+    const two_idx = while (elem != null) : (elem = it.next()) {
+        if (elem.? == 2)
+            break idx;
+        idx += 1;
+    } else unreachable;
+
+    expectEqual(queue.removeIndex(two_idx), 2);
+    expectEqual(queue.remove(), 1);
+    expectEqual(queue.remove(), 3);
+    expectEqual(queue.removeOrNull(), null);
+}