From 57b66a9223d27bf4b35d3ae002e90385195216c3 Mon Sep 17 00:00:00 2001 From: Matthew Saltz Date: Mon, 16 Sep 2024 19:53:04 -0500 Subject: [PATCH] std.PriorityQueue: Convert to an 'unmanaged'-style API Resolves: #21432. Removed the allocator field from the PriorityQueue struct and adapted functions to take an allocator where needed. Updated tests accordingly. --- lib/std/priority_queue.zig | 373 +++++++++++++++++++------------------ 1 file changed, 196 insertions(+), 177 deletions(-) diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig index a5ea649c467a..05fe862962ec 100644 --- a/lib/std/priority_queue.zig +++ b/lib/std/priority_queue.zig @@ -18,33 +18,26 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF return struct { const Self = @This(); - items: []T, - cap: usize, - allocator: Allocator, + items: []T = undefined, + /// The capacity of the queue. This may be read directly, but must not + /// be modified directly. + capacity: usize = 0, context: Context, - /// Initialize and return a priority queue. - pub fn init(allocator: Allocator, context: Context) Self { - return Self{ - .items = &[_]T{}, - .cap = 0, - .allocator = allocator, - .context = context, - }; - } - /// Free memory used by the queue. - pub fn deinit(self: Self) void { - self.allocator.free(self.allocatedSlice()); + pub fn deinit(self: Self, allocator: std.mem.Allocator) void { + allocator.free(self.allocatedSlice()); } /// Insert a new element, maintaining priority. - pub fn add(self: *Self, elem: T) !void { - try self.ensureUnusedCapacity(1); - addUnchecked(self, elem); + pub fn add(self: *Self, allocator: std.mem.Allocator, elem: T) !void { + try self.ensureUnusedCapacity(allocator, 1); + self.addAssumeCapacity(elem); } - fn addUnchecked(self: *Self, elem: T) void { + /// Insert a new element, maintaining priority. Assumes there is enough + /// capacity in the queue for the additional item. + pub fn addAssumeCapacity(self: *Self, elem: T) void { self.items.len += 1; self.items[self.items.len - 1] = elem; siftUp(self, self.items.len - 1); @@ -64,10 +57,16 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF } /// Add each element in `items` to the queue. - pub fn addSlice(self: *Self, items: []const T) !void { - try self.ensureUnusedCapacity(items.len); + pub fn addSlice(self: *Self, allocator: std.mem.Allocator, items: []const T) !void { + try self.ensureUnusedCapacity(allocator, items.len); + self.addSliceAssumeCapacity(items); + } + + /// Add each element in `items` to the queue. Assumes there is enough + /// capacity in the queue for the additional items. + pub fn addSliceAssumeCapacity(self: *Self, items: []const T) void { for (items) |e| { - self.addUnchecked(e); + self.addAssumeCapacity(e); } } @@ -122,17 +121,11 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF return self.items.len; } - /// Return the number of elements that can be added to the - /// queue before more memory is allocated. - pub fn capacity(self: Self) usize { - return self.cap; - } - /// Returns a slice of all the items plus the extra capacity, whose memory /// contents are `undefined`. fn allocatedSlice(self: Self) []T { // `items.len` is the length, not the capacity. - return self.items.ptr[0..self.cap]; + return self.items.ptr[0..self.capacity]; } fn siftDown(self: *Self, target_index: usize) void { @@ -155,14 +148,14 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF self.items[index] = target_element; } - /// PriorityQueue takes ownership of the passed in slice. The slice must have been - /// allocated with `allocator`. - /// Deinitialize with `deinit`. - pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self { + /// PriorityQueue takes ownership of the passed in slice. + /// + /// Deinitialize with `deinit(allocator)`, using the allocator used to + /// allocate the passed in slice. + pub fn fromOwnedSlice(items: []T, context: Context) Self { var self = Self{ .items = items, - .cap = items.len, - .allocator = allocator, + .capacity = items.len, .context = context, }; @@ -175,40 +168,43 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF } /// Ensure that the queue can fit at least `new_capacity` items. - pub fn ensureTotalCapacity(self: *Self, new_capacity: usize) !void { - var better_capacity = self.cap; + pub fn ensureTotalCapacity(self: *Self, allocator: std.mem.Allocator, new_capacity: usize) !void { + var better_capacity = self.capacity; if (better_capacity >= new_capacity) return; while (true) { better_capacity += better_capacity / 2 + 8; if (better_capacity >= new_capacity) break; } const old_memory = self.allocatedSlice(); - const new_memory = try self.allocator.realloc(old_memory, better_capacity); + const new_memory = try allocator.realloc(old_memory, better_capacity); self.items.ptr = new_memory.ptr; - self.cap = new_memory.len; + self.capacity = new_memory.len; } /// Ensure that the queue can fit at least `additional_count` **more** item. - pub fn ensureUnusedCapacity(self: *Self, additional_count: usize) !void { - return self.ensureTotalCapacity(self.items.len + additional_count); + pub fn ensureUnusedCapacity(self: *Self, allocator: std.mem.Allocator, additional_count: usize) !void { + return self.ensureTotalCapacity( + allocator, + self.items.len + additional_count, + ); } /// Reduce allocated capacity to `new_capacity`. - pub fn shrinkAndFree(self: *Self, new_capacity: usize) void { - assert(new_capacity <= self.cap); + pub fn shrinkAndFree(self: *Self, allocator: std.mem.Allocator, new_capacity: usize) void { + assert(new_capacity <= self.capacity); // Cannot shrink to smaller than the current queue size without invalidating the heap property assert(new_capacity >= self.items.len); const old_memory = self.allocatedSlice(); - const new_memory = self.allocator.realloc(old_memory, new_capacity) catch |e| switch (e) { + const new_memory = allocator.realloc(old_memory, new_capacity) catch |e| switch (e) { error.OutOfMemory => { // no problem, capacity is still correct then. return; }, }; self.items.ptr = new_memory.ptr; - self.cap = new_memory.len; + self.capacity = new_memory.len; } pub fn update(self: *Self, elem: T, new_elem: T) !void { @@ -267,7 +263,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF print("{}, ", .{e}); } print("len: {} ", .{self.items.len}); - print("capacity: {}", .{self.cap}); + print("capacity: {}", .{self.capacity}); print(" }}\n", .{}); } }; @@ -286,15 +282,16 @@ const PQlt = PriorityQueue(u32, void, lessThan); const PQgt = PriorityQueue(u32, void, greaterThan); test "add and remove min heap" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); - - try queue.add(54); - try queue.add(12); - try queue.add(7); - try queue.add(23); - try queue.add(25); - try queue.add(13); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); + + try queue.add(allocator, 54); + try queue.add(allocator, 12); + try queue.add(allocator, 7); + try queue.add(allocator, 23); + try queue.add(allocator, 25); + try queue.add(allocator, 13); try expectEqual(@as(u32, 7), queue.remove()); try expectEqual(@as(u32, 12), queue.remove()); try expectEqual(@as(u32, 13), queue.remove()); @@ -304,15 +301,16 @@ test "add and remove min heap" { } test "add and remove same min heap" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); - - try queue.add(1); - try queue.add(1); - try queue.add(2); - try queue.add(2); - try queue.add(1); - try queue.add(1); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); + + try queue.add(allocator, 1); + try queue.add(allocator, 1); + try queue.add(allocator, 2); + try queue.add(allocator, 2); + try queue.add(allocator, 1); + try queue.add(allocator, 1); try expectEqual(@as(u32, 1), queue.remove()); try expectEqual(@as(u32, 1), queue.remove()); try expectEqual(@as(u32, 1), queue.remove()); @@ -322,43 +320,46 @@ test "add and remove same min heap" { } test "removeOrNull on empty" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); try expect(queue.removeOrNull() == null); } test "edge case 3 elements" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); - try queue.add(9); - try queue.add(3); - try queue.add(2); + try queue.add(allocator, 9); + try queue.add(allocator, 3); + try queue.add(allocator, 2); try expectEqual(@as(u32, 2), queue.remove()); try expectEqual(@as(u32, 3), queue.remove()); try expectEqual(@as(u32, 9), queue.remove()); } test "peek" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); try expect(queue.peek() == null); - try queue.add(9); - try queue.add(3); - try queue.add(2); + try queue.add(allocator, 9); + try queue.add(allocator, 3); + try queue.add(allocator, 2); try expectEqual(@as(u32, 2), queue.peek().?); try expectEqual(@as(u32, 2), queue.peek().?); } test "sift up with odd indices" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); + const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; - for (items) |e| { - try queue.add(e); - } + try queue.addSlice(allocator, items[0..]); const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 }; for (sorted_items) |e| { @@ -367,10 +368,12 @@ test "sift up with odd indices" { } test "addSlice" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); + const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; - try queue.addSlice(items[0..]); + try queue.addSlice(allocator, items[0..]); const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 }; for (sorted_items) |e| { @@ -379,19 +382,21 @@ test "addSlice" { } test "fromOwnedSlice trivial case 0" { + const allocator = testing.allocator; const items = [0]u32{}; - const queue_items = try testing.allocator.dupe(u32, &items); - var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {}); - defer queue.deinit(); + const queue_items = try allocator.dupe(u32, &items); + var queue = PQlt.fromOwnedSlice(queue_items[0..], {}); + defer queue.deinit(allocator); try expectEqual(@as(usize, 0), queue.count()); try expect(queue.removeOrNull() == null); } test "fromOwnedSlice trivial case 1" { + const allocator = testing.allocator; const items = [1]u32{1}; - const queue_items = try testing.allocator.dupe(u32, &items); - var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {}); - defer queue.deinit(); + const queue_items = try allocator.dupe(u32, &items); + var queue = PQlt.fromOwnedSlice(queue_items[0..], {}); + defer queue.deinit(allocator); try expectEqual(@as(usize, 1), queue.count()); try expectEqual(items[0], queue.remove()); @@ -399,10 +404,11 @@ test "fromOwnedSlice trivial case 1" { } test "fromOwnedSlice" { + const allocator = testing.allocator; const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; - const heap_items = try testing.allocator.dupe(u32, items[0..]); - var queue = PQlt.fromOwnedSlice(testing.allocator, heap_items[0..], {}); - defer queue.deinit(); + const heap_items = try allocator.dupe(u32, items[0..]); + var queue = PQlt.fromOwnedSlice(heap_items[0..], {}); + defer queue.deinit(allocator); const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 }; for (sorted_items) |e| { @@ -411,15 +417,16 @@ test "fromOwnedSlice" { } test "add and remove max heap" { - var queue = PQgt.init(testing.allocator, {}); - defer queue.deinit(); - - try queue.add(54); - try queue.add(12); - try queue.add(7); - try queue.add(23); - try queue.add(25); - try queue.add(13); + const allocator = testing.allocator; + var queue: PQgt = .{ .context = {} }; + defer queue.deinit(allocator); + + try queue.add(allocator, 54); + try queue.add(allocator, 12); + try queue.add(allocator, 7); + try queue.add(allocator, 23); + try queue.add(allocator, 25); + try queue.add(allocator, 13); try expectEqual(@as(u32, 54), queue.remove()); try expectEqual(@as(u32, 25), queue.remove()); try expectEqual(@as(u32, 23), queue.remove()); @@ -429,15 +436,16 @@ test "add and remove max heap" { } test "add and remove same max heap" { - var queue = PQgt.init(testing.allocator, {}); - defer queue.deinit(); - - try queue.add(1); - try queue.add(1); - try queue.add(2); - try queue.add(2); - try queue.add(1); - try queue.add(1); + const allocator = testing.allocator; + var queue: PQgt = .{ .context = {} }; + defer queue.deinit(allocator); + + try queue.add(allocator, 1); + try queue.add(allocator, 1); + try queue.add(allocator, 2); + try queue.add(allocator, 2); + try queue.add(allocator, 1); + try queue.add(allocator, 1); try expectEqual(@as(u32, 2), queue.remove()); try expectEqual(@as(u32, 2), queue.remove()); try expectEqual(@as(u32, 1), queue.remove()); @@ -447,16 +455,16 @@ test "add and remove same max heap" { } test "iterator" { - var queue = PQlt.init(testing.allocator, {}); - var map = std.AutoHashMap(u32, void).init(testing.allocator); - defer { - queue.deinit(); - map.deinit(); - } + const allocator = testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); + + var map = std.AutoHashMap(u32, void).init(allocator); + defer map.deinit(); const items = [_]u32{ 54, 12, 7, 23, 25, 13 }; for (items) |e| { - _ = try queue.add(e); + _ = try queue.add(allocator, e); try map.put(e, {}); } @@ -469,13 +477,12 @@ test "iterator" { } test "remove at index" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); const items = [_]u32{ 2, 1, 8, 9, 3, 4, 5 }; - for (items) |e| { - _ = try queue.add(e); - } + try queue.addSlice(allocator, items[0..]); var it = queue.iterator(); var idx: usize = 0; @@ -495,8 +502,9 @@ test "remove at index" { } test "iterator while empty" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); var it = queue.iterator(); @@ -504,20 +512,21 @@ test "iterator while empty" { } test "shrinkAndFree" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); - try queue.ensureTotalCapacity(4); - try expect(queue.capacity() >= 4); + try queue.ensureTotalCapacity(allocator, 4); + try expect(queue.capacity >= 4); - try queue.add(1); - try queue.add(2); - try queue.add(3); - try expect(queue.capacity() >= 4); + try queue.add(allocator, 1); + try queue.add(allocator, 2); + try queue.add(allocator, 3); + try expect(queue.capacity >= 4); try expectEqual(@as(usize, 3), queue.count()); - queue.shrinkAndFree(3); - try expectEqual(@as(usize, 3), queue.capacity()); + queue.shrinkAndFree(allocator, 3); + try expectEqual(@as(usize, 3), queue.capacity); try expectEqual(@as(usize, 3), queue.count()); try expectEqual(@as(u32, 1), queue.remove()); @@ -527,12 +536,13 @@ test "shrinkAndFree" { } test "update min heap" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); - try queue.add(55); - try queue.add(44); - try queue.add(11); + try queue.add(allocator, 55); + try queue.add(allocator, 44); + try queue.add(allocator, 11); try queue.update(55, 5); try queue.update(44, 4); try queue.update(11, 1); @@ -542,13 +552,14 @@ test "update min heap" { } test "update same min heap" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); - - try queue.add(1); - try queue.add(1); - try queue.add(2); - try queue.add(2); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); + + try queue.add(allocator, 1); + try queue.add(allocator, 1); + try queue.add(allocator, 2); + try queue.add(allocator, 2); try queue.update(1, 5); try queue.update(2, 4); try expectEqual(@as(u32, 1), queue.remove()); @@ -558,12 +569,13 @@ test "update same min heap" { } test "update max heap" { - var queue = PQgt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = std.testing.allocator; + var queue: PQgt = .{ .context = {} }; + defer queue.deinit(allocator); - try queue.add(55); - try queue.add(44); - try queue.add(11); + try queue.add(allocator, 55); + try queue.add(allocator, 44); + try queue.add(allocator, 11); try queue.update(55, 5); try queue.update(44, 1); try queue.update(11, 4); @@ -573,13 +585,14 @@ test "update max heap" { } test "update same max heap" { - var queue = PQgt.init(testing.allocator, {}); - defer queue.deinit(); - - try queue.add(1); - try queue.add(1); - try queue.add(2); - try queue.add(2); + const allocator = std.testing.allocator; + var queue: PQgt = .{ .context = {} }; + defer queue.deinit(allocator); + + try queue.add(allocator, 1); + try queue.add(allocator, 1); + try queue.add(allocator, 2); + try queue.add(allocator, 2); try queue.update(1, 5); try queue.update(2, 4); try expectEqual(@as(u32, 5), queue.remove()); @@ -589,19 +602,24 @@ test "update same max heap" { } test "update after remove" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); - try queue.add(1); + try queue.add(allocator, 1); try expectEqual(@as(u32, 1), queue.remove()); try expectError(error.ElementNotFound, queue.update(1, 1)); } test "siftUp in remove" { - var queue = PQlt.init(testing.allocator, {}); - defer queue.deinit(); + const allocator = std.testing.allocator; + var queue: PQlt = .{ .context = {} }; + defer queue.deinit(allocator); - try queue.addSlice(&.{ 0, 1, 100, 2, 3, 101, 102, 4, 5, 6, 7, 103, 104, 105, 106, 8 }); + try queue.addSlice( + allocator, + &.{ 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.count()], 102).?); @@ -618,18 +636,19 @@ fn contextLessThan(context: []const u32, a: usize, b: usize) Order { const CPQlt = PriorityQueue(usize, []const u32, contextLessThan); test "add and remove min heap with context comparator" { - const context = [_]u32{ 5, 3, 4, 2, 2, 8, 0 }; - - var queue = CPQlt.init(testing.allocator, context[0..]); - defer queue.deinit(); - - try queue.add(0); - try queue.add(1); - try queue.add(2); - try queue.add(3); - try queue.add(4); - try queue.add(5); - try queue.add(6); + const allocator = std.testing.allocator; + var queue: CPQlt = .{ + .context = &[_]u32{ 5, 3, 4, 2, 2, 8, 0 }, + }; + defer queue.deinit(allocator); + + try queue.add(allocator, 0); + try queue.add(allocator, 1); + try queue.add(allocator, 2); + try queue.add(allocator, 3); + try queue.add(allocator, 4); + try queue.add(allocator, 5); + try queue.add(allocator, 6); try expectEqual(@as(usize, 6), queue.remove()); try expectEqual(@as(usize, 4), queue.remove()); try expectEqual(@as(usize, 3), queue.remove());