From a727a508cd18d4af1773841e00630f5b47e2b95f Mon Sep 17 00:00:00 2001 From: Zander Khan Date: Sat, 16 Jan 2021 12:01:06 +0000 Subject: std: Add Priority Dequeue --- lib/std/priority_dequeue.zig | 879 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 879 insertions(+) create mode 100644 lib/std/priority_dequeue.zig (limited to 'lib/std/priority_dequeue.zig') diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig new file mode 100644 index 0000000000..a742b27050 --- /dev/null +++ b/lib/std/priority_dequeue.zig @@ -0,0 +1,879 @@ +const std = @import("std"); +const Allocator = std.mem.Allocator; +const sort = std.sort; +const assert = std.debug.assert; +const warn = std.debug.warn; +const testing = std.testing; +const expect = testing.expect; +const expectEqual = testing.expectEqual; +const expectError = testing.expectError; + +/// Priority Dequeue for storing generic data. Initialize with `init`. +pub fn PriorityDequeue(comptime T: type) type { + return struct { + const Self = @This(); + + items: []T, + len: usize, + allocator: *Allocator, + lessThanFn: fn (a: T, b: T) bool, + + /// Initialize and return a new dequeue. Provide `lessThanFn` + /// that returns `true` when its first argument should + /// get min-popped before its second argument. For example, + /// to make `popMin` return the minimum value, provide + /// + /// `fn lessThanFn(a: T, b: T) bool { return a < b; }` + pub fn init(allocator: *Allocator, lessThanFn: fn (T, T) bool) Self { + return Self{ + .items = &[_]T{}, + .len = 0, + .allocator = allocator, + .lessThanFn = lessThanFn, + }; + } + + fn lessThan(self: Self, a: T, b: T) bool { + return self.lessThanFn(a, b); + } + + fn greaterThan(self: Self, a: T, b: T) bool { + return self.lessThanFn(b, a); + } + + /// Free memory used by the dequeue. + pub fn deinit(self: Self) void { + self.allocator.free(self.items); + } + + /// Insert a new element, maintaining priority. + pub fn add(self: *Self, elem: T) !void { + try ensureCapacity(self, self.len + 1); + addUnchecked(self, elem); + } + + /// Add each element in `items` to the dequeue. + pub fn addSlice(self: *Self, items: []const T) !void { + try self.ensureCapacity(self.len + items.len); + for (items) |e| { + self.addUnchecked(e); + } + } + + fn addUnchecked(self: *Self, elem: T) void { + self.items[self.len] = elem; + + if (self.len > 0) { + const start = self.getStartForSiftUp(elem, self.len); + self.siftUp(start); + } + + self.len += 1; + } + + fn isMinLayer(index: usize) bool { + // In the min-max heap structure: + // The first element is on a min layer; + // next two are on a max layer; + // next four are on a min layer, and so on. + const leading_zeros = @clz(usize, index + 1); + const highest_set_bit = 63 - leading_zeros; + return (highest_set_bit & 1) == 0; + } + + fn nextIsMinLayer(self: Self) bool { + return isMinLayer(self.len); + } + + const StartIndexAndLayer = struct { + index: usize, + min_layer: bool, + }; + + fn getStartForSiftUp(self: Self, child: T, index: usize) StartIndexAndLayer { + var child_index = index; + var parent_index = parentIndex(child_index); + const parent = self.items[parent_index]; + + const min_layer = self.nextIsMinLayer(); + if ((min_layer and self.greaterThan(child, parent)) or (!min_layer and self.lessThan(child, parent))) { + // We must swap the item with it's parent if it is on the "wrong" layer + self.items[parent_index] = child; + self.items[child_index] = parent; + return .{ + .index = parent_index, + .min_layer = !min_layer, + }; + } else { + return .{ + .index = child_index, + .min_layer = min_layer, + }; + } + } + + fn siftUp(self: *Self, start: StartIndexAndLayer) void { + if (start.min_layer) { + doSiftUp(self, start.index, lessThan); + } else { + doSiftUp(self, start.index, greaterThan); + } + } + + fn doSiftUp(self: *Self, start_index: usize, compare: fn (Self, T, T) bool) void { + var child_index = start_index; + while (child_index > 2) { + var grandparent_index = grandparentIndex(child_index); + const child = self.items[child_index]; + const grandparent = self.items[grandparent_index]; + + // If the grandparent is already better, we have gone as far as we need to + if (!compare(self.*, child, grandparent)) break; + + // Otherwise swap the item with it's grandparent + self.items[grandparent_index] = child; + self.items[child_index] = grandparent; + child_index = grandparent_index; + } + } + + /// Look at the smallest element in the dequeue. Returns + /// `null` if empty. + pub fn peekMin(self: *Self) ?T { + return if (self.len > 0) self.items[0] else null; + } + + /// Look at the largest element in the dequeue. Returns + /// `null` if empty. + pub fn peekMax(self: *Self) ?T { + if (self.len == 0) return null; + if (self.len == 1) return self.items[0]; + if (self.len == 2) return self.items[1]; + return self.bestItemAtIndices(1, 2, greaterThan).item; + } + + fn maxIndex(self: Self) ?usize { + if (self.len == 0) return null; + if (self.len == 1) return 0; + if (self.len == 2) return 1; + return self.bestItemAtIndices(1, 2, greaterThan).index; + } + + /// Pop the smallest element from the dequeue. Returns + /// `null` if empty. + pub fn removeMinOrNull(self: *Self) ?T { + return if (self.len > 0) self.removeMin() else null; + } + + /// Remove and return the smallest element from the + /// dequeue. + pub fn removeMin(self: *Self) T { + return self.removeIndex(0); + } + + /// Pop the largest element from the dequeue. Returns + /// `null` if empty. + pub fn removeMaxOrNull(self: *Self) ?T { + return if (self.len > 0) self.removeMax() else null; + } + + /// Remove and return the largest element from the + /// dequeue. + pub fn removeMax(self: *Self) T { + return self.removeIndex(self.maxIndex().?); + } + + /// 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 { + const item = self.items[index]; + const last = self.items[self.len - 1]; + + self.items[index] = last; + self.len -= 1; + siftDown(self, index); + + return item; + } + + fn siftDown(self: *Self, index: usize) void { + if (isMinLayer(index)) { + self.doSiftDown(index, lessThan); + } else { + self.doSiftDown(index, greaterThan); + } + } + + fn doSiftDown(self: *Self, start_index: usize, compare: fn (Self, T, T) bool) void { + var index = start_index; + const half = self.len >> 1; + while (true) { + const first_grandchild_index = firstGrandchildIndex(index); + const last_grandchild_index = first_grandchild_index + 3; + + const elem = self.items[index]; + + if (last_grandchild_index < self.len) { + // All four grandchildren exist + const index2 = first_grandchild_index + 1; + const index3 = index2 + 1; + + // Find the best grandchild + const best_left = self.bestItemAtIndices(first_grandchild_index, index2, compare); + const best_right = self.bestItemAtIndices(index3, last_grandchild_index, compare); + const best_grandchild = self.bestItem(best_left, best_right, compare); + + // If the item is better than it's best grandchild, we are done + if (compare(self.*, elem, best_grandchild.item) or elem == best_grandchild.item) return; + + // Otherwise, swap them + self.items[best_grandchild.index] = elem; + self.items[index] = best_grandchild.item; + index = best_grandchild.index; + + // We might need to swap the element with it's parent + self.swapIfParentIsBetter(elem, index, compare); + } else { + // The children or grandchildren are the last layer + const first_child_index = firstChildIndex(index); + if (first_child_index > self.len) return; + + const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, compare); + + // If the best descendant is still larger, we are done + if (compare(self.*, elem, best_descendent.item) or elem == best_descendent.item) return; + + // Otherwise swap them + self.items[best_descendent.index] = elem; + self.items[index] = best_descendent.item; + index = best_descendent.index; + + // If we didn't swap a grandchild, we are done + if (index < first_grandchild_index) return; + + // We might need to swap the element with it's parent + self.swapIfParentIsBetter(elem, index, compare); + return; + } + + // If we are now in the last layer, we are done + if (index >= half) return; + } + } + + fn swapIfParentIsBetter(self: *Self, child: T, child_index: usize, compare: fn (Self, T, T) bool) void { + const parent_index = parentIndex(child_index); + const parent = self.items[parent_index]; + + if (compare(self.*, parent, child)) { + self.items[parent_index] = child; + self.items[child_index] = parent; + } + } + + const ItemAndIndex = struct { + item: T, + index: usize, + }; + + fn getItem(self: Self, index: usize) ItemAndIndex { + return .{ + .item = self.items[index], + .index = index, + }; + } + + fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, compare: fn (Self, T, T) bool) ItemAndIndex { + if (compare(self, item1.item, item2.item)) { + return item1; + } else { + return item2; + } + } + + fn bestItemAtIndices(self: Self, index1: usize, index2: usize, compare: fn (Self, T, T) bool) ItemAndIndex { + var item1 = self.getItem(index1); + var item2 = self.getItem(index2); + return self.bestItem(item1, item2, compare); + } + + fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, compare: fn (Self, T, T) bool) ItemAndIndex { + const second_child_index = first_child_index + 1; + if (first_grandchild_index >= self.len) { + // No grandchildren, find the best child (second may not exist) + if (second_child_index >= self.len) { + return .{ + .item = self.items[first_child_index], + .index = first_child_index, + }; + } else { + return self.bestItemAtIndices(first_child_index, second_child_index, compare); + } + } + + const second_grandchild_index = first_grandchild_index + 1; + if (second_grandchild_index >= self.len) { + // One grandchild, so we know there is a second child. Compare first grandchild and second child + return self.bestItemAtIndices(first_grandchild_index, second_child_index, compare); + } + + const best_left_grandchild_index = self.bestItemAtIndices(first_grandchild_index, second_grandchild_index, compare).index; + const third_grandchild_index = second_grandchild_index + 1; + if (third_grandchild_index >= self.len) { + // Two grandchildren, and we know the best. Compare this to second child. + return self.bestItemAtIndices(best_left_grandchild_index, second_child_index, compare); + } else { + // Three grandchildren, compare the min of the first two with the third + return self.bestItemAtIndices(best_left_grandchild_index, third_grandchild_index, compare); + } + } + + /// Return the number of elements remaining in the heap + pub fn count(self: Self) usize { + return self.len; + } + + /// Return the number of elements that can be added to the + /// dequeue before more memory is allocated. + pub fn capacity(self: Self) usize { + return self.items.len; + } + + /// Heap takes ownership of the passed in slice. The slice must have been + /// allocated with `allocator`. + /// De-initialize with `deinit`. + pub fn fromOwnedSlice(allocator: *Allocator, lessThanFn: fn (T, T) bool, items: []T) Self { + var dequeue = Self{ + .items = items, + .len = items.len, + .allocator = allocator, + .lessThanFn = lessThanFn, + }; + const half = (dequeue.len >> 1) - 1; + var i: usize = 0; + while (i <= half) : (i += 1) { + const index = half - i; + dequeue.siftDown(index); + } + return dequeue; + } + + pub fn ensureCapacity(self: *Self, 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; + } + self.items = try self.allocator.realloc(self.items, better_capacity); + } + + pub fn resize(self: *Self, new_len: usize) !void { + try self.ensureCapacity(new_len); + self.len = new_len; + } + + pub fn shrink(self: *Self, new_len: usize) void { + // TODO take advantage of the new realloc semantics + assert(new_len <= self.len); + self.len = new_len; + } + + pub fn update(self: *Self, elem: T, new_elem: T) !void { + var old_index: usize = std.mem.indexOfScalar(T, self.items, elem) orelse return error.ElementNotFound; + _ = self.removeIndex(old_index); + self.addUnchecked(new_elem); + } + + pub const Iterator = struct { + heap: *PriorityDequeue(T), + count: usize, + + pub fn next(it: *Iterator) ?T { + if (it.count >= it.heap.len) return null; + const out = it.count; + it.count += 1; + return it.heap.items[out]; + } + + pub fn reset(it: *Iterator) void { + it.count = 0; + } + }; + + /// Return an iterator that walks the heap without consuming + /// it. Invalidated if the heap is modified. + pub fn iterator(self: *Self) Iterator { + return Iterator{ + .heap = self, + .count = 0, + }; + } + + fn dump(self: *Self) void { + warn("{{ ", .{}); + warn("items: ", .{}); + for (self.items) |e, i| { + if (i >= self.len) break; + warn("{}, ", .{e}); + } + warn("array: ", .{}); + for (self.items) |e, i| { + warn("{}, ", .{e}); + } + warn("len: {} ", .{self.len}); + warn("capacity: {}", .{self.capacity()}); + warn(" }}\n", .{}); + } + + fn parentIndex(index: usize) usize { + return (index - 1) >> 1; + } + + fn grandparentIndex(index: usize) usize { + return parentIndex(parentIndex(index)); + } + + fn firstChildIndex(index: usize) usize { + return (index << 1) + 1; + } + + fn firstGrandchildIndex(index: usize) usize { + return firstChildIndex(firstChildIndex(index)); + } + }; +} + +fn lessThanComparison(a: u32, b: u32) bool { + return a < b; +} + +const Heap = PriorityDequeue(u32); + +test "std.PriorityDequeue: add and remove min" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + try heap.add(54); + try heap.add(12); + try heap.add(7); + try heap.add(23); + try heap.add(25); + try heap.add(13); + + expectEqual(@as(u32, 7), heap.removeMin()); + expectEqual(@as(u32, 12), heap.removeMin()); + expectEqual(@as(u32, 13), heap.removeMin()); + expectEqual(@as(u32, 23), heap.removeMin()); + expectEqual(@as(u32, 25), heap.removeMin()); + expectEqual(@as(u32, 54), heap.removeMin()); +} + +test "std.PriorityDequeue: add and remove max" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + try heap.add(54); + try heap.add(12); + try heap.add(7); + try heap.add(23); + try heap.add(25); + try heap.add(13); + + expectEqual(@as(u32, 54), heap.removeMax()); + expectEqual(@as(u32, 25), heap.removeMax()); + expectEqual(@as(u32, 23), heap.removeMax()); + expectEqual(@as(u32, 13), heap.removeMax()); + expectEqual(@as(u32, 12), heap.removeMax()); + expectEqual(@as(u32, 7), heap.removeMax()); +} + +test "std.PriorityDequeue: add and remove same min" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + try heap.add(1); + try heap.add(1); + try heap.add(2); + try heap.add(2); + try heap.add(1); + try heap.add(1); + + expectEqual(@as(u32, 1), heap.removeMin()); + expectEqual(@as(u32, 1), heap.removeMin()); + expectEqual(@as(u32, 1), heap.removeMin()); + expectEqual(@as(u32, 1), heap.removeMin()); + expectEqual(@as(u32, 2), heap.removeMin()); + expectEqual(@as(u32, 2), heap.removeMin()); +} + +test "std.PriorityDequeue: add and remove same max" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + try heap.add(1); + try heap.add(1); + try heap.add(2); + try heap.add(2); + try heap.add(1); + try heap.add(1); + + expectEqual(@as(u32, 2), heap.removeMax()); + expectEqual(@as(u32, 2), heap.removeMax()); + expectEqual(@as(u32, 1), heap.removeMax()); + expectEqual(@as(u32, 1), heap.removeMax()); + expectEqual(@as(u32, 1), heap.removeMax()); + expectEqual(@as(u32, 1), heap.removeMax()); +} + +test "std.PriorityDequeue: removeOrNull empty" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + expect(heap.removeMinOrNull() == null); + expect(heap.removeMaxOrNull() == null); +} + +test "std.PriorityDequeue: edge case 3 elements" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + try heap.add(9); + try heap.add(3); + try heap.add(2); + + expectEqual(@as(u32, 2), heap.removeMin()); + expectEqual(@as(u32, 3), heap.removeMin()); + expectEqual(@as(u32, 9), heap.removeMin()); +} + +test "std.PriorityDequeue: edge case 3 elements max" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + try heap.add(9); + try heap.add(3); + try heap.add(2); + + expectEqual(@as(u32, 9), heap.removeMax()); + expectEqual(@as(u32, 3), heap.removeMax()); + expectEqual(@as(u32, 2), heap.removeMax()); +} + +test "std.PriorityDequeue: peekMin" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + expect(heap.peekMin() == null); + + try heap.add(9); + try heap.add(3); + try heap.add(2); + + expect(heap.peekMin().? == 2); + expect(heap.peekMin().? == 2); +} + +test "std.PriorityDequeue: peekMax" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + expect(heap.peekMin() == null); + + try heap.add(9); + try heap.add(3); + try heap.add(2); + + expect(heap.peekMax().? == 9); + expect(heap.peekMax().? == 9); +} + +test "std.PriorityDequeue: sift up with odd indices" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + 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 heap.add(e); + } + + 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| { + expectEqual(e, heap.removeMin()); + } +} + +test "std.PriorityDequeue: sift up with odd indices" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + 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 heap.add(e); + } + + const sorted_items = [_]u32{ 25, 24, 24, 22, 21, 16, 15, 15, 14, 13, 12, 11, 7, 7, 6, 5, 2, 1 }; + for (sorted_items) |e| { + expectEqual(e, heap.removeMax()); + } +} + +test "std.PriorityDequeue: addSlice min" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; + try heap.addSlice(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| { + expectEqual(e, heap.removeMin()); + } +} + +test "std.PriorityDequeue: addSlice max" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; + try heap.addSlice(items[0..]); + + const sorted_items = [_]u32{ 25, 24, 24, 22, 21, 16, 15, 15, 14, 13, 12, 11, 7, 7, 6, 5, 2, 1 }; + for (sorted_items) |e| { + expectEqual(e, heap.removeMax()); + } +} + +test "std.PriorityDequeue: fromOwnedSlice" { + 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 heap = Heap.fromOwnedSlice(testing.allocator, lessThanComparison, heap_items[0..]); + defer heap.deinit(); + + 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| { + expectEqual(e, heap.removeMin()); + } +} + +test "std.PriorityDequeue: update min heap" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + try heap.add(55); + try heap.add(44); + try heap.add(11); + try heap.update(55, 5); + try heap.update(44, 4); + try heap.update(11, 1); + expectEqual(@as(u32, 1), heap.removeMin()); + expectEqual(@as(u32, 4), heap.removeMin()); + expectEqual(@as(u32, 5), heap.removeMin()); +} + +test "std.PriorityDequeue: update same min heap" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + try heap.add(1); + try heap.add(1); + try heap.add(2); + try heap.add(2); + try heap.update(1, 5); + try heap.update(2, 4); + expectEqual(@as(u32, 1), heap.removeMin()); + expectEqual(@as(u32, 2), heap.removeMin()); + expectEqual(@as(u32, 4), heap.removeMin()); + expectEqual(@as(u32, 5), heap.removeMin()); +} + +test "std.PriorityDequeue: update max heap" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + try heap.add(55); + try heap.add(44); + try heap.add(11); + try heap.update(55, 5); + try heap.update(44, 1); + try heap.update(11, 4); + + expectEqual(@as(u32, 5), heap.removeMax()); + expectEqual(@as(u32, 4), heap.removeMax()); + expectEqual(@as(u32, 1), heap.removeMax()); +} + +test "std.PriorityDequeue: update same max heap" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + try heap.add(1); + try heap.add(1); + try heap.add(2); + try heap.add(2); + try heap.update(1, 5); + try heap.update(2, 4); + expectEqual(@as(u32, 5), heap.removeMax()); + expectEqual(@as(u32, 4), heap.removeMax()); + expectEqual(@as(u32, 2), heap.removeMax()); + expectEqual(@as(u32, 1), heap.removeMax()); +} + +test "std.PriorityDequeue: iterator" { + var heap = Heap.init(testing.allocator, lessThanComparison); + var map = std.AutoHashMap(u32, void).init(testing.allocator); + defer { + heap.deinit(); + map.deinit(); + } + + const items = [_]u32{ 54, 12, 7, 23, 25, 13 }; + for (items) |e| { + _ = try heap.add(e); + _ = try map.put(e, {}); + } + + var it = heap.iterator(); + while (it.next()) |e| { + _ = map.remove(e); + } + + expectEqual(@as(usize, 0), map.count()); +} + +test "std.PriorityDequeue: remove at index" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + try heap.add(3); + try heap.add(2); + try heap.add(1); + + var it = heap.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(heap.removeIndex(two_idx), 2); + expectEqual(heap.removeMin(), 1); + expectEqual(heap.removeMin(), 3); + expectEqual(heap.removeMinOrNull(), null); +} + +test "std.PriorityDequeue: iterator while empty" { + var heap = Heap.init(testing.allocator, lessThanComparison); + defer heap.deinit(); + + var it = heap.iterator(); + + expectEqual(it.next(), null); +} + +test "std.PriorityDequeue: fuzz testing min" { + var prng = std.rand.DefaultPrng.init(0x12345678); + + const test_case_count = 100; + const heap_size = 1_000; + + var i: usize = 0; + while (i < test_case_count) : (i += 1) { + try fuzzTestMin(&prng.random, heap_size); + } +} + +fn fuzzTestMin(rng: *std.rand.Random, comptime heap_size: usize) !void { + const allocator = testing.allocator; + const items = try generateRandomSlice(allocator, rng, heap_size); + + var heap = Heap.fromOwnedSlice(allocator, lessThanComparison, items); + defer heap.deinit(); + + var last_removed: ?u32 = null; + while (heap.removeMinOrNull()) |next| { + if (last_removed) |last| { + expect(last <= next); + } + last_removed = next; + } +} + +test "std.PriorityDequeue: fuzz testing max" { + var prng = std.rand.DefaultPrng.init(0x87654321); + + const test_case_count = 100; + const heap_size = 1_000; + + var i: usize = 0; + while (i < test_case_count) : (i += 1) { + try fuzzTestMax(&prng.random, heap_size); + } +} + +fn fuzzTestMax(rng: *std.rand.Random, heap_size: usize) !void { + const allocator = testing.allocator; + const items = try generateRandomSlice(allocator, rng, heap_size); + + var heap = Heap.fromOwnedSlice(testing.allocator, lessThanComparison, items); + defer heap.deinit(); + + var last_removed: ?u32 = null; + while (heap.removeMaxOrNull()) |next| { + if (last_removed) |last| { + expect(last >= next); + } + last_removed = next; + } +} + +test "std.PriorityDequeue: fuzz testing min and max" { + var prng = std.rand.DefaultPrng.init(0x87654321); + + const test_case_count = 100; + const heap_size = 1_000; + + var i: usize = 0; + while (i < test_case_count) : (i += 1) { + try fuzzTestMinMax(&prng.random, heap_size); + } +} + +fn fuzzTestMinMax(rng: *std.rand.Random, heap_size: usize) !void { + const allocator = testing.allocator; + const items = try generateRandomSlice(allocator, rng, heap_size); + + var heap = Heap.fromOwnedSlice(allocator, lessThanComparison, items); + defer heap.deinit(); + + var last_min: ?u32 = null; + var last_max: ?u32 = null; + var i: usize = 0; + while (i < heap_size) : (i += 1) { + if (i % 2 == 0) { + const next = heap.removeMin(); + if (last_min) |last| { + expect(last <= next); + } + last_min = next; + } else { + const next = heap.removeMax(); + if (last_max) |last| { + expect(last >= next); + } + last_max = next; + } + } +} + +fn generateRandomSlice(allocator: *std.mem.Allocator, rng: *std.rand.Random, size: usize) ![]u32 { + var array = std.ArrayList(u32).init(allocator); + try array.ensureCapacity(size); + + var i: usize = 0; + while (i < size) : (i += 1) { + const elem = rng.int(u32); + try array.append(elem); + } + + return array.toOwnedSlice(); +} -- cgit v1.2.3 From 349ccc0bd07df49c9fa5934036e848d7885ebacb Mon Sep 17 00:00:00 2001 From: Zander Khan Date: Sat, 16 Jan 2021 17:48:14 +0000 Subject: Add license to top of file --- lib/std/priority_dequeue.zig | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) (limited to 'lib/std/priority_dequeue.zig') diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig index a742b27050..9df8c07f7e 100644 --- a/lib/std/priority_dequeue.zig +++ b/lib/std/priority_dequeue.zig @@ -1,6 +1,10 @@ -const std = @import("std"); +// SPDX-License-Identifier: MIT +// Copyright (c) 2015-2021 Zig Contributors +// This file is part of [zig](https://ziglang.org/), which is MIT licensed. +// The MIT license requires this copyright notice to be included in all copies +// and substantial portions of the software. +const std = @import("std.zig"); const Allocator = std.mem.Allocator; -const sort = std.sort; const assert = std.debug.assert; const warn = std.debug.warn; const testing = std.testing; -- cgit v1.2.3 From ecee1cae454ad5924f4698caaa712348858d9179 Mon Sep 17 00:00:00 2001 From: Zander Khan Date: Sat, 16 Jan 2021 17:48:55 +0000 Subject: Remove magic number --- lib/std/priority_dequeue.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'lib/std/priority_dequeue.zig') diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig index 9df8c07f7e..5dd861bcbd 100644 --- a/lib/std/priority_dequeue.zig +++ b/lib/std/priority_dequeue.zig @@ -81,7 +81,7 @@ pub fn PriorityDequeue(comptime T: type) type { // next two are on a max layer; // next four are on a min layer, and so on. const leading_zeros = @clz(usize, index + 1); - const highest_set_bit = 63 - leading_zeros; + const highest_set_bit = @bitSizeOf(usize) - 1 - leading_zeros; return (highest_set_bit & 1) == 0; } -- cgit v1.2.3 From a162a2194744ed9d9c095f0160e463a7a1b4e36c Mon Sep 17 00:00:00 2001 From: Zander Khan Date: Sat, 16 Jan 2021 17:52:14 +0000 Subject: Ensure we cannot remove an item outside the current length of the queue --- lib/std/priority_dequeue.zig | 1 + 1 file changed, 1 insertion(+) (limited to 'lib/std/priority_dequeue.zig') diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig index 5dd861bcbd..540f37960a 100644 --- a/lib/std/priority_dequeue.zig +++ b/lib/std/priority_dequeue.zig @@ -191,6 +191,7 @@ pub fn PriorityDequeue(comptime T: type) type { /// same order as iterator, which is not necessarily priority /// order. pub fn removeIndex(self: *Self, index: usize) T { + assert(self.len > index); const item = self.items[index]; const last = self.items[self.len - 1]; -- cgit v1.2.3 From 4d098034141d1f6ea5874d147227979af8aaae3c Mon Sep 17 00:00:00 2001 From: Zander Khan Date: Sat, 16 Jan 2021 18:06:44 +0000 Subject: Fix edge cases in fromOwnedSlice --- lib/std/priority_dequeue.zig | 31 +++++++++++++++++++++++++++---- lib/std/priority_queue.zig | 23 +++++++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) (limited to 'lib/std/priority_dequeue.zig') diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig index 540f37960a..81f11605f5 100644 --- a/lib/std/priority_dequeue.zig +++ b/lib/std/priority_dequeue.zig @@ -349,19 +349,22 @@ pub fn PriorityDequeue(comptime T: type) type { /// allocated with `allocator`. /// De-initialize with `deinit`. pub fn fromOwnedSlice(allocator: *Allocator, lessThanFn: fn (T, T) bool, items: []T) Self { - var dequeue = Self{ + var queue = Self{ .items = items, .len = items.len, .allocator = allocator, .lessThanFn = lessThanFn, }; - const half = (dequeue.len >> 1) - 1; + + if (queue.len <= 1) return queue; + + const half = (queue.len >> 1) - 1; var i: usize = 0; while (i <= half) : (i += 1) { const index = half - i; - dequeue.siftDown(index); + queue.siftDown(index); } - return dequeue; + return queue; } pub fn ensureCapacity(self: *Self, new_capacity: usize) !void { @@ -646,6 +649,26 @@ test "std.PriorityDequeue: addSlice max" { } } +test "std.PriorityDequeue: fromOwnedSlice trivial case 0" { + const items = [0]u32{}; + const heap_items = try testing.allocator.dupe(u32, &items); + var heap = Heap.fromOwnedSlice(testing.allocator, lessThanComparison, heap_items[0..]); + defer heap.deinit(); + expectEqual(@as(usize, 0), heap.len); + expect(heap.removeMinOrNull() == null); +} + +test "std.PriorityDequeue: fromOwnedSlice trivial case 1" { + const items = [1]u32{1}; + const heap_items = try testing.allocator.dupe(u32, &items); + var heap = Heap.fromOwnedSlice(testing.allocator, lessThanComparison, heap_items[0..]); + defer heap.deinit(); + + expectEqual(@as(usize, 1), heap.len); + expectEqual(items[0], heap.removeMin()); + expect(heap.removeMinOrNull() == null); +} + test "std.PriorityDequeue: fromOwnedSlice" { 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..]); diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig index 24ad0ec2b1..a5f9209e86 100644 --- a/lib/std/priority_queue.zig +++ b/lib/std/priority_queue.zig @@ -166,6 +166,9 @@ pub fn PriorityQueue(comptime T: type) type { .allocator = allocator, .compareFn = compareFn, }; + + if (queue.len <= 1) return queue; + const half = (queue.len >> 1) - 1; var i: usize = 0; while (i <= half) : (i += 1) { @@ -352,6 +355,26 @@ test "std.PriorityQueue: addSlice" { } } +test "std.PriorityQueue: fromOwnedSlice trivial case 0" { + const items = [0]u32{}; + const queue_items = try testing.allocator.dupe(u32, &items); + var queue = PQ.fromOwnedSlice(testing.allocator, lessThan, queue_items[0..]); + defer queue.deinit(); + expectEqual(@as(usize, 0), queue.len); + expect(queue.removeOrNull() == null); +} + +test "std.PriorityQueue: fromOwnedSlice trivial case 1" { + const items = [1]u32{1}; + const queue_items = try testing.allocator.dupe(u32, &items); + var queue = PQ.fromOwnedSlice(testing.allocator, lessThan, queue_items[0..]); + defer queue.deinit(); + + expectEqual(@as(usize, 1), queue.len); + expectEqual(items[0], queue.remove()); + expect(queue.removeOrNull() == null); +} + test "std.PriorityQueue: fromOwnedSlice" { 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..]); -- cgit v1.2.3 From 4600b489a6daf3990afc8464dd7f2b9a1a4efa72 Mon Sep 17 00:00:00 2001 From: Zander Khan Date: Sat, 16 Jan 2021 18:09:44 +0000 Subject: Rename heap to queue in tests for consistency --- lib/std/priority_dequeue.zig | 466 +++++++++++++++++++++---------------------- 1 file changed, 233 insertions(+), 233 deletions(-) (limited to 'lib/std/priority_dequeue.zig') diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig index 81f11605f5..9dbb457bac 100644 --- a/lib/std/priority_dequeue.zig +++ b/lib/std/priority_dequeue.zig @@ -334,7 +334,7 @@ pub fn PriorityDequeue(comptime T: type) type { } } - /// Return the number of elements remaining in the heap + /// Return the number of elements remaining in the dequeue pub fn count(self: Self) usize { return self.len; } @@ -345,7 +345,7 @@ pub fn PriorityDequeue(comptime T: type) type { return self.items.len; } - /// Heap takes ownership of the passed in slice. The slice must have been + /// Dequeue takes ownership of the passed in slice. The slice must have been /// allocated with `allocator`. /// De-initialize with `deinit`. pub fn fromOwnedSlice(allocator: *Allocator, lessThanFn: fn (T, T) bool, items: []T) Self { @@ -395,14 +395,14 @@ pub fn PriorityDequeue(comptime T: type) type { } pub const Iterator = struct { - heap: *PriorityDequeue(T), + queue: *PriorityDequeue(T), count: usize, pub fn next(it: *Iterator) ?T { - if (it.count >= it.heap.len) return null; + if (it.count >= it.queue.len) return null; const out = it.count; it.count += 1; - return it.heap.items[out]; + return it.queue.items[out]; } pub fn reset(it: *Iterator) void { @@ -410,11 +410,11 @@ pub fn PriorityDequeue(comptime T: type) type { } }; - /// Return an iterator that walks the heap without consuming - /// it. Invalidated if the heap is modified. + /// Return an iterator that walks the queue without consuming + /// it. Invalidated if the queue is modified. pub fn iterator(self: *Self) Iterator { return Iterator{ - .heap = self, + .queue = self, .count = 0, }; } @@ -457,308 +457,308 @@ fn lessThanComparison(a: u32, b: u32) bool { return a < b; } -const Heap = PriorityDequeue(u32); +const PDQ = PriorityDequeue(u32); test "std.PriorityDequeue: add and remove min" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); - try heap.add(54); - try heap.add(12); - try heap.add(7); - try heap.add(23); - try heap.add(25); - try heap.add(13); + try queue.add(54); + try queue.add(12); + try queue.add(7); + try queue.add(23); + try queue.add(25); + try queue.add(13); - expectEqual(@as(u32, 7), heap.removeMin()); - expectEqual(@as(u32, 12), heap.removeMin()); - expectEqual(@as(u32, 13), heap.removeMin()); - expectEqual(@as(u32, 23), heap.removeMin()); - expectEqual(@as(u32, 25), heap.removeMin()); - expectEqual(@as(u32, 54), heap.removeMin()); + expectEqual(@as(u32, 7), queue.removeMin()); + expectEqual(@as(u32, 12), queue.removeMin()); + expectEqual(@as(u32, 13), queue.removeMin()); + expectEqual(@as(u32, 23), queue.removeMin()); + expectEqual(@as(u32, 25), queue.removeMin()); + expectEqual(@as(u32, 54), queue.removeMin()); } test "std.PriorityDequeue: add and remove max" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); - try heap.add(54); - try heap.add(12); - try heap.add(7); - try heap.add(23); - try heap.add(25); - try heap.add(13); + try queue.add(54); + try queue.add(12); + try queue.add(7); + try queue.add(23); + try queue.add(25); + try queue.add(13); - expectEqual(@as(u32, 54), heap.removeMax()); - expectEqual(@as(u32, 25), heap.removeMax()); - expectEqual(@as(u32, 23), heap.removeMax()); - expectEqual(@as(u32, 13), heap.removeMax()); - expectEqual(@as(u32, 12), heap.removeMax()); - expectEqual(@as(u32, 7), heap.removeMax()); + expectEqual(@as(u32, 54), queue.removeMax()); + expectEqual(@as(u32, 25), queue.removeMax()); + expectEqual(@as(u32, 23), queue.removeMax()); + expectEqual(@as(u32, 13), queue.removeMax()); + expectEqual(@as(u32, 12), queue.removeMax()); + expectEqual(@as(u32, 7), queue.removeMax()); } test "std.PriorityDequeue: add and remove same min" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); - try heap.add(1); - try heap.add(1); - try heap.add(2); - try heap.add(2); - try heap.add(1); - try heap.add(1); + try queue.add(1); + try queue.add(1); + try queue.add(2); + try queue.add(2); + try queue.add(1); + try queue.add(1); - expectEqual(@as(u32, 1), heap.removeMin()); - expectEqual(@as(u32, 1), heap.removeMin()); - expectEqual(@as(u32, 1), heap.removeMin()); - expectEqual(@as(u32, 1), heap.removeMin()); - expectEqual(@as(u32, 2), heap.removeMin()); - expectEqual(@as(u32, 2), heap.removeMin()); + expectEqual(@as(u32, 1), queue.removeMin()); + expectEqual(@as(u32, 1), queue.removeMin()); + expectEqual(@as(u32, 1), queue.removeMin()); + expectEqual(@as(u32, 1), queue.removeMin()); + expectEqual(@as(u32, 2), queue.removeMin()); + expectEqual(@as(u32, 2), queue.removeMin()); } test "std.PriorityDequeue: add and remove same max" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); - try heap.add(1); - try heap.add(1); - try heap.add(2); - try heap.add(2); - try heap.add(1); - try heap.add(1); + try queue.add(1); + try queue.add(1); + try queue.add(2); + try queue.add(2); + try queue.add(1); + try queue.add(1); - expectEqual(@as(u32, 2), heap.removeMax()); - expectEqual(@as(u32, 2), heap.removeMax()); - expectEqual(@as(u32, 1), heap.removeMax()); - expectEqual(@as(u32, 1), heap.removeMax()); - expectEqual(@as(u32, 1), heap.removeMax()); - expectEqual(@as(u32, 1), heap.removeMax()); + expectEqual(@as(u32, 2), queue.removeMax()); + expectEqual(@as(u32, 2), queue.removeMax()); + expectEqual(@as(u32, 1), queue.removeMax()); + expectEqual(@as(u32, 1), queue.removeMax()); + expectEqual(@as(u32, 1), queue.removeMax()); + expectEqual(@as(u32, 1), queue.removeMax()); } test "std.PriorityDequeue: removeOrNull empty" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); - expect(heap.removeMinOrNull() == null); - expect(heap.removeMaxOrNull() == null); + expect(queue.removeMinOrNull() == null); + expect(queue.removeMaxOrNull() == null); } test "std.PriorityDequeue: edge case 3 elements" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); - try heap.add(9); - try heap.add(3); - try heap.add(2); + try queue.add(9); + try queue.add(3); + try queue.add(2); - expectEqual(@as(u32, 2), heap.removeMin()); - expectEqual(@as(u32, 3), heap.removeMin()); - expectEqual(@as(u32, 9), heap.removeMin()); + expectEqual(@as(u32, 2), queue.removeMin()); + expectEqual(@as(u32, 3), queue.removeMin()); + expectEqual(@as(u32, 9), queue.removeMin()); } test "std.PriorityDequeue: edge case 3 elements max" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); - try heap.add(9); - try heap.add(3); - try heap.add(2); + try queue.add(9); + try queue.add(3); + try queue.add(2); - expectEqual(@as(u32, 9), heap.removeMax()); - expectEqual(@as(u32, 3), heap.removeMax()); - expectEqual(@as(u32, 2), heap.removeMax()); + expectEqual(@as(u32, 9), queue.removeMax()); + expectEqual(@as(u32, 3), queue.removeMax()); + expectEqual(@as(u32, 2), queue.removeMax()); } test "std.PriorityDequeue: peekMin" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); - expect(heap.peekMin() == null); + expect(queue.peekMin() == null); - try heap.add(9); - try heap.add(3); - try heap.add(2); + try queue.add(9); + try queue.add(3); + try queue.add(2); - expect(heap.peekMin().? == 2); - expect(heap.peekMin().? == 2); + expect(queue.peekMin().? == 2); + expect(queue.peekMin().? == 2); } test "std.PriorityDequeue: peekMax" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); - expect(heap.peekMin() == null); + expect(queue.peekMin() == null); - try heap.add(9); - try heap.add(3); - try heap.add(2); + try queue.add(9); + try queue.add(3); + try queue.add(2); - expect(heap.peekMax().? == 9); - expect(heap.peekMax().? == 9); + expect(queue.peekMax().? == 9); + expect(queue.peekMax().? == 9); } test "std.PriorityDequeue: sift up with odd indices" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); 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 heap.add(e); + try queue.add(e); } 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| { - expectEqual(e, heap.removeMin()); + expectEqual(e, queue.removeMin()); } } test "std.PriorityDequeue: sift up with odd indices" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); 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 heap.add(e); + try queue.add(e); } const sorted_items = [_]u32{ 25, 24, 24, 22, 21, 16, 15, 15, 14, 13, 12, 11, 7, 7, 6, 5, 2, 1 }; for (sorted_items) |e| { - expectEqual(e, heap.removeMax()); + expectEqual(e, queue.removeMax()); } } test "std.PriorityDequeue: addSlice min" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; - try heap.addSlice(items[0..]); + try queue.addSlice(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| { - expectEqual(e, heap.removeMin()); + expectEqual(e, queue.removeMin()); } } test "std.PriorityDequeue: addSlice max" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; - try heap.addSlice(items[0..]); + try queue.addSlice(items[0..]); const sorted_items = [_]u32{ 25, 24, 24, 22, 21, 16, 15, 15, 14, 13, 12, 11, 7, 7, 6, 5, 2, 1 }; for (sorted_items) |e| { - expectEqual(e, heap.removeMax()); + expectEqual(e, queue.removeMax()); } } test "std.PriorityDequeue: fromOwnedSlice trivial case 0" { const items = [0]u32{}; - const heap_items = try testing.allocator.dupe(u32, &items); - var heap = Heap.fromOwnedSlice(testing.allocator, lessThanComparison, heap_items[0..]); - defer heap.deinit(); - expectEqual(@as(usize, 0), heap.len); - expect(heap.removeMinOrNull() == null); + const queue_items = try testing.allocator.dupe(u32, &items); + var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, queue_items[0..]); + defer queue.deinit(); + expectEqual(@as(usize, 0), queue.len); + expect(queue.removeMinOrNull() == null); } test "std.PriorityDequeue: fromOwnedSlice trivial case 1" { const items = [1]u32{1}; - const heap_items = try testing.allocator.dupe(u32, &items); - var heap = Heap.fromOwnedSlice(testing.allocator, lessThanComparison, heap_items[0..]); - defer heap.deinit(); + const queue_items = try testing.allocator.dupe(u32, &items); + var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, queue_items[0..]); + defer queue.deinit(); - expectEqual(@as(usize, 1), heap.len); - expectEqual(items[0], heap.removeMin()); - expect(heap.removeMinOrNull() == null); + expectEqual(@as(usize, 1), queue.len); + expectEqual(items[0], queue.removeMin()); + expect(queue.removeMinOrNull() == null); } test "std.PriorityDequeue: fromOwnedSlice" { 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 heap = Heap.fromOwnedSlice(testing.allocator, lessThanComparison, heap_items[0..]); - defer heap.deinit(); + const queue_items = try testing.allocator.dupe(u32, items[0..]); + var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, queue_items[0..]); + defer queue.deinit(); 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| { - expectEqual(e, heap.removeMin()); + expectEqual(e, queue.removeMin()); } } -test "std.PriorityDequeue: update min heap" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); - - try heap.add(55); - try heap.add(44); - try heap.add(11); - try heap.update(55, 5); - try heap.update(44, 4); - try heap.update(11, 1); - expectEqual(@as(u32, 1), heap.removeMin()); - expectEqual(@as(u32, 4), heap.removeMin()); - expectEqual(@as(u32, 5), heap.removeMin()); -} - -test "std.PriorityDequeue: update same min heap" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); - - try heap.add(1); - try heap.add(1); - try heap.add(2); - try heap.add(2); - try heap.update(1, 5); - try heap.update(2, 4); - expectEqual(@as(u32, 1), heap.removeMin()); - expectEqual(@as(u32, 2), heap.removeMin()); - expectEqual(@as(u32, 4), heap.removeMin()); - expectEqual(@as(u32, 5), heap.removeMin()); -} - -test "std.PriorityDequeue: update max heap" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); - - try heap.add(55); - try heap.add(44); - try heap.add(11); - try heap.update(55, 5); - try heap.update(44, 1); - try heap.update(11, 4); - - expectEqual(@as(u32, 5), heap.removeMax()); - expectEqual(@as(u32, 4), heap.removeMax()); - expectEqual(@as(u32, 1), heap.removeMax()); -} - -test "std.PriorityDequeue: update same max heap" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); - - try heap.add(1); - try heap.add(1); - try heap.add(2); - try heap.add(2); - try heap.update(1, 5); - try heap.update(2, 4); - expectEqual(@as(u32, 5), heap.removeMax()); - expectEqual(@as(u32, 4), heap.removeMax()); - expectEqual(@as(u32, 2), heap.removeMax()); - expectEqual(@as(u32, 1), heap.removeMax()); +test "std.PriorityDequeue: update min queue" { + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); + + try queue.add(55); + try queue.add(44); + try queue.add(11); + try queue.update(55, 5); + try queue.update(44, 4); + try queue.update(11, 1); + expectEqual(@as(u32, 1), queue.removeMin()); + expectEqual(@as(u32, 4), queue.removeMin()); + expectEqual(@as(u32, 5), queue.removeMin()); +} + +test "std.PriorityDequeue: update same min queue" { + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); + + try queue.add(1); + try queue.add(1); + try queue.add(2); + try queue.add(2); + try queue.update(1, 5); + try queue.update(2, 4); + expectEqual(@as(u32, 1), queue.removeMin()); + expectEqual(@as(u32, 2), queue.removeMin()); + expectEqual(@as(u32, 4), queue.removeMin()); + expectEqual(@as(u32, 5), queue.removeMin()); +} + +test "std.PriorityDequeue: update max queue" { + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); + + try queue.add(55); + try queue.add(44); + try queue.add(11); + try queue.update(55, 5); + try queue.update(44, 1); + try queue.update(11, 4); + + expectEqual(@as(u32, 5), queue.removeMax()); + expectEqual(@as(u32, 4), queue.removeMax()); + expectEqual(@as(u32, 1), queue.removeMax()); +} + +test "std.PriorityDequeue: update same max queue" { + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); + + try queue.add(1); + try queue.add(1); + try queue.add(2); + try queue.add(2); + try queue.update(1, 5); + try queue.update(2, 4); + expectEqual(@as(u32, 5), queue.removeMax()); + expectEqual(@as(u32, 4), queue.removeMax()); + expectEqual(@as(u32, 2), queue.removeMax()); + expectEqual(@as(u32, 1), queue.removeMax()); } test "std.PriorityDequeue: iterator" { - var heap = Heap.init(testing.allocator, lessThanComparison); + var queue = PDQ.init(testing.allocator, lessThanComparison); var map = std.AutoHashMap(u32, void).init(testing.allocator); defer { - heap.deinit(); + queue.deinit(); map.deinit(); } const items = [_]u32{ 54, 12, 7, 23, 25, 13 }; for (items) |e| { - _ = try heap.add(e); + _ = try queue.add(e); _ = try map.put(e, {}); } - var it = heap.iterator(); + var it = queue.iterator(); while (it.next()) |e| { _ = map.remove(e); } @@ -767,14 +767,14 @@ test "std.PriorityDequeue: iterator" { } test "std.PriorityDequeue: remove at index" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); - try heap.add(3); - try heap.add(2); - try heap.add(1); + try queue.add(3); + try queue.add(2); + try queue.add(1); - var it = heap.iterator(); + var it = queue.iterator(); var elem = it.next(); var idx: usize = 0; const two_idx = while (elem != null) : (elem = it.next()) { @@ -783,17 +783,17 @@ test "std.PriorityDequeue: remove at index" { idx += 1; } else unreachable; - expectEqual(heap.removeIndex(two_idx), 2); - expectEqual(heap.removeMin(), 1); - expectEqual(heap.removeMin(), 3); - expectEqual(heap.removeMinOrNull(), null); + expectEqual(queue.removeIndex(two_idx), 2); + expectEqual(queue.removeMin(), 1); + expectEqual(queue.removeMin(), 3); + expectEqual(queue.removeMinOrNull(), null); } test "std.PriorityDequeue: iterator while empty" { - var heap = Heap.init(testing.allocator, lessThanComparison); - defer heap.deinit(); + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); - var it = heap.iterator(); + var it = queue.iterator(); expectEqual(it.next(), null); } @@ -802,23 +802,23 @@ test "std.PriorityDequeue: fuzz testing min" { var prng = std.rand.DefaultPrng.init(0x12345678); const test_case_count = 100; - const heap_size = 1_000; + const queue_size = 1_000; var i: usize = 0; while (i < test_case_count) : (i += 1) { - try fuzzTestMin(&prng.random, heap_size); + try fuzzTestMin(&prng.random, queue_size); } } -fn fuzzTestMin(rng: *std.rand.Random, comptime heap_size: usize) !void { +fn fuzzTestMin(rng: *std.rand.Random, comptime queue_size: usize) !void { const allocator = testing.allocator; - const items = try generateRandomSlice(allocator, rng, heap_size); + const items = try generateRandomSlice(allocator, rng, queue_size); - var heap = Heap.fromOwnedSlice(allocator, lessThanComparison, items); - defer heap.deinit(); + var queue = PDQ.fromOwnedSlice(allocator, lessThanComparison, items); + defer queue.deinit(); var last_removed: ?u32 = null; - while (heap.removeMinOrNull()) |next| { + while (queue.removeMinOrNull()) |next| { if (last_removed) |last| { expect(last <= next); } @@ -830,23 +830,23 @@ test "std.PriorityDequeue: fuzz testing max" { var prng = std.rand.DefaultPrng.init(0x87654321); const test_case_count = 100; - const heap_size = 1_000; + const queue_size = 1_000; var i: usize = 0; while (i < test_case_count) : (i += 1) { - try fuzzTestMax(&prng.random, heap_size); + try fuzzTestMax(&prng.random, queue_size); } } -fn fuzzTestMax(rng: *std.rand.Random, heap_size: usize) !void { +fn fuzzTestMax(rng: *std.rand.Random, queue_size: usize) !void { const allocator = testing.allocator; - const items = try generateRandomSlice(allocator, rng, heap_size); + const items = try generateRandomSlice(allocator, rng, queue_size); - var heap = Heap.fromOwnedSlice(testing.allocator, lessThanComparison, items); - defer heap.deinit(); + var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, items); + defer queue.deinit(); var last_removed: ?u32 = null; - while (heap.removeMaxOrNull()) |next| { + while (queue.removeMaxOrNull()) |next| { if (last_removed) |last| { expect(last >= next); } @@ -858,33 +858,33 @@ test "std.PriorityDequeue: fuzz testing min and max" { var prng = std.rand.DefaultPrng.init(0x87654321); const test_case_count = 100; - const heap_size = 1_000; + const queue_size = 1_000; var i: usize = 0; while (i < test_case_count) : (i += 1) { - try fuzzTestMinMax(&prng.random, heap_size); + try fuzzTestMinMax(&prng.random, queue_size); } } -fn fuzzTestMinMax(rng: *std.rand.Random, heap_size: usize) !void { +fn fuzzTestMinMax(rng: *std.rand.Random, queue_size: usize) !void { const allocator = testing.allocator; - const items = try generateRandomSlice(allocator, rng, heap_size); + const items = try generateRandomSlice(allocator, rng, queue_size); - var heap = Heap.fromOwnedSlice(allocator, lessThanComparison, items); - defer heap.deinit(); + var queue = PDQ.fromOwnedSlice(allocator, lessThanComparison, items); + defer queue.deinit(); var last_min: ?u32 = null; var last_max: ?u32 = null; var i: usize = 0; - while (i < heap_size) : (i += 1) { + while (i < queue_size) : (i += 1) { if (i % 2 == 0) { - const next = heap.removeMin(); + const next = queue.removeMin(); if (last_min) |last| { expect(last <= next); } last_min = next; } else { - const next = heap.removeMax(); + const next = queue.removeMax(); if (last_max) |last| { expect(last >= next); } -- cgit v1.2.3 From c6986f29f94ee404ae3a3221449dc4af5599ca1f Mon Sep 17 00:00:00 2001 From: Zander Khan Date: Sat, 16 Jan 2021 18:11:26 +0000 Subject: Fix update might change an element no longer in the queue --- lib/std/priority_dequeue.zig | 2 +- lib/std/priority_queue.zig | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) (limited to 'lib/std/priority_dequeue.zig') diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig index 9dbb457bac..0d789bbdc8 100644 --- a/lib/std/priority_dequeue.zig +++ b/lib/std/priority_dequeue.zig @@ -389,7 +389,7 @@ pub fn PriorityDequeue(comptime T: type) type { } pub fn update(self: *Self, elem: T, new_elem: T) !void { - var old_index: usize = std.mem.indexOfScalar(T, self.items, elem) orelse return error.ElementNotFound; + var old_index: usize = std.mem.indexOfScalar(T, self.items[0 .. self.len - 1], elem) orelse return error.ElementNotFound; _ = self.removeIndex(old_index); self.addUnchecked(new_elem); } diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig index a5f9209e86..6e286f1cea 100644 --- a/lib/std/priority_queue.zig +++ b/lib/std/priority_queue.zig @@ -199,7 +199,7 @@ pub fn PriorityQueue(comptime T: type) type { } pub fn update(self: *Self, elem: T, new_elem: T) !void { - var update_index: usize = std.mem.indexOfScalar(T, self.items, elem) orelse return error.ElementNotFound; + var update_index: usize = std.mem.indexOfScalar(T, self.items[0 .. self.len - 1], elem) orelse return error.ElementNotFound; const old_elem: T = self.items[update_index]; self.items[update_index] = new_elem; if (self.compareFn(new_elem, old_elem)) { -- cgit v1.2.3 From e1ab425bcead727a73e4512aeca1ba9112b2c88e Mon Sep 17 00:00:00 2001 From: Zander Khan Date: Sat, 16 Jan 2021 18:43:13 +0000 Subject: Fix slice length when updating --- lib/std/priority_dequeue.zig | 2 +- lib/std/priority_queue.zig | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) (limited to 'lib/std/priority_dequeue.zig') diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig index 0d789bbdc8..9655458890 100644 --- a/lib/std/priority_dequeue.zig +++ b/lib/std/priority_dequeue.zig @@ -389,7 +389,7 @@ pub fn PriorityDequeue(comptime T: type) type { } pub fn update(self: *Self, elem: T, new_elem: T) !void { - var old_index: usize = std.mem.indexOfScalar(T, self.items[0 .. self.len - 1], elem) orelse return error.ElementNotFound; + var old_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], elem) orelse return error.ElementNotFound; _ = self.removeIndex(old_index); self.addUnchecked(new_elem); } diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig index 6e286f1cea..dc3070d1b3 100644 --- a/lib/std/priority_queue.zig +++ b/lib/std/priority_queue.zig @@ -199,7 +199,7 @@ pub fn PriorityQueue(comptime T: type) type { } pub fn update(self: *Self, elem: T, new_elem: T) !void { - var update_index: usize = std.mem.indexOfScalar(T, self.items[0 .. self.len - 1], elem) orelse return error.ElementNotFound; + var update_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], elem) orelse return error.ElementNotFound; const old_elem: T = self.items[update_index]; self.items[update_index] = new_elem; if (self.compareFn(new_elem, old_elem)) { -- cgit v1.2.3 From 9a09ebb1b92f33913bdcc16e9d978fb6e5a820f1 Mon Sep 17 00:00:00 2001 From: Zander Khan Date: Sun, 17 Jan 2021 14:41:20 +0000 Subject: Replace `shrink` with `shrinkAndFree` and `shrinkRetainingCapacity` --- lib/std/priority_dequeue.zig | 53 +++++++++++++++++++++++++++++++++++++++++--- lib/std/priority_queue.zig | 53 +++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 100 insertions(+), 6 deletions(-) (limited to 'lib/std/priority_dequeue.zig') diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig index 9655458890..bb119438d2 100644 --- a/lib/std/priority_dequeue.zig +++ b/lib/std/priority_dequeue.zig @@ -382,9 +382,29 @@ pub fn PriorityDequeue(comptime T: type) type { self.len = new_len; } - pub fn shrink(self: *Self, new_len: usize) void { - // TODO take advantage of the new realloc semantics - assert(new_len <= self.len); + /// Reduce allocated capacity to `new_len`. + pub fn shrinkAndFree(self: *Self, new_len: usize) void { + assert(new_len <= self.items.len); + + // Cannot shrink to smaller than the current queue size without invalidating the heap property + assert(new_len >= self.len); + + self.items = self.allocator.realloc(self.items[0..], new_len) catch |e| switch (e) { + error.OutOfMemory => { // no problem, capacity is still correct then. + self.items.len = new_len; + return; + }, + }; + self.len = new_len; + } + + /// Reduce length to `new_len`. + pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void { + assert(new_len <= self.items.len); + + // Cannot shrink to smaller than the current queue size without invalidating the heap property + assert(new_len >= self.len); + self.len = new_len; } @@ -798,6 +818,33 @@ test "std.PriorityDequeue: iterator while empty" { expectEqual(it.next(), null); } +test "std.PriorityDequeue: shrinkRetainingCapacity and shrinkAndFree" { + var queue = PDQ.init(testing.allocator, lessThanComparison); + defer queue.deinit(); + + try queue.ensureCapacity(4); + expect(queue.capacity() >= 4); + + try queue.add(1); + try queue.add(2); + try queue.add(3); + expect(queue.capacity() >= 4); + expectEqual(@as(usize, 3), queue.len); + + queue.shrinkRetainingCapacity(3); + expect(queue.capacity() >= 4); + expectEqual(@as(usize, 3), queue.len); + + queue.shrinkAndFree(3); + expectEqual(@as(usize, 3), queue.capacity()); + expectEqual(@as(usize, 3), queue.len); + + expectEqual(@as(u32, 3), queue.removeMax()); + expectEqual(@as(u32, 2), queue.removeMax()); + expectEqual(@as(u32, 1), queue.removeMax()); + expect(queue.removeMaxOrNull() == null); +} + test "std.PriorityDequeue: fuzz testing min" { var prng = std.rand.DefaultPrng.init(0x12345678); diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig index dc3070d1b3..844d37580c 100644 --- a/lib/std/priority_queue.zig +++ b/lib/std/priority_queue.zig @@ -192,9 +192,29 @@ pub fn PriorityQueue(comptime T: type) type { self.len = new_len; } - pub fn shrink(self: *Self, new_len: usize) void { - // TODO take advantage of the new realloc semantics - assert(new_len <= self.len); + /// Reduce allocated capacity to `new_len`. + pub fn shrinkAndFree(self: *Self, new_len: usize) void { + assert(new_len <= self.items.len); + + // Cannot shrink to smaller than the current queue size without invalidating the heap property + assert(new_len >= self.len); + + self.items = self.allocator.realloc(self.items[0..], new_len) catch |e| switch (e) { + error.OutOfMemory => { // no problem, capacity is still correct then. + self.items.len = new_len; + return; + }, + }; + self.len = new_len; + } + + /// Reduce length to `new_len`. + pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void { + assert(new_len <= self.items.len); + + // Cannot shrink to smaller than the current queue size without invalidating the heap property + assert(new_len >= self.len); + self.len = new_len; } @@ -477,6 +497,33 @@ test "std.PriorityQueue: iterator while empty" { expectEqual(it.next(), null); } +test "std.PriorityQueue: shrinkRetainingCapacity and shrinkAndFree" { + var queue = PQ.init(testing.allocator, lessThan); + defer queue.deinit(); + + try queue.ensureCapacity(4); + expect(queue.capacity() >= 4); + + try queue.add(1); + try queue.add(2); + try queue.add(3); + expect(queue.capacity() >= 4); + expectEqual(@as(usize, 3), queue.len); + + queue.shrinkRetainingCapacity(3); + expect(queue.capacity() >= 4); + expectEqual(@as(usize, 3), queue.len); + + queue.shrinkAndFree(3); + expectEqual(@as(usize, 3), queue.capacity()); + expectEqual(@as(usize, 3), queue.len); + + expectEqual(@as(u32, 1), queue.remove()); + expectEqual(@as(u32, 2), queue.remove()); + expectEqual(@as(u32, 3), queue.remove()); + expect(queue.removeOrNull() == null); +} + test "std.PriorityQueue: update min heap" { var queue = PQ.init(testing.allocator, lessThan); defer queue.deinit(); -- cgit v1.2.3 From 5bfd9238de82b8f66ba4834bcf7aa120e42fe6ca Mon Sep 17 00:00:00 2001 From: Zander Khan Date: Sun, 17 Jan 2021 14:43:38 +0000 Subject: Remove `resize`. Adding uninitialized memory at the end of the `items` would break the heap property. --- lib/std/priority_dequeue.zig | 5 ----- lib/std/priority_queue.zig | 5 ----- 2 files changed, 10 deletions(-) (limited to 'lib/std/priority_dequeue.zig') diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig index bb119438d2..2e19da1d5d 100644 --- a/lib/std/priority_dequeue.zig +++ b/lib/std/priority_dequeue.zig @@ -377,11 +377,6 @@ pub fn PriorityDequeue(comptime T: type) type { self.items = try self.allocator.realloc(self.items, better_capacity); } - pub fn resize(self: *Self, new_len: usize) !void { - try self.ensureCapacity(new_len); - self.len = new_len; - } - /// Reduce allocated capacity to `new_len`. pub fn shrinkAndFree(self: *Self, new_len: usize) void { assert(new_len <= self.items.len); diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig index 844d37580c..f67b390608 100644 --- a/lib/std/priority_queue.zig +++ b/lib/std/priority_queue.zig @@ -187,11 +187,6 @@ pub fn PriorityQueue(comptime T: type) type { self.items = try self.allocator.realloc(self.items, better_capacity); } - pub fn resize(self: *Self, new_len: usize) !void { - try self.ensureCapacity(new_len); - self.len = new_len; - } - /// Reduce allocated capacity to `new_len`. pub fn shrinkAndFree(self: *Self, new_len: usize) void { assert(new_len <= self.items.len); -- cgit v1.2.3 From ce22c70586137bd20de101f3b1ec3a44fbdb54e7 Mon Sep 17 00:00:00 2001 From: Zander Khan Date: Mon, 18 Jan 2021 19:02:11 +0000 Subject: Change `compareFn` to `fn (a: T, b: T) std.math.Order` --- lib/std/priority_dequeue.zig | 131 +++++++++++++++++++++++++------------------ lib/std/priority_queue.zig | 43 +++++++------- 2 files changed, 100 insertions(+), 74 deletions(-) (limited to 'lib/std/priority_dequeue.zig') diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig index 2e19da1d5d..6894dcc997 100644 --- a/lib/std/priority_dequeue.zig +++ b/lib/std/priority_dequeue.zig @@ -7,6 +7,7 @@ const std = @import("std.zig"); const Allocator = std.mem.Allocator; const assert = std.debug.assert; const warn = std.debug.warn; +const Order = std.math.Order; const testing = std.testing; const expect = testing.expect; const expectEqual = testing.expectEqual; @@ -20,31 +21,26 @@ pub fn PriorityDequeue(comptime T: type) type { items: []T, len: usize, allocator: *Allocator, - lessThanFn: fn (a: T, b: T) bool, - - /// Initialize and return a new dequeue. Provide `lessThanFn` - /// that returns `true` when its first argument should - /// get min-popped before its second argument. For example, - /// to make `popMin` return the minimum value, provide + compareFn: fn (a: T, b: T) Order, + + /// Initialize and return a new priority dequeue. Provide `compareFn` + /// that returns `Order.lt` when its first argument should + /// get min-popped before its second argument, `Order.eq` if the + /// arguments are of equal priority, or `Order.gt` if the second + /// argument should be min-popped first. Popping the max element works + /// in reverse. For example, to make `popMin` return the smallest + /// number, provide /// - /// `fn lessThanFn(a: T, b: T) bool { return a < b; }` - pub fn init(allocator: *Allocator, lessThanFn: fn (T, T) bool) Self { + /// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }` + pub fn init(allocator: *Allocator, compareFn: fn (T, T) Order) Self { return Self{ .items = &[_]T{}, .len = 0, .allocator = allocator, - .lessThanFn = lessThanFn, + .compareFn = compareFn, }; } - fn lessThan(self: Self, a: T, b: T) bool { - return self.lessThanFn(a, b); - } - - fn greaterThan(self: Self, a: T, b: T) bool { - return self.lessThanFn(b, a); - } - /// Free memory used by the dequeue. pub fn deinit(self: Self) void { self.allocator.free(self.items); @@ -100,7 +96,8 @@ pub fn PriorityDequeue(comptime T: type) type { const parent = self.items[parent_index]; const min_layer = self.nextIsMinLayer(); - if ((min_layer and self.greaterThan(child, parent)) or (!min_layer and self.lessThan(child, parent))) { + const order = self.compareFn(child, parent); + if ((min_layer and order == .gt) or (!min_layer and order == .lt)) { // We must swap the item with it's parent if it is on the "wrong" layer self.items[parent_index] = child; self.items[child_index] = parent; @@ -118,21 +115,21 @@ pub fn PriorityDequeue(comptime T: type) type { fn siftUp(self: *Self, start: StartIndexAndLayer) void { if (start.min_layer) { - doSiftUp(self, start.index, lessThan); + doSiftUp(self, start.index, .lt); } else { - doSiftUp(self, start.index, greaterThan); + doSiftUp(self, start.index, .gt); } } - fn doSiftUp(self: *Self, start_index: usize, compare: fn (Self, T, T) bool) void { + fn doSiftUp(self: *Self, start_index: usize, target_order: Order) void { var child_index = start_index; while (child_index > 2) { var grandparent_index = grandparentIndex(child_index); const child = self.items[child_index]; const grandparent = self.items[grandparent_index]; - // If the grandparent is already better, we have gone as far as we need to - if (!compare(self.*, child, grandparent)) break; + // If the grandparent is already better or equal, we have gone as far as we need to + if (self.compareFn(child, grandparent) != target_order) break; // Otherwise swap the item with it's grandparent self.items[grandparent_index] = child; @@ -153,14 +150,14 @@ pub fn PriorityDequeue(comptime T: type) type { if (self.len == 0) return null; if (self.len == 1) return self.items[0]; if (self.len == 2) return self.items[1]; - return self.bestItemAtIndices(1, 2, greaterThan).item; + return self.bestItemAtIndices(1, 2, .gt).item; } fn maxIndex(self: Self) ?usize { if (self.len == 0) return null; if (self.len == 1) return 0; if (self.len == 2) return 1; - return self.bestItemAtIndices(1, 2, greaterThan).index; + return self.bestItemAtIndices(1, 2, .gt).index; } /// Pop the smallest element from the dequeue. Returns @@ -204,13 +201,13 @@ pub fn PriorityDequeue(comptime T: type) type { fn siftDown(self: *Self, index: usize) void { if (isMinLayer(index)) { - self.doSiftDown(index, lessThan); + self.doSiftDown(index, .lt); } else { - self.doSiftDown(index, greaterThan); + self.doSiftDown(index, .gt); } } - fn doSiftDown(self: *Self, start_index: usize, compare: fn (Self, T, T) bool) void { + fn doSiftDown(self: *Self, start_index: usize, target_order: Order) void { var index = start_index; const half = self.len >> 1; while (true) { @@ -225,12 +222,12 @@ pub fn PriorityDequeue(comptime T: type) type { const index3 = index2 + 1; // Find the best grandchild - const best_left = self.bestItemAtIndices(first_grandchild_index, index2, compare); - const best_right = self.bestItemAtIndices(index3, last_grandchild_index, compare); - const best_grandchild = self.bestItem(best_left, best_right, compare); + const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order); + const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order); + const best_grandchild = self.bestItem(best_left, best_right, target_order); - // If the item is better than it's best grandchild, we are done - if (compare(self.*, elem, best_grandchild.item) or elem == best_grandchild.item) return; + // If the item is better than or equal to its best grandchild, we are done + if (self.compareFn(best_grandchild.item, elem) != target_order) return; // Otherwise, swap them self.items[best_grandchild.index] = elem; @@ -238,16 +235,16 @@ pub fn PriorityDequeue(comptime T: type) type { index = best_grandchild.index; // We might need to swap the element with it's parent - self.swapIfParentIsBetter(elem, index, compare); + self.swapIfParentIsBetter(elem, index, target_order); } else { // The children or grandchildren are the last layer const first_child_index = firstChildIndex(index); if (first_child_index > self.len) return; - const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, compare); + const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order); - // If the best descendant is still larger, we are done - if (compare(self.*, elem, best_descendent.item) or elem == best_descendent.item) return; + // If the item is better than or equal to its best descendant, we are done + if (self.compareFn(best_descendent.item, elem) != target_order) return; // Otherwise swap them self.items[best_descendent.index] = elem; @@ -258,7 +255,7 @@ pub fn PriorityDequeue(comptime T: type) type { if (index < first_grandchild_index) return; // We might need to swap the element with it's parent - self.swapIfParentIsBetter(elem, index, compare); + self.swapIfParentIsBetter(elem, index, target_order); return; } @@ -267,11 +264,11 @@ pub fn PriorityDequeue(comptime T: type) type { } } - fn swapIfParentIsBetter(self: *Self, child: T, child_index: usize, compare: fn (Self, T, T) bool) void { + fn swapIfParentIsBetter(self: *Self, child: T, child_index: usize, target_order: Order) void { const parent_index = parentIndex(child_index); const parent = self.items[parent_index]; - if (compare(self.*, parent, child)) { + if (self.compareFn(parent, child) == target_order) { self.items[parent_index] = child; self.items[child_index] = parent; } @@ -289,21 +286,21 @@ pub fn PriorityDequeue(comptime T: type) type { }; } - fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, compare: fn (Self, T, T) bool) ItemAndIndex { - if (compare(self, item1.item, item2.item)) { + fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex { + if (self.compareFn(item1.item, item2.item) == target_order) { return item1; } else { return item2; } } - fn bestItemAtIndices(self: Self, index1: usize, index2: usize, compare: fn (Self, T, T) bool) ItemAndIndex { + fn bestItemAtIndices(self: Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex { var item1 = self.getItem(index1); var item2 = self.getItem(index2); - return self.bestItem(item1, item2, compare); + return self.bestItem(item1, item2, target_order); } - fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, compare: fn (Self, T, T) bool) ItemAndIndex { + fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, target_order: Order) ItemAndIndex { const second_child_index = first_child_index + 1; if (first_grandchild_index >= self.len) { // No grandchildren, find the best child (second may not exist) @@ -313,24 +310,24 @@ pub fn PriorityDequeue(comptime T: type) type { .index = first_child_index, }; } else { - return self.bestItemAtIndices(first_child_index, second_child_index, compare); + return self.bestItemAtIndices(first_child_index, second_child_index, target_order); } } const second_grandchild_index = first_grandchild_index + 1; if (second_grandchild_index >= self.len) { // One grandchild, so we know there is a second child. Compare first grandchild and second child - return self.bestItemAtIndices(first_grandchild_index, second_child_index, compare); + return self.bestItemAtIndices(first_grandchild_index, second_child_index, target_order); } - const best_left_grandchild_index = self.bestItemAtIndices(first_grandchild_index, second_grandchild_index, compare).index; + const best_left_grandchild_index = self.bestItemAtIndices(first_grandchild_index, second_grandchild_index, target_order).index; const third_grandchild_index = second_grandchild_index + 1; if (third_grandchild_index >= self.len) { // Two grandchildren, and we know the best. Compare this to second child. - return self.bestItemAtIndices(best_left_grandchild_index, second_child_index, compare); + return self.bestItemAtIndices(best_left_grandchild_index, second_child_index, target_order); } else { // Three grandchildren, compare the min of the first two with the third - return self.bestItemAtIndices(best_left_grandchild_index, third_grandchild_index, compare); + return self.bestItemAtIndices(best_left_grandchild_index, third_grandchild_index, target_order); } } @@ -348,12 +345,12 @@ pub fn PriorityDequeue(comptime T: type) type { /// Dequeue takes ownership of the passed in slice. The slice must have been /// allocated with `allocator`. /// De-initialize with `deinit`. - pub fn fromOwnedSlice(allocator: *Allocator, lessThanFn: fn (T, T) bool, items: []T) Self { + pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (T, T) Order, items: []T) Self { var queue = Self{ .items = items, .len = items.len, .allocator = allocator, - .lessThanFn = lessThanFn, + .compareFn = compareFn, }; if (queue.len <= 1) return queue; @@ -468,8 +465,8 @@ pub fn PriorityDequeue(comptime T: type) type { }; } -fn lessThanComparison(a: u32, b: u32) bool { - return a < b; +fn lessThanComparison(a: u32, b: u32) Order { + return std.math.order(a, b); } const PDQ = PriorityDequeue(u32); @@ -493,6 +490,32 @@ test "std.PriorityDequeue: add and remove min" { expectEqual(@as(u32, 54), queue.removeMin()); } +test "std.PriorityDequeue: add and remove min structs" { + const S = struct { + size: u32, + }; + var queue = PriorityDequeue(S).init(testing.allocator, struct { + fn order(a: S, b: S) Order { + return std.math.order(a.size, b.size); + } + }.order); + defer queue.deinit(); + + try queue.add(.{ .size = 54 }); + try queue.add(.{ .size = 12 }); + try queue.add(.{ .size = 7 }); + try queue.add(.{ .size = 23 }); + try queue.add(.{ .size = 25 }); + try queue.add(.{ .size = 13 }); + + expectEqual(@as(u32, 7), queue.removeMin().size); + expectEqual(@as(u32, 12), queue.removeMin().size); + expectEqual(@as(u32, 13), queue.removeMin().size); + expectEqual(@as(u32, 23), queue.removeMin().size); + expectEqual(@as(u32, 25), queue.removeMin().size); + expectEqual(@as(u32, 54), queue.removeMin().size); +} + test "std.PriorityDequeue: add and remove max" { var queue = PDQ.init(testing.allocator, lessThanComparison); defer queue.deinit(); diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig index f67b390608..c4b24588df 100644 --- a/lib/std/priority_queue.zig +++ b/lib/std/priority_queue.zig @@ -7,6 +7,7 @@ const std = @import("std.zig"); const Allocator = std.mem.Allocator; const assert = std.debug.assert; const warn = std.debug.warn; +const Order = std.math.Order; const testing = std.testing; const expect = testing.expect; const expectEqual = testing.expectEqual; @@ -20,15 +21,17 @@ pub fn PriorityQueue(comptime T: type) type { items: []T, len: usize, allocator: *Allocator, - compareFn: fn (a: T, b: T) bool, - - /// Initialize and return a priority queue. Provide - /// `compareFn` that returns `true` when its first argument - /// should get popped before its second argument. For example, - /// to make `pop` return the minimum value, provide + compareFn: fn (a: T, b: T) Order, + + /// Initialize and return a priority queue. Provide `compareFn` + /// that returns `Order.lt` when its first argument should + /// get popped before its second argument, `Order.eq` if the + /// arguments are of equal priority, or `Order.gt` if the second + /// argument should be popped first. For example, to make `pop` + /// return the smallest number, provide /// - /// `fn lessThan(a: T, b: T) bool { return a < b; }` - pub fn init(allocator: *Allocator, compareFn: fn (a: T, b: T) bool) Self { + /// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }` + pub fn init(allocator: *Allocator, compareFn: fn (a: T, b: T) Order) Self { return Self{ .items = &[_]T{}, .len = 0, @@ -61,7 +64,7 @@ pub fn PriorityQueue(comptime T: type) type { const child = self.items[child_index]; const parent = self.items[parent_index]; - if (!self.compareFn(child, parent)) break; + if (self.compareFn(child, parent) != .lt) break; self.items[parent_index] = child; self.items[child_index] = parent; @@ -133,14 +136,14 @@ pub fn PriorityQueue(comptime T: type) type { var smallest = self.items[index]; if (left) |e| { - if (self.compareFn(e, smallest)) { + if (self.compareFn(e, smallest) == .lt) { smallest_index = left_index; smallest = e; } } if (right) |e| { - if (self.compareFn(e, smallest)) { + if (self.compareFn(e, smallest) == .lt) { smallest_index = right_index; smallest = e; } @@ -159,7 +162,7 @@ pub fn PriorityQueue(comptime T: type) type { /// PriorityQueue takes ownership of the passed in slice. The slice must have been /// allocated with `allocator`. /// Deinitialize with `deinit`. - pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (a: T, b: T) bool, items: []T) Self { + pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (a: T, b: T) Order, items: []T) Self { var queue = Self{ .items = items, .len = items.len, @@ -217,10 +220,10 @@ pub fn PriorityQueue(comptime T: type) type { var update_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], elem) orelse return error.ElementNotFound; const old_elem: T = self.items[update_index]; self.items[update_index] = new_elem; - if (self.compareFn(new_elem, old_elem)) { - siftUp(self, update_index); - } else { - siftDown(self, update_index); + switch (self.compareFn(new_elem, old_elem)) { + .lt => siftUp(self, update_index), + .gt => siftDown(self, update_index), + .eq => {}, // Nothing to do as the items have equal priority } } @@ -267,12 +270,12 @@ pub fn PriorityQueue(comptime T: type) type { }; } -fn lessThan(a: u32, b: u32) bool { - return a < b; +fn lessThan(a: u32, b: u32) Order { + return std.math.order(a, b); } -fn greaterThan(a: u32, b: u32) bool { - return a > b; +fn greaterThan(a: u32, b: u32) Order { + return lessThan(a, b).invert(); } const PQ = PriorityQueue(u32); -- cgit v1.2.3