motiejus/zig
1
Fork 0
Code Releases Activity

Move compareFn from init to type constructor in PriorityQueue and PriorityDequeue.

Browse Source 
This change significantly improves performance for simple compare functions and modifies
the API to be more consistent with e.g. `HashMap`.
This commit is contained in:
Max Hollmann
2021-10-17 01:10:01 +02:00
committed by Andrew Kelley
parent 22beaf5afc
commit 53b87fa78a
2 changed files with 88 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
lib/std/priority_queue.zig
View File

@@ -9,29 +9,26 @@ const expectEqual = testing.expectEqual;
const expectError = testing.expectError;
/// Priority queue for storing generic data. Initialize with `init`.
pub fn PriorityQueue(comptime T: type) type {
/// 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) Order { return std.math.order(a, b); }`
pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order) type {
return struct {
const Self = @This();
items: []T,
len: usize,
allocator: *Allocator,
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) Order { return std.math.order(a, b); }`
pub fn init(allocator: *Allocator, compareFn: fn (a: T, b: T) Order) Self {
/// Initialize and return a priority queue.
pub fn init(allocator: *Allocator) Self {
return Self{
.items = &[_]T{},
.len = 0,
.allocator = allocator,
.compareFn = compareFn,
};
}
@@ -59,7 +56,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) != .lt) break;
if (compareFn(child, parent) != .lt) break;
self.items[parent_index] = child;
self.items[child_index] = parent;
@@ -131,14 +128,14 @@ pub fn PriorityQueue(comptime T: type) type {
var smallest = self.items[index];
if (left) |e| {
if (self.compareFn(e, smallest) == .lt) {
if (compareFn(e, smallest) == .lt) {
smallest_index = left_index;
smallest = e;
}
}
if (right) |e| {
if (self.compareFn(e, smallest) == .lt) {
if (compareFn(e, smallest) == .lt) {
smallest_index = right_index;
smallest = e;
}
@@ -157,12 +154,11 @@ 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) Order, items: []T) Self {
pub fn fromOwnedSlice(allocator: *Allocator, items: []T) Self {
var queue = Self{
.items = items,
.len = items.len,
.allocator = allocator,
.compareFn = compareFn,
};
if (queue.len <= 1) return queue;
@@ -213,7 +209,7 @@ 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;
switch (self.compareFn(new_elem, old_elem)) {
switch (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
@@ -221,7 +217,7 @@ pub fn PriorityQueue(comptime T: type) type {
}
pub const Iterator = struct {
queue: *PriorityQueue(T),
queue: *PriorityQueue(T, compareFn),
count: usize,
pub fn next(it: *Iterator) ?T {
@@ -271,10 +267,11 @@ fn greaterThan(a: u32, b: u32) Order {
return lessThan(a, b).invert();
}
const PQ = PriorityQueue(u32);
const PQlt = PriorityQueue(u32, lessThan);
const PQgt = PriorityQueue(u32, greaterThan);
test "std.PriorityQueue: add and remove min heap" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
defer queue.deinit();
try queue.add(54);
@@ -292,7 +289,7 @@ test "std.PriorityQueue: add and remove min heap" {
}
test "std.PriorityQueue: add and remove same min heap" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
defer queue.deinit();
try queue.add(1);
@@ -310,14 +307,14 @@ test "std.PriorityQueue: add and remove same min heap" {
}
test "std.PriorityQueue: removeOrNull on empty" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
defer queue.deinit();
try expect(queue.removeOrNull() == null);
}
test "std.PriorityQueue: edge case 3 elements" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
defer queue.deinit();
try queue.add(9);
@@ -329,7 +326,7 @@ test "std.PriorityQueue: edge case 3 elements" {
}
test "std.PriorityQueue: peek" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
defer queue.deinit();
try expect(queue.peek() == null);
@@ -341,7 +338,7 @@ test "std.PriorityQueue: peek" {
}
test "std.PriorityQueue: sift up with odd indices" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
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| {
@@ -355,7 +352,7 @@ test "std.PriorityQueue: sift up with odd indices" {
}
test "std.PriorityQueue: addSlice" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
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 queue.addSlice(items[0..]);
@@ -369,7 +366,7 @@ 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..]);
var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..]);
defer queue.deinit();
try expectEqual(@as(usize, 0), queue.len);
try expect(queue.removeOrNull() == null);
@@ -378,7 +375,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 0" {
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..]);
var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..]);
defer queue.deinit();
try expectEqual(@as(usize, 1), queue.len);
@@ -389,7 +386,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 1" {
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..]);
var queue = PQ.fromOwnedSlice(testing.allocator, lessThan, heap_items[0..]);
var queue = PQlt.fromOwnedSlice(testing.allocator, heap_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 };
@@ -399,7 +396,7 @@ test "std.PriorityQueue: fromOwnedSlice" {
}
test "std.PriorityQueue: add and remove max heap" {
var queue = PQ.init(testing.allocator, greaterThan);
var queue = PQgt.init(testing.allocator);
defer queue.deinit();
try queue.add(54);
@@ -417,7 +414,7 @@ test "std.PriorityQueue: add and remove max heap" {
}
test "std.PriorityQueue: add and remove same max heap" {
var queue = PQ.init(testing.allocator, greaterThan);
var queue = PQgt.init(testing.allocator);
defer queue.deinit();
try queue.add(1);
@@ -435,7 +432,7 @@ test "std.PriorityQueue: add and remove same max heap" {
}
test "std.PriorityQueue: iterator" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
var map = std.AutoHashMap(u32, void).init(testing.allocator);
defer {
queue.deinit();
@@ -457,7 +454,7 @@ test "std.PriorityQueue: iterator" {
}
test "std.PriorityQueue: remove at index" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
defer queue.deinit();
try queue.add(3);
@@ -480,7 +477,7 @@ test "std.PriorityQueue: remove at index" {
}
test "std.PriorityQueue: iterator while empty" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
defer queue.deinit();
var it = queue.iterator();
@@ -489,7 +486,7 @@ test "std.PriorityQueue: iterator while empty" {
}
test "std.PriorityQueue: shrinkAndFree" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
defer queue.deinit();
try queue.ensureTotalCapacity(4);
@@ -512,7 +509,7 @@ test "std.PriorityQueue: shrinkAndFree" {
}
test "std.PriorityQueue: update min heap" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
defer queue.deinit();
try queue.add(55);
@@ -527,7 +524,7 @@ test "std.PriorityQueue: update min heap" {
}
test "std.PriorityQueue: update same min heap" {
var queue = PQ.init(testing.allocator, lessThan);
var queue = PQlt.init(testing.allocator);
defer queue.deinit();
try queue.add(1);
@@ -543,7 +540,7 @@ test "std.PriorityQueue: update same min heap" {
}
test "std.PriorityQueue: update max heap" {
var queue = PQ.init(testing.allocator, greaterThan);
var queue = PQgt.init(testing.allocator);
defer queue.deinit();
try queue.add(55);
@@ -558,7 +555,7 @@ test "std.PriorityQueue: update max heap" {
}
test "std.PriorityQueue: update same max heap" {
var queue = PQ.init(testing.allocator, greaterThan);
var queue = PQgt.init(testing.allocator);
defer queue.deinit();
try queue.add(1);