authorgravatar for kbutcher6200@gmail.comkprotty <kbutcher6200@gmail.com> 2022-04-15 16:30:45-05:00
committergravatar for kbutcher6200@gmail.comkprotty <kbutcher6200@gmail.com> 2022-04-15 16:30:45-05:00
log4d0303b99249adcc4acdd51fa9a10884f1e5ed61
tree5dd02beef1e81ecfd2e5823bb0fac3228b4fdef3
parent7284eb22dce8086674c449ccf434ecdced9e4e0d

treap: fix + determinstically randomize test cases.

1 files changed, 128 insertions(+), 24 deletions(-)

lib/std/treap.zig+128-24
...@@ -103,9 +103,13 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {...@@ -103,9 +103,13 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
103103
104 /// An Entry represents a slot in the treap associated with a given key.104 /// An Entry represents a slot in the treap associated with a given key.
105 pub const Entry = struct {105 pub const Entry = struct {
106 /// The associated key for this entry.
106 key: Key,107 key: Key,
108 /// A reference to the treap this entry is apart of.
107 treap: *Self,109 treap: *Self,
110 /// The current node at this entry.
108 node: ?*Node,111 node: ?*Node,
112 /// The current state of the entry.
109 context: union(enum) {113 context: union(enum) {
110 /// A find() was called for this entry and the position in the treap is known.114 /// A find() was called for this entry and the position in the treap is known.
111 inserted_under: ?*Node,115 inserted_under: ?*Node,
...@@ -113,11 +117,6 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {...@@ -113,11 +117,6 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
113 removed,117 removed,
114 },118 },
115119
116 /// Returns the current Node at this Entry in the treap if there is one.
117 pub fn get(self: Entry) ?*Node {
118 return self.node;
119 }
120
121 /// Update's the Node at this Entry in the treap with the new node.120 /// Update's the Node at this Entry in the treap with the new node.
122 pub fn set(self: *Entry, new_node: ?*Node) void {121 pub fn set(self: *Entry, new_node: ?*Node) void {
123 // Update the entry's node reference after updating the treap below.122 // 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 {...@@ -182,7 +181,7 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
182 while (node.parent) |p| {181 while (node.parent) |p| {
183 if (p.priority <= node.priority) break;182 if (p.priority <= node.priority) break;
184183
185 const is_right = p.children[1] == @as(?*Node, node);184 const is_right = p.children[1] == node;
186 assert(p.children[@boolToInt(is_right)] == node);185 assert(p.children[@boolToInt(is_right)] == node);
187186
188 const rotate_right = !is_right;187 const rotate_right = !is_right;
...@@ -214,8 +213,8 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {...@@ -214,8 +213,8 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
214 // rotate the node down to be a leaf of the tree for removal, respecting priorities.213 // rotate the node down to be a leaf of the tree for removal, respecting priorities.
215 while (node.children[0] orelse node.children[1]) |_| {214 while (node.children[0] orelse node.children[1]) |_| {
216 self.rotate(node, rotate_right: {215 self.rotate(node, rotate_right: {
217 const right = node.children[0] orelse break :rotate_right true;216 const right = node.children[1] orelse break :rotate_right true;
218 const left = node.children[1] orelse break :rotate_right false;217 const left = node.children[0] orelse break :rotate_right false;
219 break :rotate_right (left.priority < right.priority);218 break :rotate_right (left.priority < right.priority);
220 });219 });
221 }220 }
...@@ -244,10 +243,13 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {...@@ -244,10 +243,13 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
244 const target = node.children[@boolToInt(!right)] orelse unreachable;243 const target = node.children[@boolToInt(!right)] orelse unreachable;
245 const adjacent = target.children[@boolToInt(right)];244 const adjacent = target.children[@boolToInt(right)];
246245
247 // do the rotation246 // rotate the children
248 target.children[@boolToInt(right)] = node;247 target.children[@boolToInt(right)] = node;
249 node.parent = target;
250 node.children[@boolToInt(!right)] = adjacent;248 node.children[@boolToInt(!right)] = adjacent;
249
250 // rotate the parents
251 node.parent = target;
252 target.parent = parent;
251 if (adjacent) |adj| adj.parent = node;253 if (adjacent) |adj| adj.parent = node;
252254
253 // fix the parent link255 // fix the parent link
...@@ -258,37 +260,139 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {...@@ -258,37 +260,139 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
258 };260 };
259}261}
260262
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
261const TestTreap = Treap(u64, std.math.order);308const TestTreap = Treap(u64, std.math.order);
262const TestNode = TestTreap.Node;309const TestNode = TestTreap.Node;
263310
264test "std.Treap: insert, find, remove" {311test "std.Treap: insert, find, replace, remove" {
265 var prng = std.rand.DefaultPrng.init(0xdeadbeef);
266 var rng = prng.random();
267
268 var treap = TestTreap{};312 var treap = TestTreap{};
269 var nodes: [6]TestNode = undefined;313 var nodes: [10]TestNode = undefined;
270314
271 for (nodes) |*node| {315 var prng = std.rand.DefaultPrng.init(0xdeadbeef);
272 const key = rng.int(u64);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);
273322
323 // make sure the current entry is empty.
274 var entry = treap.getEntryFor(key);324 var entry = treap.getEntryFor(key);
275 try testing.expectEqual(entry.key, key);325 try testing.expectEqual(entry.key, key);
276 try testing.expectEqual(entry.get(), null);326 try testing.expectEqual(entry.node, null);
277327
328 // insert the entry and make sure the fields are correct.
278 entry.set(node);329 entry.set(node);
279 try testing.expectEqual(entry.key, key);
280 try testing.expectEqual(node.key, key);330 try testing.expectEqual(node.key, key);
281 try testing.expectEqual(entry.get(), node);331 try testing.expectEqual(entry.key, key);
332 try testing.expectEqual(entry.node, node);
282 }333 }
283334
284 for (nodes) |*node| {335 // find check
336 iter.reset();
337 while (iter.next()) |node| {
285 const key = node.key;338 const key = node.key;
286339
340 // find the entry by-key and by-node after having been inserted.
287 var entry = treap.getEntryFor(node.key);341 var entry = treap.getEntryFor(node.key);
288 try testing.expectEqual(entry.key, key);342 try testing.expectEqual(entry.key, key);
289 try testing.expectEqual(entry.get(), node);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);
290364
291 var existingEntry = treap.getEntryForExisting(node);365 // replace the stub_node back to the node and ensure future finds point to the old node.
292 try testing.expectEqual(entry, existingEntry);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);
293 }397 }
294}398}