From 7284eb22dce8086674c449ccf434ecdced9e4e0d Mon Sep 17 00:00:00 2001 From: kprotty Date: Fri, 15 Apr 2022 15:26:51 -0500 Subject: [PATCH 1/3] treap: initial implementation --- CMakeLists.txt | 1 + lib/std/std.zig | 1 + lib/std/treap.zig | 294 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 296 insertions(+) create mode 100644 lib/std/treap.zig diff --git a/CMakeLists.txt b/CMakeLists.txt index 81b96d67ab..9c4982ea83 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -539,6 +539,7 @@ set(ZIG_STAGE2_SOURCES "${CMAKE_SOURCE_DIR}/lib/std/Thread/ResetEvent.zig" "${CMAKE_SOURCE_DIR}/lib/std/Thread/StaticResetEvent.zig" "${CMAKE_SOURCE_DIR}/lib/std/time.zig" + "${CMAKE_SOURCE_DIR}/lib/std/treap.zig" "${CMAKE_SOURCE_DIR}/lib/std/unicode.zig" "${CMAKE_SOURCE_DIR}/lib/std/zig.zig" "${CMAKE_SOURCE_DIR}/lib/std/zig/Ast.zig" diff --git a/lib/std/std.zig b/lib/std/std.zig index f34d108271..831a82cbbe 100644 --- a/lib/std/std.zig +++ b/lib/std/std.zig @@ -39,6 +39,7 @@ pub const StringArrayHashMapUnmanaged = array_hash_map.StringArrayHashMapUnmanag pub const TailQueue = @import("linked_list.zig").TailQueue; pub const Target = @import("target.zig").Target; pub const Thread = @import("Thread.zig"); +pub const Treap = @import("treap.zig").Treap; pub const Tz = @import("tz.zig").Tz; pub const array_hash_map = @import("array_hash_map.zig"); diff --git a/lib/std/treap.zig b/lib/std/treap.zig new file mode 100644 index 0000000000..27c5285c34 --- /dev/null +++ b/lib/std/treap.zig @@ -0,0 +1,294 @@ +const std = @import("std.zig"); +const assert = std.debug.assert; +const testing = std.testing; +const Order = std.math.Order; + +pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { + return struct { + const Self = @This(); + + // Allow for compareFn to be fn(anytype, anytype) anytype + // which allows the convenient use of std.math.order. + fn compare(a: Key, b: Key) Order { + return compareFn(a, b); + } + + root: ?*Node = null, + prng: Prng = .{}, + + /// A customized pseudo random number generator for the treap. + /// This just helps reducing the memory size of the treap itself + /// as std.rand.DefaultPrng requires larger state (while producing better entropy for randomness to be fair). + const Prng = struct { + xorshift: usize = 0, + + fn random(self: *Prng, seed: usize) usize { + // Lazily seed the prng state + if (self.xorshift == 0) { + self.xorshift = seed; + } + + // Since we're using usize, decide the shifts by the integer's bit width. + const shifts = switch (@bitSizeOf(usize)) { + 64 => .{13, 7, 17}, + 32 => .{13, 17, 5}, + 16 => .{7, 9, 8}, + else => @compileError("platform not supported"), + }; + + self.xorshift ^= self.xorshift >> shifts[0]; + self.xorshift ^= self.xorshift << shifts[1]; + self.xorshift ^= self.xorshift >> shifts[2]; + + assert(self.xorshift != 0); + return self.xorshift; + } + }; + + /// A Node represents an item or point in the treap with a uniquely associated key. + pub const Node = struct { + key: Key, + priority: usize, + parent: ?*Node, + children: [2]?*Node, + }; + + /// Returns the smallest Node by key in the treap if there is one. + /// Use `getEntryForExisting()` to replace/remove this Node from the treap. + pub fn getMin(self: Self) ?*Node { + var node = self.root; + while (node) |current| { + node = current.children[0] orelse break; + } + return node; + } + + /// Returns the largest Node by key in the treap if there is one. + /// Use `getEntryForExisting()` to replace/remove this Node from the treap. + pub fn getMax(self: Self) ?*Node { + var node = self.root; + while (node) |current| { + node = current.children[1] orelse break; + } + return node; + } + + /// Lookup the Entry for the given key in the treap. + /// The Entry act's as a slot in the treap to insert/replace/remove the node associated with the key. + pub fn getEntryFor(self: *Self, key: Key) Entry { + var parent: ?*Node = undefined; + const node = self.find(key, &parent); + + return Entry{ + .key = key, + .treap = self, + .node = node, + .context = .{ .inserted_under = parent }, + }; + } + + /// Get an entry for a Node that currently exists in the treap. + /// It is undefined behavior if the Node is not currently inserted in the treap. + /// The Entry act's as a slot in the treap to insert/replace/remove the node associated with the key. + pub fn getEntryForExisting(self: *Self, node: *Node) Entry { + assert(node.priority != 0); + + return Entry{ + .key = node.key, + .treap = self, + .node = node, + .context = .{ .inserted_under = node.parent }, + }; + } + + /// An Entry represents a slot in the treap associated with a given key. + pub const Entry = struct { + key: Key, + treap: *Self, + node: ?*Node, + context: union(enum) { + /// A find() was called for this entry and the position in the treap is known. + inserted_under: ?*Node, + /// The entry's node was removed from the treap and a lookup must occur again for modification. + 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. + defer self.node = new_node; + + if (self.node) |old| { + if (new_node) |new| { + self.treap.replace(old, new); + return; + } + + self.treap.remove(old); + self.context = .removed; + return; + } + + if (new_node) |new| { + // A previous treap.remove() could have rebalanced the nodes + // so when inserting after a removal, we have to re-lookup the parent again. + // This lookup shouldn't find a node because we're yet to insert it.. + var parent: ?*Node = undefined; + switch (self.context) { + .inserted_under => |p| parent = p, + .removed => assert(self.treap.find(self.key, &parent) == null), + } + + self.treap.insert(self.key, parent, new); + self.context = .{ .inserted_under = parent }; + } + } + }; + + fn find(self: Self, key: Key, parent_ref: *?*Node) ?*Node { + var node = self.root; + parent_ref.* = null; + + // basic binary search while tracking the parent. + while (node) |current| { + const order = compare(key, current.key); + if (order == .eq) break; + + parent_ref.* = current; + node = current.children[@boolToInt(order == .gt)]; + } + + return node; + } + + fn insert(self: *Self, key: Key, parent: ?*Node, node: *Node) void { + // generate a random priority & prepare the node to be inserted into the tree + node.key = key; + node.priority = self.prng.random(@ptrToInt(node)); + node.parent = parent; + node.children = [_]?*Node{ null, null }; + + // point the parent at the new node + const link = if (parent) |p| &p.children[@boolToInt(compare(key, p.key) == .gt)] else &self.root; + assert(link.* == null); + link.* = node; + + // rotate the node up into the tree to balance it according to its priority + while (node.parent) |p| { + if (p.priority <= node.priority) break; + + const is_right = p.children[1] == @as(?*Node, node); + assert(p.children[@boolToInt(is_right)] == node); + + const rotate_right = !is_right; + self.rotate(p, rotate_right); + } + } + + fn replace(self: *Self, old: *Node, new: *Node) void { + // copy over the values from the old node + new.key = old.key; + new.priority = old.priority; + new.parent = old.parent; + new.children = old.children; + + // point the parent at the new node + const link = if (old.parent) |p| &p.children[@boolToInt(p.children[1] == old)] else &self.root; + assert(link.* == old); + link.* = new; + + // point the children's parent at the new node + for (old.children) |child_node| { + const child = child_node orelse continue; + assert(child.parent == old); + child.parent = new; + } + } + + fn remove(self: *Self, node: *Node) void { + // 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; + break :rotate_right (left.priority < right.priority); + }); + } + + // node is a now a leaf; remove by nulling out the parent's reference to it. + const link = if (node.parent) |p| &p.children[@boolToInt(p.children[1] == node)] else &self.root; + assert(link.* == node); + link.* = null; + + // clean up after ourselves + node.key = undefined; + node.priority = 0; + node.parent = null; + node.children = [_]?*Node{ null, null }; + } + + fn rotate(self: *Self, node: *Node, right: bool) void { + // if right, converts the following: + // parent -> (node (target YY adjacent) XX) + // parent -> (target YY (node adjacent XX)) + // + // if left (!right), converts the following: + // parent -> (node (target YY adjacent) XX) + // parent -> (target YY (node adjacent XX)) + const parent = node.parent; + const target = node.children[@boolToInt(!right)] orelse unreachable; + const adjacent = target.children[@boolToInt(right)]; + + // do the rotation + target.children[@boolToInt(right)] = node; + node.parent = target; + node.children[@boolToInt(!right)] = adjacent; + if (adjacent) |adj| adj.parent = node; + + // fix the parent link + const link = if (parent) |p| &p.children[@boolToInt(p.children[1] == node)] else &self.root; + assert(link.* == node); + link.* = target; + } + }; +} + +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(); + + var treap = TestTreap{}; + var nodes: [6]TestNode = undefined; + + for (nodes) |*node| { + const key = rng.int(u64); + + var entry = treap.getEntryFor(key); + try testing.expectEqual(entry.key, key); + try testing.expectEqual(entry.get(), null); + + entry.set(node); + try testing.expectEqual(entry.key, key); + try testing.expectEqual(node.key, key); + try testing.expectEqual(entry.get(), node); + } + + for (nodes) |*node| { + const key = node.key; + + var entry = treap.getEntryFor(node.key); + try testing.expectEqual(entry.key, key); + try testing.expectEqual(entry.get(), node); + + var existingEntry = treap.getEntryForExisting(node); + try testing.expectEqual(entry, existingEntry); + } +} From 4d0303b99249adcc4acdd51fa9a10884f1e5ed61 Mon Sep 17 00:00:00 2001 From: kprotty Date: Fri, 15 Apr 2022 16:30:45 -0500 Subject: [PATCH 2/3] 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); } } From 33952dad1259d9f4af09ebbad094c9b27d565525 Mon Sep 17 00:00:00 2001 From: kprotty Date: Fri, 15 Apr 2022 17:01:01 -0500 Subject: [PATCH 3/3] treap: zig fmt --- lib/std/treap.zig | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/lib/std/treap.zig b/lib/std/treap.zig index 7d08ef64b6..a74256356c 100644 --- a/lib/std/treap.zig +++ b/lib/std/treap.zig @@ -8,7 +8,7 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { const Self = @This(); // Allow for compareFn to be fn(anytype, anytype) anytype - // which allows the convenient use of std.math.order. + // which allows the convenient use of std.math.order. fn compare(a: Key, b: Key) Order { return compareFn(a, b); } @@ -30,9 +30,9 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { // Since we're using usize, decide the shifts by the integer's bit width. const shifts = switch (@bitSizeOf(usize)) { - 64 => .{13, 7, 17}, - 32 => .{13, 17, 5}, - 16 => .{7, 9, 8}, + 64 => .{ 13, 7, 17 }, + 32 => .{ 13, 17, 5 }, + 16 => .{ 7, 9, 8 }, else => @compileError("platform not supported"), }; @@ -42,7 +42,7 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { assert(self.xorshift != 0); return self.xorshift; - } + } }; /// A Node represents an item or point in the treap with a uniquely associated key. @@ -164,7 +164,7 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { return node; } - + fn insert(self: *Self, key: Key, parent: ?*Node, node: *Node) void { // generate a random priority & prepare the node to be inserted into the tree node.key = key; @@ -232,11 +232,11 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { } fn rotate(self: *Self, node: *Node, right: bool) void { - // if right, converts the following: + // if right, converts the following: // parent -> (node (target YY adjacent) XX) // parent -> (target YY (node adjacent XX)) // - // if left (!right), converts the following: + // if left (!right), converts the following: // parent -> (node (target YY adjacent) XX) // parent -> (target YY (node adjacent XX)) const parent = node.parent; @@ -389,7 +389,7 @@ test "std.Treap: insert, find, replace, remove" { 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);