From 4d0303b99249adcc4acdd51fa9a10884f1e5ed61 Mon Sep 17 00:00:00 2001 From: kprotty Date: Fri, 15 Apr 2022 16:30:45 -0500 Subject: [PATCH] treap: fix + determinstically randomize test cases. --- lib/std/treap.zig | 152 ++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 128 insertions(+), 24 deletions(-) diff --git a/lib/std/treap.zig b/lib/std/treap.zig index 27c5285c34..7d08ef64b6 100644 --- a/lib/std/treap.zig +++ b/lib/std/treap.zig @@ -103,9 +103,13 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { /// An Entry represents a slot in the treap associated with a given key. pub const Entry = struct { + /// The associated key for this entry. key: Key, + /// A reference to the treap this entry is apart of. treap: *Self, + /// The current node at this entry. node: ?*Node, + /// The current state of the entry. context: union(enum) { /// A find() was called for this entry and the position in the treap is known. inserted_under: ?*Node, @@ -113,11 +117,6 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { removed, }, - /// Returns the current Node at this Entry in the treap if there is one. - pub fn get(self: Entry) ?*Node { - return self.node; - } - /// Update's the Node at this Entry in the treap with the new node. pub fn set(self: *Entry, new_node: ?*Node) void { // Update the entry's node reference after updating the treap below. @@ -182,7 +181,7 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { while (node.parent) |p| { if (p.priority <= node.priority) break; - const is_right = p.children[1] == @as(?*Node, node); + const is_right = p.children[1] == node; assert(p.children[@boolToInt(is_right)] == node); const rotate_right = !is_right; @@ -214,8 +213,8 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { // rotate the node down to be a leaf of the tree for removal, respecting priorities. while (node.children[0] orelse node.children[1]) |_| { self.rotate(node, rotate_right: { - const right = node.children[0] orelse break :rotate_right true; - const left = node.children[1] orelse break :rotate_right false; + const right = node.children[1] orelse break :rotate_right true; + const left = node.children[0] orelse break :rotate_right false; break :rotate_right (left.priority < right.priority); }); } @@ -244,10 +243,13 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { const target = node.children[@boolToInt(!right)] orelse unreachable; const adjacent = target.children[@boolToInt(right)]; - // do the rotation + // rotate the children target.children[@boolToInt(right)] = node; - node.parent = target; node.children[@boolToInt(!right)] = adjacent; + + // rotate the parents + node.parent = target; + target.parent = parent; if (adjacent) |adj| adj.parent = node; // fix the parent link @@ -258,37 +260,139 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { }; } +// For iterating a slice in a random order +// https://lemire.me/blog/2017/09/18/visiting-all-values-in-an-array-exactly-once-in-random-order/ +fn SliceIterRandomOrder(comptime T: type) type { + return struct { + rng: std.rand.Random, + slice: []T, + index: usize = undefined, + offset: usize = undefined, + co_prime: usize, + + const Self = @This(); + + pub fn init(slice: []T, rng: std.rand.Random) Self { + return Self{ + .rng = rng, + .slice = slice, + .co_prime = blk: { + if (slice.len == 0) break :blk 0; + var prime = slice.len / 2; + while (prime < slice.len) : (prime += 1) { + var gcd = [_]usize{ prime, slice.len }; + while (gcd[1] != 0) { + const temp = gcd; + gcd = [_]usize{ temp[1], temp[0] % temp[1] }; + } + if (gcd[0] == 1) break; + } + break :blk prime; + }, + }; + } + + pub fn reset(self: *Self) void { + self.index = 0; + self.offset = self.rng.int(usize); + } + + pub fn next(self: *Self) ?*T { + if (self.index >= self.slice.len) return null; + defer self.index += 1; + return &self.slice[((self.index *% self.co_prime) +% self.offset) % self.slice.len]; + } + }; +} + const TestTreap = Treap(u64, std.math.order); const TestNode = TestTreap.Node; -test "std.Treap: insert, find, remove" { - var prng = std.rand.DefaultPrng.init(0xdeadbeef); - var rng = prng.random(); - +test "std.Treap: insert, find, replace, remove" { var treap = TestTreap{}; - var nodes: [6]TestNode = undefined; + var nodes: [10]TestNode = undefined; - for (nodes) |*node| { - const key = rng.int(u64); + var prng = std.rand.DefaultPrng.init(0xdeadbeef); + var iter = SliceIterRandomOrder(TestNode).init(&nodes, prng.random()); + // insert check + iter.reset(); + while (iter.next()) |node| { + const key = prng.random().int(u64); + + // make sure the current entry is empty. var entry = treap.getEntryFor(key); try testing.expectEqual(entry.key, key); - try testing.expectEqual(entry.get(), null); + try testing.expectEqual(entry.node, null); + // insert the entry and make sure the fields are correct. entry.set(node); - try testing.expectEqual(entry.key, key); try testing.expectEqual(node.key, key); - try testing.expectEqual(entry.get(), node); + try testing.expectEqual(entry.key, key); + try testing.expectEqual(entry.node, node); } - for (nodes) |*node| { + // find check + iter.reset(); + while (iter.next()) |node| { const key = node.key; + // find the entry by-key and by-node after having been inserted. var entry = treap.getEntryFor(node.key); try testing.expectEqual(entry.key, key); - try testing.expectEqual(entry.get(), node); + try testing.expectEqual(entry.node, node); + try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node); + } - var existingEntry = treap.getEntryForExisting(node); - try testing.expectEqual(entry, existingEntry); + // replace check + iter.reset(); + while (iter.next()) |node| { + const key = node.key; + + // find the entry by node since we already know it exists + var entry = treap.getEntryForExisting(node); + try testing.expectEqual(entry.key, key); + try testing.expectEqual(entry.node, node); + + var stub_node: TestNode = undefined; + + // replace the node with a stub_node and ensure future finds point to the stub_node. + entry.set(&stub_node); + try testing.expectEqual(entry.node, &stub_node); + try testing.expectEqual(entry.node, treap.getEntryFor(key).node); + try testing.expectEqual(entry.node, treap.getEntryForExisting(&stub_node).node); + + // replace the stub_node back to the node and ensure future finds point to the old node. + entry.set(node); + try testing.expectEqual(entry.node, node); + try testing.expectEqual(entry.node, treap.getEntryFor(key).node); + try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node); + } + + // remove check + iter.reset(); + while (iter.next()) |node| { + const key = node.key; + + // find the entry by node since we already know it exists + var entry = treap.getEntryForExisting(node); + try testing.expectEqual(entry.key, key); + try testing.expectEqual(entry.node, node); + + // remove the node at the entry and ensure future finds point to it being removed. + entry.set(null); + try testing.expectEqual(entry.node, null); + try testing.expectEqual(entry.node, treap.getEntryFor(key).node); + + // insert the node back and ensure future finds point to the inserted node + entry.set(node); + try testing.expectEqual(entry.node, node); + try testing.expectEqual(entry.node, treap.getEntryFor(key).node); + try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node); + + // remove the node again and make sure it was cleared after the insert + entry.set(null); + try testing.expectEqual(entry.node, null); + try testing.expectEqual(entry.node, treap.getEntryFor(key).node); } }