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