authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2022-04-16 06:23:18-04:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2022-04-16 06:23:18-04:00
logec8a6544a379eb2a0ed49c539f3c11eb27f26f98
treedc6edba583585ad637a628fd4dc5b392248d3d94
parent578a792b33d1f26c839cd6c77eaa2680601caaf2
parent33952dad1259d9f4af09ebbad094c9b27d565525
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #11444 from kprotty/treap

Introduce std.Treap

3 files changed, 400 insertions(+), 0 deletions(-)

CMakeLists.txt+1
......@@ -539,6 +539,7 @@ set(ZIG_STAGE2_SOURCES
539539 "${CMAKE_SOURCE_DIR}/lib/std/Thread/ResetEvent.zig"
540540 "${CMAKE_SOURCE_DIR}/lib/std/Thread/StaticResetEvent.zig"
541541 "${CMAKE_SOURCE_DIR}/lib/std/time.zig"
542 "${CMAKE_SOURCE_DIR}/lib/std/treap.zig"
542543 "${CMAKE_SOURCE_DIR}/lib/std/unicode.zig"
543544 "${CMAKE_SOURCE_DIR}/lib/std/zig.zig"
544545 "${CMAKE_SOURCE_DIR}/lib/std/zig/Ast.zig"
lib/std/std.zig+1
......@@ -39,6 +39,7 @@ pub const StringArrayHashMapUnmanaged = array_hash_map.StringArrayHashMapUnmanag
3939pub const TailQueue = @import("linked_list.zig").TailQueue;
4040pub const Target = @import("target.zig").Target;
4141pub const Thread = @import("Thread.zig");
42pub const Treap = @import("treap.zig").Treap;
4243pub const Tz = @import("tz.zig").Tz;
4344
4445pub const array_hash_map = @import("array_hash_map.zig");
lib/std/treap.zig created+398
......@@ -0,0 +1,398 @@
1const std = @import("std.zig");
2const assert = std.debug.assert;
3const testing = std.testing;
4const Order = std.math.Order;
5
6pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
7 return struct {
8 const Self = @This();
9
10 // Allow for compareFn to be fn(anytype, anytype) anytype
11 // which allows the convenient use of std.math.order.
12 fn compare(a: Key, b: Key) Order {
13 return compareFn(a, b);
14 }
15
16 root: ?*Node = null,
17 prng: Prng = .{},
18
19 /// A customized pseudo random number generator for the treap.
20 /// This just helps reducing the memory size of the treap itself
21 /// as std.rand.DefaultPrng requires larger state (while producing better entropy for randomness to be fair).
22 const Prng = struct {
23 xorshift: usize = 0,
24
25 fn random(self: *Prng, seed: usize) usize {
26 // Lazily seed the prng state
27 if (self.xorshift == 0) {
28 self.xorshift = seed;
29 }
30
31 // Since we're using usize, decide the shifts by the integer's bit width.
32 const shifts = switch (@bitSizeOf(usize)) {
33 64 => .{ 13, 7, 17 },
34 32 => .{ 13, 17, 5 },
35 16 => .{ 7, 9, 8 },
36 else => @compileError("platform not supported"),
37 };
38
39 self.xorshift ^= self.xorshift >> shifts[0];
40 self.xorshift ^= self.xorshift << shifts[1];
41 self.xorshift ^= self.xorshift >> shifts[2];
42
43 assert(self.xorshift != 0);
44 return self.xorshift;
45 }
46 };
47
48 /// A Node represents an item or point in the treap with a uniquely associated key.
49 pub const Node = struct {
50 key: Key,
51 priority: usize,
52 parent: ?*Node,
53 children: [2]?*Node,
54 };
55
56 /// Returns the smallest Node by key in the treap if there is one.
57 /// Use `getEntryForExisting()` to replace/remove this Node from the treap.
58 pub fn getMin(self: Self) ?*Node {
59 var node = self.root;
60 while (node) |current| {
61 node = current.children[0] orelse break;
62 }
63 return node;
64 }
65
66 /// Returns the largest Node by key in the treap if there is one.
67 /// Use `getEntryForExisting()` to replace/remove this Node from the treap.
68 pub fn getMax(self: Self) ?*Node {
69 var node = self.root;
70 while (node) |current| {
71 node = current.children[1] orelse break;
72 }
73 return node;
74 }
75
76 /// Lookup the Entry for the given key in the treap.
77 /// The Entry act's as a slot in the treap to insert/replace/remove the node associated with the key.
78 pub fn getEntryFor(self: *Self, key: Key) Entry {
79 var parent: ?*Node = undefined;
80 const node = self.find(key, &parent);
81
82 return Entry{
83 .key = key,
84 .treap = self,
85 .node = node,
86 .context = .{ .inserted_under = parent },
87 };
88 }
89
90 /// Get an entry for a Node that currently exists in the treap.
91 /// It is undefined behavior if the Node is not currently inserted in the treap.
92 /// The Entry act's as a slot in the treap to insert/replace/remove the node associated with the key.
93 pub fn getEntryForExisting(self: *Self, node: *Node) Entry {
94 assert(node.priority != 0);
95
96 return Entry{
97 .key = node.key,
98 .treap = self,
99 .node = node,
100 .context = .{ .inserted_under = node.parent },
101 };
102 }
103
104 /// An Entry represents a slot in the treap associated with a given key.
105 pub const Entry = struct {
106 /// The associated key for this entry.
107 key: Key,
108 /// A reference to the treap this entry is apart of.
109 treap: *Self,
110 /// The current node at this entry.
111 node: ?*Node,
112 /// The current state of the entry.
113 context: union(enum) {
114 /// A find() was called for this entry and the position in the treap is known.
115 inserted_under: ?*Node,
116 /// The entry's node was removed from the treap and a lookup must occur again for modification.
117 removed,
118 },
119
120 /// Update's the Node at this Entry in the treap with the new node.
121 pub fn set(self: *Entry, new_node: ?*Node) void {
122 // Update the entry's node reference after updating the treap below.
123 defer self.node = new_node;
124
125 if (self.node) |old| {
126 if (new_node) |new| {
127 self.treap.replace(old, new);
128 return;
129 }
130
131 self.treap.remove(old);
132 self.context = .removed;
133 return;
134 }
135
136 if (new_node) |new| {
137 // A previous treap.remove() could have rebalanced the nodes
138 // so when inserting after a removal, we have to re-lookup the parent again.
139 // This lookup shouldn't find a node because we're yet to insert it..
140 var parent: ?*Node = undefined;
141 switch (self.context) {
142 .inserted_under => |p| parent = p,
143 .removed => assert(self.treap.find(self.key, &parent) == null),
144 }
145
146 self.treap.insert(self.key, parent, new);
147 self.context = .{ .inserted_under = parent };
148 }
149 }
150 };
151
152 fn find(self: Self, key: Key, parent_ref: *?*Node) ?*Node {
153 var node = self.root;
154 parent_ref.* = null;
155
156 // basic binary search while tracking the parent.
157 while (node) |current| {
158 const order = compare(key, current.key);
159 if (order == .eq) break;
160
161 parent_ref.* = current;
162 node = current.children[@boolToInt(order == .gt)];
163 }
164
165 return node;
166 }
167
168 fn insert(self: *Self, key: Key, parent: ?*Node, node: *Node) void {
169 // generate a random priority & prepare the node to be inserted into the tree
170 node.key = key;
171 node.priority = self.prng.random(@ptrToInt(node));
172 node.parent = parent;
173 node.children = [_]?*Node{ null, null };
174
175 // point the parent at the new node
176 const link = if (parent) |p| &p.children[@boolToInt(compare(key, p.key) == .gt)] else &self.root;
177 assert(link.* == null);
178 link.* = node;
179
180 // rotate the node up into the tree to balance it according to its priority
181 while (node.parent) |p| {
182 if (p.priority <= node.priority) break;
183
184 const is_right = p.children[1] == node;
185 assert(p.children[@boolToInt(is_right)] == node);
186
187 const rotate_right = !is_right;
188 self.rotate(p, rotate_right);
189 }
190 }
191
192 fn replace(self: *Self, old: *Node, new: *Node) void {
193 // copy over the values from the old node
194 new.key = old.key;
195 new.priority = old.priority;
196 new.parent = old.parent;
197 new.children = old.children;
198
199 // point the parent at the new node
200 const link = if (old.parent) |p| &p.children[@boolToInt(p.children[1] == old)] else &self.root;
201 assert(link.* == old);
202 link.* = new;
203
204 // point the children's parent at the new node
205 for (old.children) |child_node| {
206 const child = child_node orelse continue;
207 assert(child.parent == old);
208 child.parent = new;
209 }
210 }
211
212 fn remove(self: *Self, node: *Node) void {
213 // rotate the node down to be a leaf of the tree for removal, respecting priorities.
214 while (node.children[0] orelse node.children[1]) |_| {
215 self.rotate(node, rotate_right: {
216 const right = node.children[1] orelse break :rotate_right true;
217 const left = node.children[0] orelse break :rotate_right false;
218 break :rotate_right (left.priority < right.priority);
219 });
220 }
221
222 // node is a now a leaf; remove by nulling out the parent's reference to it.
223 const link = if (node.parent) |p| &p.children[@boolToInt(p.children[1] == node)] else &self.root;
224 assert(link.* == node);
225 link.* = null;
226
227 // clean up after ourselves
228 node.key = undefined;
229 node.priority = 0;
230 node.parent = null;
231 node.children = [_]?*Node{ null, null };
232 }
233
234 fn rotate(self: *Self, node: *Node, right: bool) void {
235 // if right, converts the following:
236 // parent -> (node (target YY adjacent) XX)
237 // parent -> (target YY (node adjacent XX))
238 //
239 // if left (!right), converts the following:
240 // parent -> (node (target YY adjacent) XX)
241 // parent -> (target YY (node adjacent XX))
242 const parent = node.parent;
243 const target = node.children[@boolToInt(!right)] orelse unreachable;
244 const adjacent = target.children[@boolToInt(right)];
245
246 // rotate the children
247 target.children[@boolToInt(right)] = node;
248 node.children[@boolToInt(!right)] = adjacent;
249
250 // rotate the parents
251 node.parent = target;
252 target.parent = parent;
253 if (adjacent) |adj| adj.parent = node;
254
255 // fix the parent link
256 const link = if (parent) |p| &p.children[@boolToInt(p.children[1] == node)] else &self.root;
257 assert(link.* == node);
258 link.* = target;
259 }
260 };
261}
262
263// For iterating a slice in a random order
264// https://lemire.me/blog/2017/09/18/visiting-all-values-in-an-array-exactly-once-in-random-order/
265fn SliceIterRandomOrder(comptime T: type) type {
266 return struct {
267 rng: std.rand.Random,
268 slice: []T,
269 index: usize = undefined,
270 offset: usize = undefined,
271 co_prime: usize,
272
273 const Self = @This();
274
275 pub fn init(slice: []T, rng: std.rand.Random) Self {
276 return Self{
277 .rng = rng,
278 .slice = slice,
279 .co_prime = blk: {
280 if (slice.len == 0) break :blk 0;
281 var prime = slice.len / 2;
282 while (prime < slice.len) : (prime += 1) {
283 var gcd = [_]usize{ prime, slice.len };
284 while (gcd[1] != 0) {
285 const temp = gcd;
286 gcd = [_]usize{ temp[1], temp[0] % temp[1] };
287 }
288 if (gcd[0] == 1) break;
289 }
290 break :blk prime;
291 },
292 };
293 }
294
295 pub fn reset(self: *Self) void {
296 self.index = 0;
297 self.offset = self.rng.int(usize);
298 }
299
300 pub fn next(self: *Self) ?*T {
301 if (self.index >= self.slice.len) return null;
302 defer self.index += 1;
303 return &self.slice[((self.index *% self.co_prime) +% self.offset) % self.slice.len];
304 }
305 };
306}
307
308const TestTreap = Treap(u64, std.math.order);
309const TestNode = TestTreap.Node;
310
311test "std.Treap: insert, find, replace, remove" {
312 var treap = TestTreap{};
313 var nodes: [10]TestNode = undefined;
314
315 var prng = std.rand.DefaultPrng.init(0xdeadbeef);
316 var iter = SliceIterRandomOrder(TestNode).init(&nodes, prng.random());
317
318 // insert check
319 iter.reset();
320 while (iter.next()) |node| {
321 const key = prng.random().int(u64);
322
323 // make sure the current entry is empty.
324 var entry = treap.getEntryFor(key);
325 try testing.expectEqual(entry.key, key);
326 try testing.expectEqual(entry.node, null);
327
328 // insert the entry and make sure the fields are correct.
329 entry.set(node);
330 try testing.expectEqual(node.key, key);
331 try testing.expectEqual(entry.key, key);
332 try testing.expectEqual(entry.node, node);
333 }
334
335 // find check
336 iter.reset();
337 while (iter.next()) |node| {
338 const key = node.key;
339
340 // find the entry by-key and by-node after having been inserted.
341 var entry = treap.getEntryFor(node.key);
342 try testing.expectEqual(entry.key, key);
343 try testing.expectEqual(entry.node, node);
344 try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node);
345 }
346
347 // replace check
348 iter.reset();
349 while (iter.next()) |node| {
350 const key = node.key;
351
352 // find the entry by node since we already know it exists
353 var entry = treap.getEntryForExisting(node);
354 try testing.expectEqual(entry.key, key);
355 try testing.expectEqual(entry.node, node);
356
357 var stub_node: TestNode = undefined;
358
359 // replace the node with a stub_node and ensure future finds point to the stub_node.
360 entry.set(&stub_node);
361 try testing.expectEqual(entry.node, &stub_node);
362 try testing.expectEqual(entry.node, treap.getEntryFor(key).node);
363 try testing.expectEqual(entry.node, treap.getEntryForExisting(&stub_node).node);
364
365 // replace the stub_node back to the node and ensure future finds point to the old node.
366 entry.set(node);
367 try testing.expectEqual(entry.node, node);
368 try testing.expectEqual(entry.node, treap.getEntryFor(key).node);
369 try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node);
370 }
371
372 // remove check
373 iter.reset();
374 while (iter.next()) |node| {
375 const key = node.key;
376
377 // find the entry by node since we already know it exists
378 var entry = treap.getEntryForExisting(node);
379 try testing.expectEqual(entry.key, key);
380 try testing.expectEqual(entry.node, node);
381
382 // remove the node at the entry and ensure future finds point to it being removed.
383 entry.set(null);
384 try testing.expectEqual(entry.node, null);
385 try testing.expectEqual(entry.node, treap.getEntryFor(key).node);
386
387 // insert the node back and ensure future finds point to the inserted node
388 entry.set(node);
389 try testing.expectEqual(entry.node, node);
390 try testing.expectEqual(entry.node, treap.getEntryFor(key).node);
391 try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node);
392
393 // remove the node again and make sure it was cleared after the insert
394 entry.set(null);
395 try testing.expectEqual(entry.node, null);
396 try testing.expectEqual(entry.node, treap.getEntryFor(key).node);
397 }
398}