diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig index e82b3d9cce..d710288105 100644 --- a/lib/std/priority_dequeue.zig +++ b/lib/std/priority_dequeue.zig @@ -390,8 +390,10 @@ pub fn PriorityDequeue(comptime T: type, comptime Context: type, comptime compar pub fn update(self: *Self, elem: T, new_elem: T) !void { const old_index = blk: { - for (self.items) |item, idx| { - if (compareFn(self.context, item, elem).compare(.eq)) break :blk idx; + var idx: usize = 0; + while (idx < self.len) : (idx += 1) { + const item = self.items[idx]; + if (compareFn(self.context, item, elem) == .eq) break :blk idx; } return error.ElementNotFound; }; @@ -778,6 +780,15 @@ test "std.PriorityDequeue: update same max queue" { try expectEqual(@as(u32, 1), queue.removeMax()); } +test "std.PriorityDequeue: update after remove" { + var queue = PDQ.init(testing.allocator, {}); + defer queue.deinit(); + + try queue.add(1); + try expectEqual(@as(u32, 1), queue.removeMin()); + try expectError(error.ElementNotFound, queue.update(1, 1)); +} + test "std.PriorityDequeue: iterator" { var queue = PDQ.init(testing.allocator, {}); var map = std.AutoHashMap(u32, void).init(testing.allocator); diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig index 3b65146fad..6e5e85a819 100644 --- a/lib/std/priority_queue.zig +++ b/lib/std/priority_queue.zig @@ -218,8 +218,10 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF pub fn update(self: *Self, elem: T, new_elem: T) !void { const update_index = blk: { - for (self.items) |item, idx| { - if (compareFn(self.context, item, elem).compare(.eq)) break :blk idx; + var idx: usize = 0; + while (idx < self.len) : (idx += 1) { + const item = self.items[idx]; + if (compareFn(self.context, item, elem) == .eq) break :blk idx; } return error.ElementNotFound; }; @@ -591,6 +593,15 @@ test "std.PriorityQueue: update same max heap" { try expectEqual(@as(u32, 1), queue.remove()); } +test "std.PriorityQueue: update after remove" { + var queue = PQlt.init(testing.allocator, {}); + defer queue.deinit(); + + try queue.add(1); + try expectEqual(@as(u32, 1), queue.remove()); + try expectError(error.ElementNotFound, queue.update(1, 1)); +} + test "std.PriorityQueue: siftUp in remove" { var queue = PQlt.init(testing.allocator, {}); defer queue.deinit();