mirrors/zig
1
0
Fork 0
mirror of https://codeberg.org/ziglang/zig.git synced 2025-12-06 13:54:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Merge pull request #10342 from Arnavion/pq-comparefn-context

Browse source 
std.priority_{de,}queue: allow comparator to take context parameter
This commit is contained in:
Andrew Kelley 2021-12-17 14:34:32 -08:00 • committed by GitHub
commit e8b39960bb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 151 additions and 86 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

131
lib/std/priority_dequeue.zig
View file

@ -8,27 +8,29 @@ const expectEqual = testing.expectEqual;
const expectError = testing.expectError; const expectError = testing.expectError;
/// Priority Dequeue for storing generic data. Initialize with `init`. /// Priority Dequeue for storing generic data. Initialize with `init`.
/// Provide `compareFn` that returns `Order.lt` when its first /// Provide `compareFn` that returns `Order.lt` when its second
/// argument should get min-popped before its second argument, /// argument should get min-popped before its third argument,
/// `Order.eq` if the arguments are of equal priority, or `Order.gt` /// `Order.eq` if the arguments are of equal priority, or `Order.gt`
/// if the second argument should be min-popped first. /// if the third argument should be min-popped second.
/// Popping the max element works in reverse. For example, /// Popping the max element works in reverse. For example,
/// to make `popMin` return the smallest number, provide /// to make `popMin` return the smallest number, provide
/// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }` /// `fn lessThan(context: void, a: T, b: T) Order { _ = context; return std.math.order(a, b); }`
pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) type { pub fn PriorityDequeue(comptime T: type, comptime Context: type, comptime compareFn: fn (context: Context, a: T, b: T) Order) type {
return struct { return struct {
const Self = @This(); const Self = @This();
items: []T, items: []T,
len: usize, len: usize,
allocator: Allocator, allocator: Allocator,
context: Context,
/// Initialize and return a new priority dequeue. /// Initialize and return a new priority dequeue.
pub fn init(allocator: Allocator) Self { pub fn init(allocator: Allocator, context: Context) Self {
return Self{ return Self{
.items = &[_]T{}, .items = &[_]T{},
.len = 0, .len = 0,
.allocator = allocator, .allocator = allocator,
.context = context,
}; };
} }
@ -87,7 +89,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
const parent = self.items[parent_index]; const parent = self.items[parent_index];
const min_layer = self.nextIsMinLayer(); const min_layer = self.nextIsMinLayer();
const order = compareFn(child, parent); const order = compareFn(self.context, child, parent);
if ((min_layer and order == .gt) or (!min_layer and order == .lt)) { 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 // We must swap the item with it's parent if it is on the "wrong" layer
self.items[parent_index] = child; self.items[parent_index] = child;
@ -120,7 +122,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
const grandparent = self.items[grandparent_index]; const grandparent = self.items[grandparent_index];
// If the grandparent is already better or equal, we have gone as far as we need to // If the grandparent is already better or equal, we have gone as far as we need to
if (compareFn(child, grandparent) != target_order) break; if (compareFn(self.context, child, grandparent) != target_order) break;
// Otherwise swap the item with it's grandparent // Otherwise swap the item with it's grandparent
self.items[grandparent_index] = child; self.items[grandparent_index] = child;
@ -215,10 +217,10 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
// Find the best grandchild // Find the best grandchild
const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order); const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order);
const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order); const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order);
const best_grandchild = Self.bestItem(best_left, best_right, target_order); const best_grandchild = self.bestItem(best_left, best_right, target_order);
// If the item is better than or equal to its best grandchild, we are done // If the item is better than or equal to its best grandchild, we are done
if (compareFn(best_grandchild.item, elem) != target_order) return; if (compareFn(self.context, best_grandchild.item, elem) != target_order) return;
// Otherwise, swap them // Otherwise, swap them
self.items[best_grandchild.index] = elem; self.items[best_grandchild.index] = elem;
@ -235,7 +237,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order); const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order);
// If the item is better than or equal to its best descendant, we are done // If the item is better than or equal to its best descendant, we are done
if (compareFn(best_descendent.item, elem) != target_order) return; if (compareFn(self.context, best_descendent.item, elem) != target_order) return;
// Otherwise swap them // Otherwise swap them
self.items[best_descendent.index] = elem; self.items[best_descendent.index] = elem;
@ -259,7 +261,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
const parent_index = parentIndex(child_index); const parent_index = parentIndex(child_index);
const parent = self.items[parent_index]; const parent = self.items[parent_index];
if (compareFn(parent, child) == target_order) { if (compareFn(self.context, parent, child) == target_order) {
self.items[parent_index] = child; self.items[parent_index] = child;
self.items[child_index] = parent; self.items[child_index] = parent;
} }
@ -277,8 +279,8 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
}; };
} }
fn bestItem(item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex { fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex {
if (compareFn(item1.item, item2.item) == target_order) { if (compareFn(self.context, item1.item, item2.item) == target_order) {
return item1; return item1;
} else { } else {
return item2; return item2;
@ -288,7 +290,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
fn bestItemAtIndices(self: Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex { fn bestItemAtIndices(self: Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex {
var item1 = self.getItem(index1); var item1 = self.getItem(index1);
var item2 = self.getItem(index2); var item2 = self.getItem(index2);
return Self.bestItem(item1, item2, target_order); return self.bestItem(item1, item2, target_order);
} }
fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, target_order: Order) ItemAndIndex { fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, target_order: Order) ItemAndIndex {
@ -336,11 +338,12 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
/// Dequeue 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`. /// allocated with `allocator`.
/// De-initialize with `deinit`. /// De-initialize with `deinit`.
pub fn fromOwnedSlice(allocator: Allocator, items: []T) Self { pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {
var queue = Self{ var queue = Self{
.items = items, .items = items,
.len = items.len, .len = items.len,
.allocator = allocator, .allocator = allocator,
.context = context,
}; };
if (queue.len <= 1) return queue; if (queue.len <= 1) return queue;
@ -394,7 +397,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
} }
pub const Iterator = struct { pub const Iterator = struct {
queue: *PriorityDequeue(T, compareFn), queue: *PriorityDequeue(T, Context, compareFn),
count: usize, count: usize,
pub fn next(it: *Iterator) ?T { pub fn next(it: *Iterator) ?T {
@ -453,14 +456,15 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
}; };
} }
fn lessThanComparison(a: u32, b: u32) Order { fn lessThanComparison(context: void, a: u32, b: u32) Order {
_ = context;
return std.math.order(a, b); return std.math.order(a, b);
} }
const PDQ = PriorityDequeue(u32, lessThanComparison); const PDQ = PriorityDequeue(u32, void, lessThanComparison);
test "std.PriorityDequeue: add and remove min" { test "std.PriorityDequeue: add and remove min" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(54); try queue.add(54);
@ -482,11 +486,12 @@ test "std.PriorityDequeue: add and remove min structs" {
const S = struct { const S = struct {
size: u32, size: u32,
}; };
var queue = PriorityDequeue(S, struct { var queue = PriorityDequeue(S, void, struct {
fn order(a: S, b: S) Order { fn order(context: void, a: S, b: S) Order {
_ = context;
return std.math.order(a.size, b.size); return std.math.order(a.size, b.size);
} }
}.order).init(testing.allocator); }.order).init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(.{ .size = 54 }); try queue.add(.{ .size = 54 });
@ -505,7 +510,7 @@ test "std.PriorityDequeue: add and remove min structs" {
} }
test "std.PriorityDequeue: add and remove max" { test "std.PriorityDequeue: add and remove max" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(54); try queue.add(54);
@ -524,7 +529,7 @@ test "std.PriorityDequeue: add and remove max" {
} }
test "std.PriorityDequeue: add and remove same min" { test "std.PriorityDequeue: add and remove same min" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(1); try queue.add(1);
@ -543,7 +548,7 @@ test "std.PriorityDequeue: add and remove same min" {
} }
test "std.PriorityDequeue: add and remove same max" { test "std.PriorityDequeue: add and remove same max" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(1); try queue.add(1);
@ -562,7 +567,7 @@ test "std.PriorityDequeue: add and remove same max" {
} }
test "std.PriorityDequeue: removeOrNull empty" { test "std.PriorityDequeue: removeOrNull empty" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try expect(queue.removeMinOrNull() == null); try expect(queue.removeMinOrNull() == null);
@ -570,7 +575,7 @@ test "std.PriorityDequeue: removeOrNull empty" {
} }
test "std.PriorityDequeue: edge case 3 elements" { test "std.PriorityDequeue: edge case 3 elements" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(9); try queue.add(9);
@ -583,7 +588,7 @@ test "std.PriorityDequeue: edge case 3 elements" {
} }
test "std.PriorityDequeue: edge case 3 elements max" { test "std.PriorityDequeue: edge case 3 elements max" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(9); try queue.add(9);
@ -596,7 +601,7 @@ test "std.PriorityDequeue: edge case 3 elements max" {
} }
test "std.PriorityDequeue: peekMin" { test "std.PriorityDequeue: peekMin" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try expect(queue.peekMin() == null); try expect(queue.peekMin() == null);
@ -610,7 +615,7 @@ test "std.PriorityDequeue: peekMin" {
} }
test "std.PriorityDequeue: peekMax" { test "std.PriorityDequeue: peekMax" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try expect(queue.peekMin() == null); try expect(queue.peekMin() == null);
@ -624,7 +629,7 @@ test "std.PriorityDequeue: peekMax" {
} }
test "std.PriorityDequeue: sift up with odd indices" { test "std.PriorityDequeue: sift up with odd indices" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
for (items) |e| { for (items) |e| {
@ -638,7 +643,7 @@ test "std.PriorityDequeue: sift up with odd indices" {
} }
test "std.PriorityDequeue: sift up with odd indices" { test "std.PriorityDequeue: sift up with odd indices" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
for (items) |e| { for (items) |e| {
@ -652,7 +657,7 @@ test "std.PriorityDequeue: sift up with odd indices" {
} }
test "std.PriorityDequeue: addSlice min" { test "std.PriorityDequeue: addSlice min" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; 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(items[0..]);
@ -664,7 +669,7 @@ test "std.PriorityDequeue: addSlice min" {
} }
test "std.PriorityDequeue: addSlice max" { test "std.PriorityDequeue: addSlice max" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; 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(items[0..]);
@ -678,7 +683,7 @@ test "std.PriorityDequeue: addSlice max" {
test "std.PriorityDequeue: fromOwnedSlice trivial case 0" { test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {
const items = [0]u32{}; const items = [0]u32{};
const queue_items = try testing.allocator.dupe(u32, &items); const queue_items = try testing.allocator.dupe(u32, &items);
var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..]); var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..], {});
defer queue.deinit(); defer queue.deinit();
try expectEqual(@as(usize, 0), queue.len); try expectEqual(@as(usize, 0), queue.len);
try expect(queue.removeMinOrNull() == null); try expect(queue.removeMinOrNull() == null);
@ -687,7 +692,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {
test "std.PriorityDequeue: fromOwnedSlice trivial case 1" { test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {
const items = [1]u32{1}; const items = [1]u32{1};
const queue_items = try testing.allocator.dupe(u32, &items); const queue_items = try testing.allocator.dupe(u32, &items);
var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..]); var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..], {});
defer queue.deinit(); defer queue.deinit();
try expectEqual(@as(usize, 1), queue.len); try expectEqual(@as(usize, 1), queue.len);
@ -698,7 +703,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {
test "std.PriorityDequeue: fromOwnedSlice" { 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 items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
const queue_items = try testing.allocator.dupe(u32, items[0..]); const queue_items = try testing.allocator.dupe(u32, items[0..]);
var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..]); var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..], {});
defer queue.deinit(); 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 }; const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
@ -708,7 +713,7 @@ test "std.PriorityDequeue: fromOwnedSlice" {
} }
test "std.PriorityDequeue: update min queue" { test "std.PriorityDequeue: update min queue" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(55); try queue.add(55);
@ -723,7 +728,7 @@ test "std.PriorityDequeue: update min queue" {
} }
test "std.PriorityDequeue: update same min queue" { test "std.PriorityDequeue: update same min queue" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(1); try queue.add(1);
@ -739,7 +744,7 @@ test "std.PriorityDequeue: update same min queue" {
} }
test "std.PriorityDequeue: update max queue" { test "std.PriorityDequeue: update max queue" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(55); try queue.add(55);
@ -755,7 +760,7 @@ test "std.PriorityDequeue: update max queue" {
} }
test "std.PriorityDequeue: update same max queue" { test "std.PriorityDequeue: update same max queue" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(1); try queue.add(1);
@ -771,7 +776,7 @@ test "std.PriorityDequeue: update same max queue" {
} }
test "std.PriorityDequeue: iterator" { test "std.PriorityDequeue: iterator" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
var map = std.AutoHashMap(u32, void).init(testing.allocator); var map = std.AutoHashMap(u32, void).init(testing.allocator);
defer { defer {
queue.deinit(); queue.deinit();
@ -793,7 +798,7 @@ test "std.PriorityDequeue: iterator" {
} }
test "std.PriorityDequeue: remove at index" { test "std.PriorityDequeue: remove at index" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(3); try queue.add(3);
@ -816,7 +821,7 @@ test "std.PriorityDequeue: remove at index" {
} }
test "std.PriorityDequeue: iterator while empty" { test "std.PriorityDequeue: iterator while empty" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
var it = queue.iterator(); var it = queue.iterator();
@ -825,7 +830,7 @@ test "std.PriorityDequeue: iterator while empty" {
} }
test "std.PriorityDequeue: shrinkAndFree" { test "std.PriorityDequeue: shrinkAndFree" {
var queue = PDQ.init(testing.allocator); var queue = PDQ.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.ensureTotalCapacity(4); try queue.ensureTotalCapacity(4);
@ -864,7 +869,7 @@ fn fuzzTestMin(rng: std.rand.Random, comptime queue_size: usize) !void {
const allocator = testing.allocator; const allocator = testing.allocator;
const items = try generateRandomSlice(allocator, rng, queue_size); const items = try generateRandomSlice(allocator, rng, queue_size);
var queue = PDQ.fromOwnedSlice(allocator, items); var queue = PDQ.fromOwnedSlice(allocator, items, {});
defer queue.deinit(); defer queue.deinit();
var last_removed: ?u32 = null; var last_removed: ?u32 = null;
@ -893,7 +898,7 @@ fn fuzzTestMax(rng: std.rand.Random, queue_size: usize) !void {
const allocator = testing.allocator; const allocator = testing.allocator;
const items = try generateRandomSlice(allocator, rng, queue_size); const items = try generateRandomSlice(allocator, rng, queue_size);
var queue = PDQ.fromOwnedSlice(testing.allocator, items); var queue = PDQ.fromOwnedSlice(testing.allocator, items, {});
defer queue.deinit(); defer queue.deinit();
var last_removed: ?u32 = null; var last_removed: ?u32 = null;
@ -922,7 +927,7 @@ fn fuzzTestMinMax(rng: std.rand.Random, queue_size: usize) !void {
const allocator = testing.allocator; const allocator = testing.allocator;
const items = try generateRandomSlice(allocator, rng, queue_size); const items = try generateRandomSlice(allocator, rng, queue_size);
var queue = PDQ.fromOwnedSlice(allocator, items); var queue = PDQ.fromOwnedSlice(allocator, items, {});
defer queue.deinit(); defer queue.deinit();
var last_min: ?u32 = null; var last_min: ?u32 = null;
@ -957,3 +962,31 @@ fn generateRandomSlice(allocator: std.mem.Allocator, rng: std.rand.Random, size:
return array.toOwnedSlice(); return array.toOwnedSlice();
} }
fn contextLessThanComparison(context: []const u32, a: usize, b: usize) Order {
return std.math.order(context[a], context[b]);
}
const CPDQ = PriorityDequeue(usize, []const u32, contextLessThanComparison);
test "std.PriorityDequeue: add and remove" {
const context = [_]u32{ 5, 3, 4, 2, 2, 8, 0 };
var queue = CPDQ.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);
try expectEqual(@as(usize, 6), queue.removeMin());
try expectEqual(@as(usize, 5), queue.removeMax());
try expectEqual(@as(usize, 3), queue.removeMin());
try expectEqual(@as(usize, 0), queue.removeMax());
try expectEqual(@as(usize, 4), queue.removeMin());
try expectEqual(@as(usize, 2), queue.removeMax());
try expectEqual(@as(usize, 1), queue.removeMin());
}

106
lib/std/priority_queue.zig
View file

@ -8,26 +8,28 @@ const expectEqual = testing.expectEqual;
const expectError = testing.expectError; const expectError = testing.expectError;
/// Priority queue for storing generic data. Initialize with `init`. /// Priority queue for storing generic data. Initialize with `init`.
/// Provide `compareFn` that returns `Order.lt` when its first /// Provide `compareFn` that returns `Order.lt` when its second
/// argument should get popped before its second argument, /// argument should get popped before its third argument,
/// `Order.eq` if the arguments are of equal priority, or `Order.gt` /// `Order.eq` if the arguments are of equal priority, or `Order.gt`
/// if the second argument should be popped first. /// if the third argument should be popped first.
/// For example, to make `pop` return the smallest number, provide /// For example, to make `pop` return the smallest number, provide
/// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }` /// `fn lessThan(context: void, a: T, b: T) Order { _ = context; return std.math.order(a, b); }`
pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order) type { pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareFn: fn (context: Context, a: T, b: T) Order) type {
return struct { return struct {
const Self = @This(); const Self = @This();
items: []T, items: []T,
len: usize, len: usize,
allocator: Allocator, allocator: Allocator,
context: Context,
/// Initialize and return a priority queue. /// Initialize and return a priority queue.
pub fn init(allocator: Allocator) Self { pub fn init(allocator: Allocator, context: Context) Self {
return Self{ return Self{
.items = &[_]T{}, .items = &[_]T{},
.len = 0, .len = 0,
.allocator = allocator, .allocator = allocator,
.context = context,
}; };
} }
@ -55,7 +57,7 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
const child = self.items[child_index]; const child = self.items[child_index];
const parent = self.items[parent_index]; const parent = self.items[parent_index];
if (compareFn(child, parent) != .lt) break; if (compareFn(self.context, child, parent) != .lt) break;
self.items[parent_index] = child; self.items[parent_index] = child;
self.items[child_index] = parent; self.items[child_index] = parent;
@ -127,14 +129,14 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
var smallest = self.items[index]; var smallest = self.items[index];
if (left) |e| { if (left) |e| {
if (compareFn(e, smallest) == .lt) { if (compareFn(self.context, e, smallest) == .lt) {
smallest_index = left_index; smallest_index = left_index;
smallest = e; smallest = e;
} }
} }
if (right) |e| { if (right) |e| {
if (compareFn(e, smallest) == .lt) { if (compareFn(self.context, e, smallest) == .lt) {
smallest_index = right_index; smallest_index = right_index;
smallest = e; smallest = e;
} }
@ -153,11 +155,12 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
/// PriorityQueue takes ownership of the passed in slice. The slice must have been /// PriorityQueue takes ownership of the passed in slice. The slice must have been
/// allocated with `allocator`. /// allocated with `allocator`.
/// Deinitialize with `deinit`. /// Deinitialize with `deinit`.
pub fn fromOwnedSlice(allocator: Allocator, items: []T) Self { pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {
var queue = Self{ var queue = Self{
.items = items, .items = items,
.len = items.len, .len = items.len,
.allocator = allocator, .allocator = allocator,
.context = context,
}; };
if (queue.len <= 1) return queue; if (queue.len <= 1) return queue;
@ -207,7 +210,7 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
var update_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], 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]; const old_elem: T = self.items[update_index];
self.items[update_index] = new_elem; self.items[update_index] = new_elem;
switch (compareFn(new_elem, old_elem)) { switch (compareFn(self.context, new_elem, old_elem)) {
.lt => siftUp(self, update_index), .lt => siftUp(self, update_index),
.gt => siftDown(self, update_index), .gt => siftDown(self, update_index),
.eq => {}, // Nothing to do as the items have equal priority .eq => {}, // Nothing to do as the items have equal priority
@ -215,7 +218,7 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
} }
pub const Iterator = struct { pub const Iterator = struct {
queue: *PriorityQueue(T, compareFn), queue: *PriorityQueue(T, Context, compareFn),
count: usize, count: usize,
pub fn next(it: *Iterator) ?T { pub fn next(it: *Iterator) ?T {
@ -258,19 +261,20 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
}; };
} }
fn lessThan(a: u32, b: u32) Order { fn lessThan(context: void, a: u32, b: u32) Order {
_ = context;
return std.math.order(a, b); return std.math.order(a, b);
} }
fn greaterThan(a: u32, b: u32) Order { fn greaterThan(context: void, a: u32, b: u32) Order {
return lessThan(a, b).invert(); return lessThan(context, a, b).invert();
} }
const PQlt = PriorityQueue(u32, lessThan); const PQlt = PriorityQueue(u32, void, lessThan);
const PQgt = PriorityQueue(u32, greaterThan); const PQgt = PriorityQueue(u32, void, greaterThan);
test "std.PriorityQueue: add and remove min heap" { test "std.PriorityQueue: add and remove min heap" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(54); try queue.add(54);
@ -288,7 +292,7 @@ test "std.PriorityQueue: add and remove min heap" {
} }
test "std.PriorityQueue: add and remove same min heap" { test "std.PriorityQueue: add and remove same min heap" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(1); try queue.add(1);
@ -306,14 +310,14 @@ test "std.PriorityQueue: add and remove same min heap" {
} }
test "std.PriorityQueue: removeOrNull on empty" { test "std.PriorityQueue: removeOrNull on empty" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try expect(queue.removeOrNull() == null); try expect(queue.removeOrNull() == null);
} }
test "std.PriorityQueue: edge case 3 elements" { test "std.PriorityQueue: edge case 3 elements" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(9); try queue.add(9);
@ -325,7 +329,7 @@ test "std.PriorityQueue: edge case 3 elements" {
} }
test "std.PriorityQueue: peek" { test "std.PriorityQueue: peek" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try expect(queue.peek() == null); try expect(queue.peek() == null);
@ -337,7 +341,7 @@ test "std.PriorityQueue: peek" {
} }
test "std.PriorityQueue: sift up with odd indices" { test "std.PriorityQueue: sift up with odd indices" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
for (items) |e| { for (items) |e| {
@ -351,7 +355,7 @@ test "std.PriorityQueue: sift up with odd indices" {
} }
test "std.PriorityQueue: addSlice" { test "std.PriorityQueue: addSlice" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 }; 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(items[0..]);
@ -365,7 +369,7 @@ test "std.PriorityQueue: addSlice" {
test "std.PriorityQueue: fromOwnedSlice trivial case 0" { test "std.PriorityQueue: fromOwnedSlice trivial case 0" {
const items = [0]u32{}; const items = [0]u32{};
const queue_items = try testing.allocator.dupe(u32, &items); const queue_items = try testing.allocator.dupe(u32, &items);
var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..]); var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {});
defer queue.deinit(); defer queue.deinit();
try expectEqual(@as(usize, 0), queue.len); try expectEqual(@as(usize, 0), queue.len);
try expect(queue.removeOrNull() == null); try expect(queue.removeOrNull() == null);
@ -374,7 +378,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 0" {
test "std.PriorityQueue: fromOwnedSlice trivial case 1" { test "std.PriorityQueue: fromOwnedSlice trivial case 1" {
const items = [1]u32{1}; const items = [1]u32{1};
const queue_items = try testing.allocator.dupe(u32, &items); const queue_items = try testing.allocator.dupe(u32, &items);
var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..]); var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {});
defer queue.deinit(); defer queue.deinit();
try expectEqual(@as(usize, 1), queue.len); try expectEqual(@as(usize, 1), queue.len);
@ -385,7 +389,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 1" {
test "std.PriorityQueue: fromOwnedSlice" { 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 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..]); const heap_items = try testing.allocator.dupe(u32, items[0..]);
var queue = PQlt.fromOwnedSlice(testing.allocator, heap_items[0..]); var queue = PQlt.fromOwnedSlice(testing.allocator, heap_items[0..], {});
defer queue.deinit(); 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 }; const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
@ -395,7 +399,7 @@ test "std.PriorityQueue: fromOwnedSlice" {
} }
test "std.PriorityQueue: add and remove max heap" { test "std.PriorityQueue: add and remove max heap" {
var queue = PQgt.init(testing.allocator); var queue = PQgt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(54); try queue.add(54);
@ -413,7 +417,7 @@ test "std.PriorityQueue: add and remove max heap" {
} }
test "std.PriorityQueue: add and remove same max heap" { test "std.PriorityQueue: add and remove same max heap" {
var queue = PQgt.init(testing.allocator); var queue = PQgt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(1); try queue.add(1);
@ -431,7 +435,7 @@ test "std.PriorityQueue: add and remove same max heap" {
} }
test "std.PriorityQueue: iterator" { test "std.PriorityQueue: iterator" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
var map = std.AutoHashMap(u32, void).init(testing.allocator); var map = std.AutoHashMap(u32, void).init(testing.allocator);
defer { defer {
queue.deinit(); queue.deinit();
@ -453,7 +457,7 @@ test "std.PriorityQueue: iterator" {
} }
test "std.PriorityQueue: remove at index" { test "std.PriorityQueue: remove at index" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(3); try queue.add(3);
@ -476,7 +480,7 @@ test "std.PriorityQueue: remove at index" {
} }
test "std.PriorityQueue: iterator while empty" { test "std.PriorityQueue: iterator while empty" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
var it = queue.iterator(); var it = queue.iterator();
@ -485,7 +489,7 @@ test "std.PriorityQueue: iterator while empty" {
} }
test "std.PriorityQueue: shrinkAndFree" { test "std.PriorityQueue: shrinkAndFree" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.ensureTotalCapacity(4); try queue.ensureTotalCapacity(4);
@ -508,7 +512,7 @@ test "std.PriorityQueue: shrinkAndFree" {
} }
test "std.PriorityQueue: update min heap" { test "std.PriorityQueue: update min heap" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(55); try queue.add(55);
@ -523,7 +527,7 @@ test "std.PriorityQueue: update min heap" {
} }
test "std.PriorityQueue: update same min heap" { test "std.PriorityQueue: update same min heap" {
var queue = PQlt.init(testing.allocator); var queue = PQlt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(1); try queue.add(1);
@ -539,7 +543,7 @@ test "std.PriorityQueue: update same min heap" {
} }
test "std.PriorityQueue: update max heap" { test "std.PriorityQueue: update max heap" {
var queue = PQgt.init(testing.allocator); var queue = PQgt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(55); try queue.add(55);
@ -554,7 +558,7 @@ test "std.PriorityQueue: update max heap" {
} }
test "std.PriorityQueue: update same max heap" { test "std.PriorityQueue: update same max heap" {
var queue = PQgt.init(testing.allocator); var queue = PQgt.init(testing.allocator, {});
defer queue.deinit(); defer queue.deinit();
try queue.add(1); try queue.add(1);
@ -568,3 +572,31 @@ test "std.PriorityQueue: update same max heap" {
try expectEqual(@as(u32, 2), queue.remove()); try expectEqual(@as(u32, 2), queue.remove());
try expectEqual(@as(u32, 1), queue.remove()); try expectEqual(@as(u32, 1), queue.remove());
} }
fn contextLessThan(context: []const u32, a: usize, b: usize) Order {
return std.math.order(context[a], context[b]);
}
const CPQlt = PriorityQueue(usize, []const u32, contextLessThan);
test "std.PriorityQueue: add and remove min heap with contextful 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);
try expectEqual(@as(usize, 6), queue.remove());
try expectEqual(@as(usize, 4), queue.remove());
try expectEqual(@as(usize, 3), queue.remove());
try expectEqual(@as(usize, 1), queue.remove());
try expectEqual(@as(usize, 2), queue.remove());
try expectEqual(@as(usize, 0), queue.remove());
try expectEqual(@as(usize, 5), queue.remove());
}