| ... | @@ -7,6 +7,7 @@ const std = @import("std.zig"); | ... | @@ -7,6 +7,7 @@ const std = @import("std.zig"); |
| 7 | const Allocator = std.mem.Allocator; | 7 | const Allocator = std.mem.Allocator; |
| 8 | const assert = std.debug.assert; | 8 | const assert = std.debug.assert; |
| 9 | const warn = std.debug.warn; | 9 | const warn = std.debug.warn; |
| | 10 | const Order = std.math.Order; |
| 10 | const testing = std.testing; | 11 | const testing = std.testing; |
| 11 | const expect = testing.expect; | 12 | const expect = testing.expect; |
| 12 | const expectEqual = testing.expectEqual; | 13 | const expectEqual = testing.expectEqual; |
| ... | @@ -20,31 +21,26 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -20,31 +21,26 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 20 | items: []T, | 21 | items: []T, |
| 21 | len: usize, | 22 | len: usize, |
| 22 | allocator: *Allocator, | 23 | allocator: *Allocator, |
| 23 | lessThanFn: fn (a: T, b: T) bool, | 24 | compareFn: fn (a: T, b: T) Order, |
| 24 | | 25 | |
| 25 | /// Initialize and return a new dequeue. Provide `lessThanFn` | 26 | /// Initialize and return a new priority dequeue. Provide `compareFn` |
| 26 | /// that returns `true` when its first argument should | 27 | /// that returns `Order.lt` when its first argument should |
| 27 | /// get min-popped before its second argument. For example, | 28 | /// get min-popped before its second argument, `Order.eq` if the |
| 28 | /// to make `popMin` return the minimum value, provide | 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 |
| 29 | /// | 33 | /// |
| 30 | /// `fn lessThanFn(a: T, b: T) bool { return a < b; }` | 34 | /// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }` |
| 31 | pub fn init(allocator: *Allocator, lessThanFn: fn (T, T) bool) Self { | 35 | pub fn init(allocator: *Allocator, compareFn: fn (T, T) Order) Self { |
| 32 | return Self{ | 36 | return Self{ |
| 33 | .items = &[_]T{}, | 37 | .items = &[_]T{}, |
| 34 | .len = 0, | 38 | .len = 0, |
| 35 | .allocator = allocator, | 39 | .allocator = allocator, |
| 36 | .lessThanFn = lessThanFn, | 40 | .compareFn = compareFn, |
| 37 | }; | 41 | }; |
| 38 | } | 42 | } |
| 39 | | 43 | |
| 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 | | | |
| 48 | /// Free memory used by the dequeue. | 44 | /// Free memory used by the dequeue. |
| 49 | pub fn deinit(self: Self) void { | 45 | pub fn deinit(self: Self) void { |
| 50 | self.allocator.free(self.items); | 46 | self.allocator.free(self.items); |
| ... | @@ -100,7 +96,8 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -100,7 +96,8 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 100 | const parent = self.items[parent_index]; | 96 | const parent = self.items[parent_index]; |
| 101 | | 97 | |
| 102 | const min_layer = self.nextIsMinLayer(); | 98 | 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)) { |
| 104 | // We must swap the item with it's parent if it is on the "wrong" layer | 101 | // We must swap the item with it's parent if it is on the "wrong" layer |
| 105 | self.items[parent_index] = child; | 102 | self.items[parent_index] = child; |
| 106 | self.items[child_index] = parent; | 103 | self.items[child_index] = parent; |
| ... | @@ -118,21 +115,21 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -118,21 +115,21 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 118 | | 115 | |
| 119 | fn siftUp(self: *Self, start: StartIndexAndLayer) void { | 116 | fn siftUp(self: *Self, start: StartIndexAndLayer) void { |
| 120 | if (start.min_layer) { | 117 | if (start.min_layer) { |
| 121 | doSiftUp(self, start.index, lessThan); | 118 | doSiftUp(self, start.index, .lt); |
| 122 | } else { | 119 | } else { |
| 123 | doSiftUp(self, start.index, greaterThan); | 120 | doSiftUp(self, start.index, .gt); |
| 124 | } | 121 | } |
| 125 | } | 122 | } |
| 126 | | 123 | |
| 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 { |
| 128 | var child_index = start_index; | 125 | var child_index = start_index; |
| 129 | while (child_index > 2) { | 126 | while (child_index > 2) { |
| 130 | var grandparent_index = grandparentIndex(child_index); | 127 | var grandparent_index = grandparentIndex(child_index); |
| 131 | const child = self.items[child_index]; | 128 | const child = self.items[child_index]; |
| 132 | const grandparent = self.items[grandparent_index]; | 129 | const grandparent = self.items[grandparent_index]; |
| 133 | | 130 | |
| 134 | // If the grandparent is already better, we have gone as far as we need to | 131 | // If the grandparent is already better or equal, we have gone as far as we need to |
| 135 | if (!compare(self.*, child, grandparent)) break; | 132 | if (self.compareFn(child, grandparent) != target_order) break; |
| 136 | | 133 | |
| 137 | // Otherwise swap the item with it's grandparent | 134 | // Otherwise swap the item with it's grandparent |
| 138 | self.items[grandparent_index] = child; | 135 | self.items[grandparent_index] = child; |
| ... | @@ -153,14 +150,14 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -153,14 +150,14 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 153 | if (self.len == 0) return null; | 150 | if (self.len == 0) return null; |
| 154 | if (self.len == 1) return self.items[0]; | 151 | if (self.len == 1) return self.items[0]; |
| 155 | if (self.len == 2) return self.items[1]; | 152 | if (self.len == 2) return self.items[1]; |
| 156 | return self.bestItemAtIndices(1, 2, greaterThan).item; | 153 | return self.bestItemAtIndices(1, 2, .gt).item; |
| 157 | } | 154 | } |
| 158 | | 155 | |
| 159 | fn maxIndex(self: Self) ?usize { | 156 | fn maxIndex(self: Self) ?usize { |
| 160 | if (self.len == 0) return null; | 157 | if (self.len == 0) return null; |
| 161 | if (self.len == 1) return 0; | 158 | if (self.len == 1) return 0; |
| 162 | if (self.len == 2) return 1; | 159 | if (self.len == 2) return 1; |
| 163 | return self.bestItemAtIndices(1, 2, greaterThan).index; | 160 | return self.bestItemAtIndices(1, 2, .gt).index; |
| 164 | } | 161 | } |
| 165 | | 162 | |
| 166 | /// Pop the smallest element from the dequeue. Returns | 163 | /// Pop the smallest element from the dequeue. Returns |
| ... | @@ -204,13 +201,13 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -204,13 +201,13 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 204 | | 201 | |
| 205 | fn siftDown(self: *Self, index: usize) void { | 202 | fn siftDown(self: *Self, index: usize) void { |
| 206 | if (isMinLayer(index)) { | 203 | if (isMinLayer(index)) { |
| 207 | self.doSiftDown(index, lessThan); | 204 | self.doSiftDown(index, .lt); |
| 208 | } else { | 205 | } else { |
| 209 | self.doSiftDown(index, greaterThan); | 206 | self.doSiftDown(index, .gt); |
| 210 | } | 207 | } |
| 211 | } | 208 | } |
| 212 | | 209 | |
| 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 { |
| 214 | var index = start_index; | 211 | var index = start_index; |
| 215 | const half = self.len >> 1; | 212 | const half = self.len >> 1; |
| 216 | while (true) { | 213 | while (true) { |
| ... | @@ -225,12 +222,12 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -225,12 +222,12 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 225 | const index3 = index2 + 1; | 222 | const index3 = index2 + 1; |
| 226 | | 223 | |
| 227 | // Find the best grandchild | 224 | // Find the best grandchild |
| 228 | const best_left = self.bestItemAtIndices(first_grandchild_index, index2, compare); | 225 | const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order); |
| 229 | const best_right = self.bestItemAtIndices(index3, last_grandchild_index, compare); | 226 | const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order); |
| 230 | const best_grandchild = self.bestItem(best_left, best_right, compare); | 227 | const best_grandchild = self.bestItem(best_left, best_right, target_order); |
| 231 | | 228 | |
| 232 | // If the item is better than it's best grandchild, we are done | 229 | // If the item is better than or equal to its best grandchild, we are done |
| 233 | if (compare(self.*, elem, best_grandchild.item) or elem == best_grandchild.item) return; | 230 | if (self.compareFn(best_grandchild.item, elem) != target_order) return; |
| 234 | | 231 | |
| 235 | // Otherwise, swap them | 232 | // Otherwise, swap them |
| 236 | self.items[best_grandchild.index] = elem; | 233 | self.items[best_grandchild.index] = elem; |
| ... | @@ -238,16 +235,16 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -238,16 +235,16 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 238 | index = best_grandchild.index; | 235 | index = best_grandchild.index; |
| 239 | | 236 | |
| 240 | // We might need to swap the element with it's parent | 237 | // We might need to swap the element with it's parent |
| 241 | self.swapIfParentIsBetter(elem, index, compare); | 238 | self.swapIfParentIsBetter(elem, index, target_order); |
| 242 | } else { | 239 | } else { |
| 243 | // The children or grandchildren are the last layer | 240 | // The children or grandchildren are the last layer |
| 244 | const first_child_index = firstChildIndex(index); | 241 | const first_child_index = firstChildIndex(index); |
| 245 | if (first_child_index > self.len) return; | 242 | if (first_child_index > self.len) return; |
| 246 | | 243 | |
| 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); |
| 248 | | 245 | |
| 249 | // If the best descendant is still larger, we are done | 246 | // If the item is better than or equal to its best descendant, we are done |
| 250 | if (compare(self.*, elem, best_descendent.item) or elem == best_descendent.item) return; | 247 | if (self.compareFn(best_descendent.item, elem) != target_order) return; |
| 251 | | 248 | |
| 252 | // Otherwise swap them | 249 | // Otherwise swap them |
| 253 | self.items[best_descendent.index] = elem; | 250 | self.items[best_descendent.index] = elem; |
| ... | @@ -258,7 +255,7 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -258,7 +255,7 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 258 | if (index < first_grandchild_index) return; | 255 | if (index < first_grandchild_index) return; |
| 259 | | 256 | |
| 260 | // We might need to swap the element with it's parent | 257 | // We might need to swap the element with it's parent |
| 261 | self.swapIfParentIsBetter(elem, index, compare); | 258 | self.swapIfParentIsBetter(elem, index, target_order); |
| 262 | return; | 259 | return; |
| 263 | } | 260 | } |
| 264 | | 261 | |
| ... | @@ -267,11 +264,11 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -267,11 +264,11 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 267 | } | 264 | } |
| 268 | } | 265 | } |
| 269 | | 266 | |
| 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 { |
| 271 | const parent_index = parentIndex(child_index); | 268 | const parent_index = parentIndex(child_index); |
| 272 | const parent = self.items[parent_index]; | 269 | const parent = self.items[parent_index]; |
| 273 | | 270 | |
| 274 | if (compare(self.*, parent, child)) { | 271 | if (self.compareFn(parent, child) == target_order) { |
| 275 | self.items[parent_index] = child; | 272 | self.items[parent_index] = child; |
| 276 | self.items[child_index] = parent; | 273 | self.items[child_index] = parent; |
| 277 | } | 274 | } |
| ... | @@ -289,21 +286,21 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -289,21 +286,21 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 289 | }; | 286 | }; |
| 290 | } | 287 | } |
| 291 | | 288 | |
| 292 | fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, compare: fn (Self, T, T) bool) ItemAndIndex { | 289 | fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex { |
| 293 | if (compare(self, item1.item, item2.item)) { | 290 | if (self.compareFn(item1.item, item2.item) == target_order) { |
| 294 | return item1; | 291 | return item1; |
| 295 | } else { | 292 | } else { |
| 296 | return item2; | 293 | return item2; |
| 297 | } | 294 | } |
| 298 | } | 295 | } |
| 299 | | 296 | |
| 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 { |
| 301 | var item1 = self.getItem(index1); | 298 | var item1 = self.getItem(index1); |
| 302 | var item2 = self.getItem(index2); | 299 | var item2 = self.getItem(index2); |
| 303 | return self.bestItem(item1, item2, compare); | 300 | return self.bestItem(item1, item2, target_order); |
| 304 | } | 301 | } |
| 305 | | 302 | |
| 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 { |
| 307 | const second_child_index = first_child_index + 1; | 304 | const second_child_index = first_child_index + 1; |
| 308 | if (first_grandchild_index >= self.len) { | 305 | if (first_grandchild_index >= self.len) { |
| 309 | // No grandchildren, find the best child (second may not exist) | 306 | // No grandchildren, find the best child (second may not exist) |
| ... | @@ -313,24 +310,24 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -313,24 +310,24 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 313 | .index = first_child_index, | 310 | .index = first_child_index, |
| 314 | }; | 311 | }; |
| 315 | } else { | 312 | } else { |
| 316 | return self.bestItemAtIndices(first_child_index, second_child_index, compare); | 313 | return self.bestItemAtIndices(first_child_index, second_child_index, target_order); |
| 317 | } | 314 | } |
| 318 | } | 315 | } |
| 319 | | 316 | |
| 320 | const second_grandchild_index = first_grandchild_index + 1; | 317 | const second_grandchild_index = first_grandchild_index + 1; |
| 321 | if (second_grandchild_index >= self.len) { | 318 | if (second_grandchild_index >= self.len) { |
| 322 | // One grandchild, so we know there is a second child. Compare first grandchild and second child | 319 | // 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); |
| 324 | } | 321 | } |
| 325 | | 322 | |
| 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; |
| 327 | const third_grandchild_index = second_grandchild_index + 1; | 324 | const third_grandchild_index = second_grandchild_index + 1; |
| 328 | if (third_grandchild_index >= self.len) { | 325 | if (third_grandchild_index >= self.len) { |
| 329 | // Two grandchildren, and we know the best. Compare this to second child. | 326 | // 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); |
| 331 | } else { | 328 | } else { |
| 332 | // Three grandchildren, compare the min of the first two with the third | 329 | // 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); |
| 334 | } | 331 | } |
| 335 | } | 332 | } |
| 336 | | 333 | |
| ... | @@ -348,12 +345,12 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -348,12 +345,12 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 348 | /// Dequeue takes ownership of the passed in slice. The slice must have been | 345 | /// Dequeue takes ownership of the passed in slice. The slice must have been |
| 349 | /// allocated with `allocator`. | 346 | /// allocated with `allocator`. |
| 350 | /// De-initialize with `deinit`. | 347 | /// 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 { |
| 352 | var queue = Self{ | 349 | var queue = Self{ |
| 353 | .items = items, | 350 | .items = items, |
| 354 | .len = items.len, | 351 | .len = items.len, |
| 355 | .allocator = allocator, | 352 | .allocator = allocator, |
| 356 | .lessThanFn = lessThanFn, | 353 | .compareFn = compareFn, |
| 357 | }; | 354 | }; |
| 358 | | 355 | |
| 359 | if (queue.len <= 1) return queue; | 356 | if (queue.len <= 1) return queue; |
| ... | @@ -468,8 +465,8 @@ pub fn PriorityDequeue(comptime T: type) type { | ... | @@ -468,8 +465,8 @@ pub fn PriorityDequeue(comptime T: type) type { |
| 468 | }; | 465 | }; |
| 469 | } | 466 | } |
| 470 | | 467 | |
| 471 | fn lessThanComparison(a: u32, b: u32) bool { | 468 | fn lessThanComparison(a: u32, b: u32) Order { |
| 472 | return a < b; | 469 | return std.math.order(a, b); |
| 473 | } | 470 | } |
| 474 | | 471 | |
| 475 | const PDQ = PriorityDequeue(u32); | 472 | const PDQ = PriorityDequeue(u32); |
| ... | @@ -493,6 +490,32 @@ test "std.PriorityDequeue: add and remove min" { | ... | @@ -493,6 +490,32 @@ test "std.PriorityDequeue: add and remove min" { |
| 493 | expectEqual(@as(u32, 54), queue.removeMin()); | 490 | expectEqual(@as(u32, 54), queue.removeMin()); |
| 494 | } | 491 | } |
| 495 | | 492 | |
| | 493 | test "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 | |
| 496 | test "std.PriorityDequeue: add and remove max" { | 519 | test "std.PriorityDequeue: add and remove max" { |
| 497 | var queue = PDQ.init(testing.allocator, lessThanComparison); | 520 | var queue = PDQ.init(testing.allocator, lessThanComparison); |
| 498 | defer queue.deinit(); | 521 | defer queue.deinit(); |