Fix a few bugs

Author: habedi

src/ordered/cartesian_tree_map.zig

@@ -110,15 +110,21 @@ pub fn CartesianTreeMap(
         /// Inserts a key-value pair with a random priority.
         ///
         /// Uses per-instance PRNG priorities to ensure expected O(log n) performance.
-        /// If the key already exists, updates its value and priority.
+        /// If the key already exists, updates its value in place; the node keeps
+        /// its current priority, so the tree shape (and balance) is unchanged.
         ///
         /// Time complexity: O(log n) expected
         ///
         /// ## Errors
         /// Returns `error.OutOfMemory` if node allocation fails.
         pub fn put(self: *Self, key: K, value: V) !void {
-            const priority = self.prng.random().int(u32);
-            try self.putWithPriority(key, value, priority);
+            // A value-only update keeps the existing priority: re-randomising it
+            // on every update would churn the tree structure for no benefit.
+            if (self.findNode(key)) |existing| {
+                existing.value = value;
+                return;
+            }
+            try self.insertNew(key, value, self.prng.random().int(u32));
         }
 
         /// Inserts a key-value pair with an explicit priority.
@@ -134,17 +140,20 @@ pub fn CartesianTreeMap(
         /// ## Errors
         /// Returns `error.OutOfMemory` if node allocation fails.
         pub fn putWithPriority(self: *Self, key: K, value: V, priority: u32) !void {
-            // Update in place if the key already exists. Resolving duplicates up
-            // front is required for correctness: `insertNode` only detects an
-            // equal key at the node it is currently visiting, so once the new
-            // priority exceeds an ancestor's priority it would split and insert
-            // a second node with the same key.
-            if (self.findNode(key)) |existing| {
-                existing.value = value;
-                existing.priority = priority;
-                return;
+            // To honor the explicit priority on an update, remove the existing
+            // node and reinsert it. Reinsertion positions the node by its new
+            // priority and preserves the max-heap invariant, whereas writing the
+            // priority in place would leave a node out of heap order relative to
+            // its parent or children.
+            if (self.contains(key)) {
+                _ = self.remove(key);
             }
+            try self.insertNew(key, value, priority);
+        }
 
+        /// Inserts a key known not to be present. Callers must ensure the key is
+        /// absent so the insertion never produces a duplicate.
+        fn insertNew(self: *Self, key: K, value: V, priority: u32) !void {
             const new_node = try self.allocator.create(Node);
             new_node.* = Node.init(key, value, priority);
 
@@ -403,9 +412,28 @@ fn i32Compare(lhs: i32, rhs: i32) std.math.Order {
 
 const MapOracle = @import("oracle.zig").MapOracle;
 
+const TreapI32 = CartesianTreeMap(i32, i32, i32Compare);
+
+/// Recursively asserts the two invariants every treap node must satisfy: BST
+/// ordering on keys and the max-heap property on priorities (a parent's
+/// priority is greater than or equal to each child's).
+fn expectTreapInvariants(node: ?*const TreapI32.Node) !void {
+    const n = node orelse return;
+    if (n.left) |l| {
+        try testing.expect(i32Compare(l.key, n.key) == .lt);
+        try testing.expect(l.priority <= n.priority);
+        try expectTreapInvariants(l);
+    }
+    if (n.right) |r| {
+        try testing.expect(i32Compare(r.key, n.key) == .gt);
+        try testing.expect(r.priority <= n.priority);
+        try expectTreapInvariants(r);
+    }
+}
+
 test "CartesianTreeMap: differential test against sorted-array oracle" {
     const allocator = testing.allocator;
-    var tree = CartesianTreeMap(i32, i32, i32Compare).init(allocator);
+    var tree = TreapI32.init(allocator);
     defer tree.deinit();
 
     var oracle: MapOracle(i32, i32, i32Compare) = .{};
@@ -456,10 +484,53 @@ test "CartesianTreeMap: differential test against sorted-array oracle" {
             while (k < @as(i32, @intCast(key_space))) : (k += 1) {
                 try testing.expectEqual(oracle.contains(k), tree.contains(k));
             }
+            // Every put, remove, and value-update must leave the treap a valid
+            // max-heap on priorities and a valid BST on keys.
+            try expectTreapInvariants(tree.root);
         }
     }
 }
 
+test "CartesianTreeMap: putWithPriority update preserves heap invariant" {
+    var tree = TreapI32.init(testing.allocator);
+    defer tree.deinit();
+
+    try tree.putWithPriority(5, 50, 10);
+    try tree.putWithPriority(3, 30, 20);
+    try tree.putWithPriority(8, 80, 5);
+    try tree.putWithPriority(1, 10, 15);
+    try tree.putWithPriority(7, 70, 25);
+    try expectTreapInvariants(tree.root);
+
+    // Raising key 5's priority above all others must lift it to the root while
+    // keeping the structure a valid treap.
+    try tree.putWithPriority(5, 55, 100);
+    try testing.expectEqual(@as(usize, 5), tree.count());
+    try testing.expectEqual(@as(i32, 55), tree.get(5).?);
+    try testing.expectEqual(@as(i32, 5), tree.root.?.key);
+    try testing.expectEqual(@as(u32, 100), tree.root.?.priority);
+    try expectTreapInvariants(tree.root);
+
+    // Lowering it again must re-sink it without breaking the invariants, and the
+    // value update must persist.
+    try tree.putWithPriority(5, 555, 1);
+    try testing.expectEqual(@as(usize, 5), tree.count());
+    try testing.expectEqual(@as(i32, 555), tree.get(5).?);
+    try expectTreapInvariants(tree.root);
+}
+
+test "CartesianTreeMap: put updates value without changing priority" {
+    var tree = TreapI32.init(testing.allocator);
+    defer tree.deinit();
+
+    try tree.putWithPriority(10, 100, 42);
+    // A value-only `put` must leave the node's priority untouched.
+    try tree.put(10, 200);
+    try testing.expectEqual(@as(usize, 1), tree.count());
+    try testing.expectEqual(@as(i32, 200), tree.get(10).?);
+    try testing.expectEqual(@as(u32, 42), tree.root.?.priority);
+}
+
 test "CartesianTreeMap basic operations" {
     var tree = CartesianTreeMap(i32, []const u8, i32Compare).init(testing.allocator);
     defer tree.deinit();

src/ordered/trie_map.zig

@@ -267,7 +267,7 @@ pub fn TrieMap(comptime V: type) type {
             errdefer stack.deinit(allocator);
             try stack.append(allocator, PrefixIteratorFrame{
                 .node = prefix_node.?,
-                .child_iter = prefix_node.?.children.iterator(),
+                .next_char = 0,
                 .visited_self = false,
             });
 
@@ -286,7 +286,9 @@ pub fn TrieMap(comptime V: type) type {
 
         pub const PrefixIteratorFrame = struct {
             node: *const TrieNode,
-            child_iter: std.HashMap(u8, *TrieNode, std.hash_map.AutoContext(u8), std.hash_map.default_max_load_percentage).Iterator,
+            // Next child byte to examine, scanned over 0..256 so children are
+            // visited in ascending byte order, yielding keys in sorted order.
+            next_char: u16,
             visited_self: bool,
         };
 
@@ -303,22 +305,34 @@ pub fn TrieMap(comptime V: type) type {
 
             pub fn next(self: *PrefixIterator) !?[]const u8 {
                 while (self.stack.items.len > 0) {
-                    var frame = &self.stack.items[self.stack.items.len - 1];
+                    const top = self.stack.items.len - 1;
 
-                    if (!frame.visited_self and frame.node.is_end) {
-                        frame.visited_self = true;
+                    if (!self.stack.items[top].visited_self and self.stack.items[top].node.is_end) {
+                        self.stack.items[top].visited_self = true;
                         return self.current_key.items;
                     }
 
-                    if (frame.child_iter.next()) |entry| {
-                        const char = entry.key_ptr.*;
-                        const child = entry.value_ptr.*;
-
-                        try self.current_key.append(self.allocator, char);
+                    // Scan ascending byte values for the next existing child.
+                    const node = self.stack.items[top].node;
+                    var next_child: ?*TrieNode = null;
+                    var next_byte: u8 = 0;
+                    while (self.stack.items[top].next_char < 256) {
+                        const char: u8 = @intCast(self.stack.items[top].next_char);
+                        self.stack.items[top].next_char += 1;
+                        if (node.children.get(char)) |child| {
+                            next_child = child;
+                            next_byte = char;
+                            break;
+                        }
+                    }
 
+                    if (next_child) |child| {
+                        try self.current_key.append(self.allocator, next_byte);
+                        // This append may reallocate `stack`, so `top` is
+                        // recomputed on the next loop iteration rather than reused.
                         try self.stack.append(self.allocator, PrefixIteratorFrame{
                             .node = child,
-                            .child_iter = child.children.iterator(),
+                            .next_char = 0,
                             .visited_self = false,
                         });
                     } else {
@@ -693,3 +707,25 @@ test "TrieMap: special characters" {
     try std.testing.expectEqual(@as(i32, 2), trie.get("test_case").?.*);
     try std.testing.expectEqual(@as(i32, 3), trie.get("foo.bar").?.*);
 }
+
+test "TrieMap: keysWithPrefix yields keys in sorted order" {
+    const allocator = std.testing.allocator;
+    var trie = try TrieMap(i32).init(allocator);
+    defer trie.deinit();
+
+    // Insert in an order that is neither sorted nor hash order.
+    const keys = [_][]const u8{ "bandit", "ban", "bandana", "band", "banana", "bee", "apex" };
+    for (keys, 0..) |k, i| try trie.put(k, @intCast(i));
+
+    var iter = try trie.keysWithPrefix(allocator, "ban");
+    defer iter.deinit();
+
+    // Keys under the "ban" prefix must come out in lexicographic order.
+    const expected = [_][]const u8{ "ban", "banana", "band", "bandana", "bandit" };
+    var idx: usize = 0;
+    while (try iter.next()) |k| : (idx += 1) {
+        try std.testing.expect(idx < expected.len);
+        try std.testing.expectEqualStrings(expected[idx], k);
+    }
+    try std.testing.expectEqual(expected.len, idx);
+}