Commit f8ea292d09

Frechdachs <Frechdachs@users.noreply.github.com>
2022-12-13 20:55:41
std: Fix update() method in PriorityQueue and PriorityDequeue (#13908)
Previously the update() method would iterate over its capacity, which may contain uninitialized memory or already removed elements.
1 parent b5222f8
Changed files (2)
lib
std
priority_dequeue.zig
priority_queue.zig
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);
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();