authorgravatar for git@zander.xyzZander Khan <git@zander.xyz> 2021-01-18 19:02:11+00:00
committergravatar for git@zander.xyzZander Khan <git@zander.xyz> 2021-01-18 19:02:11+00:00
logce22c70586137bd20de101f3b1ec3a44fbdb54e7
treeb5f0d8fe30e27207502baf1abfe96c9ec87f5eca
parent5bfd9238de82b8f66ba4834bcf7aa120e42fe6ca

Change `compareFn` to `fn (a: T, b: T) std.math.Order`

2 files changed, 100 insertions(+), 74 deletions(-)

lib/std/priority_dequeue.zig+77-54
......@@ -7,6 +7,7 @@ const std = @import("std.zig");
77const Allocator = std.mem.Allocator;
88const assert = std.debug.assert;
99const warn = std.debug.warn;
10const Order = std.math.Order;
1011const testing = std.testing;
1112const expect = testing.expect;
1213const expectEqual = testing.expectEqual;
......@@ -20,31 +21,26 @@ pub fn PriorityDequeue(comptime T: type) type {
2021 items: []T,
2122 len: usize,
2223 allocator: *Allocator,
23 lessThanFn: fn (a: T, b: T) bool,
24
25 /// Initialize and return a new dequeue. Provide `lessThanFn`
26 /// that returns `true` when its first argument should
27 /// get min-popped before its second argument. For example,
28 /// to make `popMin` return the minimum value, provide
24 compareFn: fn (a: T, b: T) Order,
25
26 /// Initialize and return a new priority dequeue. Provide `compareFn`
27 /// that returns `Order.lt` when its first argument should
28 /// get min-popped before its second argument, `Order.eq` if the
29 /// arguments are of equal priority, or `Order.gt` if the second
30 /// argument should be min-popped first. Popping the max element works
31 /// in reverse. For example, to make `popMin` return the smallest
32 /// number, provide
2933 ///
30 /// `fn lessThanFn(a: T, b: T) bool { return a < b; }`
31 pub fn init(allocator: *Allocator, lessThanFn: fn (T, T) bool) Self {
34 /// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`
35 pub fn init(allocator: *Allocator, compareFn: fn (T, T) Order) Self {
3236 return Self{
3337 .items = &[_]T{},
3438 .len = 0,
3539 .allocator = allocator,
36 .lessThanFn = lessThanFn,
40 .compareFn = compareFn,
3741 };
3842 }
3943
40 fn lessThan(self: Self, a: T, b: T) bool {
41 return self.lessThanFn(a, b);
42 }
43
44 fn greaterThan(self: Self, a: T, b: T) bool {
45 return self.lessThanFn(b, a);
46 }
47
4844 /// Free memory used by the dequeue.
4945 pub fn deinit(self: Self) void {
5046 self.allocator.free(self.items);
......@@ -100,7 +96,8 @@ pub fn PriorityDequeue(comptime T: type) type {
10096 const parent = self.items[parent_index];
10197
10298 const min_layer = self.nextIsMinLayer();
103 if ((min_layer and self.greaterThan(child, parent)) or (!min_layer and self.lessThan(child, parent))) {
99 const order = self.compareFn(child, parent);
100 if ((min_layer and order == .gt) or (!min_layer and order == .lt)) {
104101 // We must swap the item with it's parent if it is on the "wrong" layer
105102 self.items[parent_index] = child;
106103 self.items[child_index] = parent;
......@@ -118,21 +115,21 @@ pub fn PriorityDequeue(comptime T: type) type {
118115
119116 fn siftUp(self: *Self, start: StartIndexAndLayer) void {
120117 if (start.min_layer) {
121 doSiftUp(self, start.index, lessThan);
118 doSiftUp(self, start.index, .lt);
122119 } else {
123 doSiftUp(self, start.index, greaterThan);
120 doSiftUp(self, start.index, .gt);
124121 }
125122 }
126123
127 fn doSiftUp(self: *Self, start_index: usize, compare: fn (Self, T, T) bool) void {
124 fn doSiftUp(self: *Self, start_index: usize, target_order: Order) void {
128125 var child_index = start_index;
129126 while (child_index > 2) {
130127 var grandparent_index = grandparentIndex(child_index);
131128 const child = self.items[child_index];
132129 const grandparent = self.items[grandparent_index];
133130
134 // If the grandparent is already better, we have gone as far as we need to
135 if (!compare(self.*, child, grandparent)) break;
131 // If the grandparent is already better or equal, we have gone as far as we need to
132 if (self.compareFn(child, grandparent) != target_order) break;
136133
137134 // Otherwise swap the item with it's grandparent
138135 self.items[grandparent_index] = child;
......@@ -153,14 +150,14 @@ pub fn PriorityDequeue(comptime T: type) type {
153150 if (self.len == 0) return null;
154151 if (self.len == 1) return self.items[0];
155152 if (self.len == 2) return self.items[1];
156 return self.bestItemAtIndices(1, 2, greaterThan).item;
153 return self.bestItemAtIndices(1, 2, .gt).item;
157154 }
158155
159156 fn maxIndex(self: Self) ?usize {
160157 if (self.len == 0) return null;
161158 if (self.len == 1) return 0;
162159 if (self.len == 2) return 1;
163 return self.bestItemAtIndices(1, 2, greaterThan).index;
160 return self.bestItemAtIndices(1, 2, .gt).index;
164161 }
165162
166163 /// Pop the smallest element from the dequeue. Returns
......@@ -204,13 +201,13 @@ pub fn PriorityDequeue(comptime T: type) type {
204201
205202 fn siftDown(self: *Self, index: usize) void {
206203 if (isMinLayer(index)) {
207 self.doSiftDown(index, lessThan);
204 self.doSiftDown(index, .lt);
208205 } else {
209 self.doSiftDown(index, greaterThan);
206 self.doSiftDown(index, .gt);
210207 }
211208 }
212209
213 fn doSiftDown(self: *Self, start_index: usize, compare: fn (Self, T, T) bool) void {
210 fn doSiftDown(self: *Self, start_index: usize, target_order: Order) void {
214211 var index = start_index;
215212 const half = self.len >> 1;
216213 while (true) {
......@@ -225,12 +222,12 @@ pub fn PriorityDequeue(comptime T: type) type {
225222 const index3 = index2 + 1;
226223
227224 // Find the best grandchild
228 const best_left = self.bestItemAtIndices(first_grandchild_index, index2, compare);
229 const best_right = self.bestItemAtIndices(index3, last_grandchild_index, compare);
230 const best_grandchild = self.bestItem(best_left, best_right, compare);
225 const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order);
226 const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order);
227 const best_grandchild = self.bestItem(best_left, best_right, target_order);
231228
232 // If the item is better than it's best grandchild, we are done
233 if (compare(self.*, elem, best_grandchild.item) or elem == best_grandchild.item) return;
229 // If the item is better than or equal to its best grandchild, we are done
230 if (self.compareFn(best_grandchild.item, elem) != target_order) return;
234231
235232 // Otherwise, swap them
236233 self.items[best_grandchild.index] = elem;
......@@ -238,16 +235,16 @@ pub fn PriorityDequeue(comptime T: type) type {
238235 index = best_grandchild.index;
239236
240237 // We might need to swap the element with it's parent
241 self.swapIfParentIsBetter(elem, index, compare);
238 self.swapIfParentIsBetter(elem, index, target_order);
242239 } else {
243240 // The children or grandchildren are the last layer
244241 const first_child_index = firstChildIndex(index);
245242 if (first_child_index > self.len) return;
246243
247 const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, compare);
244 const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order);
248245
249 // If the best descendant is still larger, we are done
250 if (compare(self.*, elem, best_descendent.item) or elem == best_descendent.item) return;
246 // If the item is better than or equal to its best descendant, we are done
247 if (self.compareFn(best_descendent.item, elem) != target_order) return;
251248
252249 // Otherwise swap them
253250 self.items[best_descendent.index] = elem;
......@@ -258,7 +255,7 @@ pub fn PriorityDequeue(comptime T: type) type {
258255 if (index < first_grandchild_index) return;
259256
260257 // We might need to swap the element with it's parent
261 self.swapIfParentIsBetter(elem, index, compare);
258 self.swapIfParentIsBetter(elem, index, target_order);
262259 return;
263260 }
264261
......@@ -267,11 +264,11 @@ pub fn PriorityDequeue(comptime T: type) type {
267264 }
268265 }
269266
270 fn swapIfParentIsBetter(self: *Self, child: T, child_index: usize, compare: fn (Self, T, T) bool) void {
267 fn swapIfParentIsBetter(self: *Self, child: T, child_index: usize, target_order: Order) void {
271268 const parent_index = parentIndex(child_index);
272269 const parent = self.items[parent_index];
273270
274 if (compare(self.*, parent, child)) {
271 if (self.compareFn(parent, child) == target_order) {
275272 self.items[parent_index] = child;
276273 self.items[child_index] = parent;
277274 }
......@@ -289,21 +286,21 @@ pub fn PriorityDequeue(comptime T: type) type {
289286 };
290287 }
291288
292 fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, compare: fn (Self, T, T) bool) ItemAndIndex {
293 if (compare(self, item1.item, item2.item)) {
289 fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex {
290 if (self.compareFn(item1.item, item2.item) == target_order) {
294291 return item1;
295292 } else {
296293 return item2;
297294 }
298295 }
299296
300 fn bestItemAtIndices(self: Self, index1: usize, index2: usize, compare: fn (Self, T, T) bool) ItemAndIndex {
297 fn bestItemAtIndices(self: Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex {
301298 var item1 = self.getItem(index1);
302299 var item2 = self.getItem(index2);
303 return self.bestItem(item1, item2, compare);
300 return self.bestItem(item1, item2, target_order);
304301 }
305302
306 fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, compare: fn (Self, T, T) bool) ItemAndIndex {
303 fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, target_order: Order) ItemAndIndex {
307304 const second_child_index = first_child_index + 1;
308305 if (first_grandchild_index >= self.len) {
309306 // No grandchildren, find the best child (second may not exist)
......@@ -313,24 +310,24 @@ pub fn PriorityDequeue(comptime T: type) type {
313310 .index = first_child_index,
314311 };
315312 } else {
316 return self.bestItemAtIndices(first_child_index, second_child_index, compare);
313 return self.bestItemAtIndices(first_child_index, second_child_index, target_order);
317314 }
318315 }
319316
320317 const second_grandchild_index = first_grandchild_index + 1;
321318 if (second_grandchild_index >= self.len) {
322319 // One grandchild, so we know there is a second child. Compare first grandchild and second child
323 return self.bestItemAtIndices(first_grandchild_index, second_child_index, compare);
320 return self.bestItemAtIndices(first_grandchild_index, second_child_index, target_order);
324321 }
325322
326 const best_left_grandchild_index = self.bestItemAtIndices(first_grandchild_index, second_grandchild_index, compare).index;
323 const best_left_grandchild_index = self.bestItemAtIndices(first_grandchild_index, second_grandchild_index, target_order).index;
327324 const third_grandchild_index = second_grandchild_index + 1;
328325 if (third_grandchild_index >= self.len) {
329326 // Two grandchildren, and we know the best. Compare this to second child.
330 return self.bestItemAtIndices(best_left_grandchild_index, second_child_index, compare);
327 return self.bestItemAtIndices(best_left_grandchild_index, second_child_index, target_order);
331328 } else {
332329 // Three grandchildren, compare the min of the first two with the third
333 return self.bestItemAtIndices(best_left_grandchild_index, third_grandchild_index, compare);
330 return self.bestItemAtIndices(best_left_grandchild_index, third_grandchild_index, target_order);
334331 }
335332 }
336333
......@@ -348,12 +345,12 @@ pub fn PriorityDequeue(comptime T: type) type {
348345 /// Dequeue takes ownership of the passed in slice. The slice must have been
349346 /// allocated with `allocator`.
350347 /// De-initialize with `deinit`.
351 pub fn fromOwnedSlice(allocator: *Allocator, lessThanFn: fn (T, T) bool, items: []T) Self {
348 pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (T, T) Order, items: []T) Self {
352349 var queue = Self{
353350 .items = items,
354351 .len = items.len,
355352 .allocator = allocator,
356 .lessThanFn = lessThanFn,
353 .compareFn = compareFn,
357354 };
358355
359356 if (queue.len <= 1) return queue;
......@@ -468,8 +465,8 @@ pub fn PriorityDequeue(comptime T: type) type {
468465 };
469466}
470467
471fn lessThanComparison(a: u32, b: u32) bool {
472 return a < b;
468fn lessThanComparison(a: u32, b: u32) Order {
469 return std.math.order(a, b);
473470}
474471
475472const PDQ = PriorityDequeue(u32);
......@@ -493,6 +490,32 @@ test "std.PriorityDequeue: add and remove min" {
493490 expectEqual(@as(u32, 54), queue.removeMin());
494491}
495492
493test "std.PriorityDequeue: add and remove min structs" {
494 const S = struct {
495 size: u32,
496 };
497 var queue = PriorityDequeue(S).init(testing.allocator, struct {
498 fn order(a: S, b: S) Order {
499 return std.math.order(a.size, b.size);
500 }
501 }.order);
502 defer queue.deinit();
503
504 try queue.add(.{ .size = 54 });
505 try queue.add(.{ .size = 12 });
506 try queue.add(.{ .size = 7 });
507 try queue.add(.{ .size = 23 });
508 try queue.add(.{ .size = 25 });
509 try queue.add(.{ .size = 13 });
510
511 expectEqual(@as(u32, 7), queue.removeMin().size);
512 expectEqual(@as(u32, 12), queue.removeMin().size);
513 expectEqual(@as(u32, 13), queue.removeMin().size);
514 expectEqual(@as(u32, 23), queue.removeMin().size);
515 expectEqual(@as(u32, 25), queue.removeMin().size);
516 expectEqual(@as(u32, 54), queue.removeMin().size);
517}
518
496519test "std.PriorityDequeue: add and remove max" {
497520 var queue = PDQ.init(testing.allocator, lessThanComparison);
498521 defer queue.deinit();
lib/std/priority_queue.zig+23-20
......@@ -7,6 +7,7 @@ const std = @import("std.zig");
77const Allocator = std.mem.Allocator;
88const assert = std.debug.assert;
99const warn = std.debug.warn;
10const Order = std.math.Order;
1011const testing = std.testing;
1112const expect = testing.expect;
1213const expectEqual = testing.expectEqual;
......@@ -20,15 +21,17 @@ pub fn PriorityQueue(comptime T: type) type {
2021 items: []T,
2122 len: usize,
2223 allocator: *Allocator,
23 compareFn: fn (a: T, b: T) bool,
24
25 /// Initialize and return a priority queue. Provide
26 /// `compareFn` that returns `true` when its first argument
27 /// should get popped before its second argument. For example,
28 /// to make `pop` return the minimum value, provide
24 compareFn: fn (a: T, b: T) Order,
25
26 /// Initialize and return a priority queue. Provide `compareFn`
27 /// that returns `Order.lt` when its first argument should
28 /// get popped before its second argument, `Order.eq` if the
29 /// arguments are of equal priority, or `Order.gt` if the second
30 /// argument should be popped first. For example, to make `pop`
31 /// return the smallest number, provide
2932 ///
30 /// `fn lessThan(a: T, b: T) bool { return a < b; }`
31 pub fn init(allocator: *Allocator, compareFn: fn (a: T, b: T) bool) Self {
33 /// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`
34 pub fn init(allocator: *Allocator, compareFn: fn (a: T, b: T) Order) Self {
3235 return Self{
3336 .items = &[_]T{},
3437 .len = 0,
......@@ -61,7 +64,7 @@ pub fn PriorityQueue(comptime T: type) type {
6164 const child = self.items[child_index];
6265 const parent = self.items[parent_index];
6366
64 if (!self.compareFn(child, parent)) break;
67 if (self.compareFn(child, parent) != .lt) break;
6568
6669 self.items[parent_index] = child;
6770 self.items[child_index] = parent;
......@@ -133,14 +136,14 @@ pub fn PriorityQueue(comptime T: type) type {
133136 var smallest = self.items[index];
134137
135138 if (left) |e| {
136 if (self.compareFn(e, smallest)) {
139 if (self.compareFn(e, smallest) == .lt) {
137140 smallest_index = left_index;
138141 smallest = e;
139142 }
140143 }
141144
142145 if (right) |e| {
143 if (self.compareFn(e, smallest)) {
146 if (self.compareFn(e, smallest) == .lt) {
144147 smallest_index = right_index;
145148 smallest = e;
146149 }
......@@ -159,7 +162,7 @@ pub fn PriorityQueue(comptime T: type) type {
159162 /// PriorityQueue takes ownership of the passed in slice. The slice must have been
160163 /// allocated with `allocator`.
161164 /// Deinitialize with `deinit`.
162 pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (a: T, b: T) bool, items: []T) Self {
165 pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (a: T, b: T) Order, items: []T) Self {
163166 var queue = Self{
164167 .items = items,
165168 .len = items.len,
......@@ -217,10 +220,10 @@ pub fn PriorityQueue(comptime T: type) type {
217220 var update_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], elem) orelse return error.ElementNotFound;
218221 const old_elem: T = self.items[update_index];
219222 self.items[update_index] = new_elem;
220 if (self.compareFn(new_elem, old_elem)) {
221 siftUp(self, update_index);
222 } else {
223 siftDown(self, update_index);
223 switch (self.compareFn(new_elem, old_elem)) {
224 .lt => siftUp(self, update_index),
225 .gt => siftDown(self, update_index),
226 .eq => {}, // Nothing to do as the items have equal priority
224227 }
225228 }
226229
......@@ -267,12 +270,12 @@ pub fn PriorityQueue(comptime T: type) type {
267270 };
268271}
269272
270fn lessThan(a: u32, b: u32) bool {
271 return a < b;
273fn lessThan(a: u32, b: u32) Order {
274 return std.math.order(a, b);
272275}
273276
274fn greaterThan(a: u32, b: u32) bool {
275 return a > b;
277fn greaterThan(a: u32, b: u32) Order {
278 return lessThan(a, b).invert();
276279}
277280
278281const PQ = PriorityQueue(u32);