From f36fa7f122f7ccbf629825ced8236ba2aa1e367f Mon Sep 17 00:00:00 2001
From: DualFroz <me@dualfroz.com>
Date: Fri, 4 Sep 2026 14:25:50 +0200
Subject: [PATCH] feat: add descU32 and Weighted to heap

---
 src/heap.zig | 109 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 109 insertions(+)

diff --git a/src/heap.zig b/src/heap.zig
index 05aed9e..8a02873 100644
--- a/src/heap.zig
+++ b/src/heap.zig
@@ -154,3 +154,112 @@ test "heap: replaceTop matches pop-then-push" {
     }
 }
 
+test "heap: random stress matches sort reference and stays valid" {
+    var prng = std.Random.DefaultPrng.init(0xC0BA17);
+    const random = prng.random();
+    var heap = U32Heap.init(std.testing.allocator, {});
+    defer heap.deinit();
+
+    var reference = std.ArrayList(u32).init(std.testing.allocator);
+    defer reference.deinit();
+
+    var i: usize = 0;
+    while (i < 500) : (i += 1) {
+        const v = random.intRangeAtMost(u32, 0, 1000);
+        try heap.push(v);
+        try reference.append(v);
+        try std.testing.expect(heap.isValid());
+    }
+
+    std.mem.sort(u32, reference.items, {}, std.sort.asc(u32));
+    for (reference.items) |expected| {
+        try std.testing.expectEqual(@as(?u32, expected), heap.pop());
+        try std.testing.expect(heap.isValid());
+    }
+    try std.testing.expectEqual(@as(usize, 0), heap.len());
+}
+
+test "heap: interleaved push and pop stays valid" {
+    var prng = std.Random.DefaultPrng.init(0xFEED);
+    const random = prng.random();
+    var heap = U32Heap.init(std.testing.allocator, {});
+    defer heap.deinit();
+
+    var count: usize = 0;
+    var i: usize = 0;
+    while (i < 2000) : (i += 1) {
+        if (count == 0 or random.boolean()) {
+            try heap.push(random.intRangeAtMost(u32, 0, 100));
+            count += 1;
+        } else {
+            const popped = heap.pop();
+            try std.testing.expect(popped != null);
+            count -= 1;
+        }
+        try std.testing.expect(heap.isValid());
+        try std.testing.expectEqual(count, heap.len());
+    }
+}
+
+fn descU32(_: void, a: u32, b: u32) bool {
+    return a > b;
+}
+const MaxU32Heap = MinHeap(u32, void, descU32);
+
+test "heap: inverted comparator behaves as a max-heap" {
+    var heap = MaxU32Heap.init(std.testing.allocator, {});
+    defer heap.deinit();
+    for ([_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 }) |v| try heap.push(v);
+
+    var sorted = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 };
+    std.mem.sort(u32, &sorted, {}, std.sort.desc(u32));
+    for (sorted) |expected| try std.testing.expectEqual(@as(?u32, expected), heap.pop());
+}
+
+const Weighted = struct { key: []const u8, seq: u32 };
+
+fn keyLess(_: void, a: Weighted, b: Weighted) bool {
+    return std.mem.order(u8, a.key, b.key) == .lt;
+}
+const KeyHeap = MinHeap(Weighted, void, keyLess);
+
+test "heap: k-way merge of pre-sorted lists matches full sort" {
+    const allocator = std.testing.allocator;
+    const lists = [_][]const []const u8{
+        &[_][]const u8{ "apple", "cherry", "fig", "kiwi" },
+        &[_][]const u8{ "banana", "date", "eggplant" },
+        &[_][]const u8{"lemon"},
+        &[_][]const u8{},
+    };
+
+    const Cursor = struct { list: []const []const u8, idx: usize };
+    var cursors: [lists.len]Cursor = undefined;
+    for (lists, 0..) |list, i| cursors[i] = .{ .list = list, .idx = 0 };
+
+    var heap = KeyHeap.init(allocator, {});
+    defer heap.deinit();
+    for (cursors, 0..) |c, i| {
+        if (c.idx < c.list.len) try heap.push(.{ .key = c.list[c.idx], .seq = @intCast(i) });
+    }
+
+    var merged = std.ArrayList([]const u8).init(allocator);
+    defer merged.deinit();
+    while (heap.pop()) |top| {
+        try merged.append(top.key);
+        var c = &cursors[top.seq];
+        c.idx += 1;
+        if (c.idx < c.list.len) try heap.push(.{ .key = c.list[c.idx], .seq = top.seq });
+    }
+
+    var expected = std.ArrayList([]const u8).init(allocator);
+    defer expected.deinit();
+    for (lists) |list| try expected.appendSlice(list);
+    std.mem.sort([]const u8, expected.items, {}, struct {
+        fn lt(_: void, a: []const u8, b: []const u8) bool {
+            return std.mem.order(u8, a, b) == .lt;
+        }
+    }.lt);
+
+    try std.testing.expectEqual(expected.items.len, merged.items.len);
+    for (expected.items, merged.items) |e, m| try std.testing.expectEqualStrings(e, m);
+}