github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/interval/generic/example_interval_btree.go (about) 1 // Code generated by go_generics. DO NOT EDIT. 2 3 // Copyright 2020 The Cockroach Authors. 4 // 5 // Use of this software is governed by the Business Source License 6 // included in the file licenses/BSL.txt. 7 // 8 // As of the Change Date specified in that file, in accordance with 9 // the Business Source License, use of this software will be governed 10 // by the Apache License, Version 2.0, included in the file 11 // licenses/APL.txt. 12 13 package generic 14 15 import ( 16 "bytes" 17 "sort" 18 "strings" 19 "sync" 20 "sync/atomic" 21 "unsafe" 22 ) 23 24 // nilT is a nil instance of the Template type. 25 var nilT *example 26 27 const ( 28 degree = 16 29 maxItems = 2*degree - 1 30 minItems = degree - 1 31 ) 32 33 // cmp returns a value indicating the sort order relationship between 34 // a and b. The comparison is performed lexicographically on 35 // (a.Key(), a.EndKey(), a.ID()) 36 // and 37 // (b.Key(), b.EndKey(), b.ID()) 38 // tuples. 39 // 40 // Given c = cmp(a, b): 41 // 42 // c == -1 if (a.Key(), a.EndKey(), a.ID()) < (b.Key(), b.EndKey(), b.ID()) 43 // c == 0 if (a.Key(), a.EndKey(), a.ID()) == (b.Key(), b.EndKey(), b.ID()) 44 // c == 1 if (a.Key(), a.EndKey(), a.ID()) > (b.Key(), b.EndKey(), b.ID()) 45 // 46 func cmp(a, b *example) int { 47 c := bytes.Compare(a.Key(), b.Key()) 48 if c != 0 { 49 return c 50 } 51 c = bytes.Compare(a.EndKey(), b.EndKey()) 52 if c != 0 { 53 return c 54 } 55 if a.ID() < b.ID() { 56 return -1 57 } else if a.ID() > b.ID() { 58 return 1 59 } else { 60 return 0 61 } 62 } 63 64 // keyBound represents the upper-bound of a key range. 65 type keyBound struct { 66 key []byte 67 inc bool 68 } 69 70 func (b keyBound) compare(o keyBound) int { 71 c := bytes.Compare(b.key, o.key) 72 if c != 0 { 73 return c 74 } 75 if b.inc == o.inc { 76 return 0 77 } 78 if b.inc { 79 return 1 80 } 81 return -1 82 } 83 84 func (b keyBound) contains(a *example) bool { 85 c := bytes.Compare(a.Key(), b.key) 86 if c == 0 { 87 return b.inc 88 } 89 return c < 0 90 } 91 92 func upperBound(c *example) keyBound { 93 if len(c.EndKey()) != 0 { 94 return keyBound{key: c.EndKey()} 95 } 96 return keyBound{key: c.Key(), inc: true} 97 } 98 99 type leafNode struct { 100 ref int32 101 count int16 102 leaf bool 103 max keyBound 104 items [maxItems]*example 105 } 106 107 type node struct { 108 leafNode 109 children [maxItems + 1]*node 110 } 111 112 //go:nocheckptr casts a ptr to a smaller struct to a ptr to a larger struct. 113 func leafToNode(ln *leafNode) *node { 114 return (*node)(unsafe.Pointer(ln)) 115 } 116 117 func nodeToLeaf(n *node) *leafNode { 118 return (*leafNode)(unsafe.Pointer(n)) 119 } 120 121 var leafPool = sync.Pool{ 122 New: func() interface{} { 123 return new(leafNode) 124 }, 125 } 126 127 var nodePool = sync.Pool{ 128 New: func() interface{} { 129 return new(node) 130 }, 131 } 132 133 func newLeafNode() *node { 134 n := leafToNode(leafPool.Get().(*leafNode)) 135 n.leaf = true 136 n.ref = 1 137 return n 138 } 139 140 func newNode() *node { 141 n := nodePool.Get().(*node) 142 n.ref = 1 143 return n 144 } 145 146 // mut creates and returns a mutable node reference. If the node is not shared 147 // with any other trees then it can be modified in place. Otherwise, it must be 148 // cloned to ensure unique ownership. In this way, we enforce a copy-on-write 149 // policy which transparently incorporates the idea of local mutations, like 150 // Clojure's transients or Haskell's ST monad, where nodes are only copied 151 // during the first time that they are modified between Clone operations. 152 // 153 // When a node is cloned, the provided pointer will be redirected to the new 154 // mutable node. 155 func mut(n **node) *node { 156 if atomic.LoadInt32(&(*n).ref) == 1 { 157 // Exclusive ownership. Can mutate in place. 158 return *n 159 } 160 // If we do not have unique ownership over the node then we 161 // clone it to gain unique ownership. After doing so, we can 162 // release our reference to the old node. We pass recursive 163 // as true because even though we just observed the node's 164 // reference count to be greater than 1, we might be racing 165 // with another call to decRef on this node. 166 c := (*n).clone() 167 (*n).decRef(true /* recursive */) 168 *n = c 169 return *n 170 } 171 172 // incRef acquires a reference to the node. 173 func (n *node) incRef() { 174 atomic.AddInt32(&n.ref, 1) 175 } 176 177 // decRef releases a reference to the node. If requested, the method 178 // will recurse into child nodes and decrease their refcounts as well. 179 func (n *node) decRef(recursive bool) { 180 if atomic.AddInt32(&n.ref, -1) > 0 { 181 // Other references remain. Can't free. 182 return 183 } 184 // Clear and release node into memory pool. 185 if n.leaf { 186 ln := nodeToLeaf(n) 187 *ln = leafNode{} 188 leafPool.Put(ln) 189 } else { 190 // Release child references first, if requested. 191 if recursive { 192 for i := int16(0); i <= n.count; i++ { 193 n.children[i].decRef(true /* recursive */) 194 } 195 } 196 *n = node{} 197 nodePool.Put(n) 198 } 199 } 200 201 // clone creates a clone of the receiver with a single reference count. 202 func (n *node) clone() *node { 203 var c *node 204 if n.leaf { 205 c = newLeafNode() 206 } else { 207 c = newNode() 208 } 209 // NB: copy field-by-field without touching n.ref to avoid 210 // triggering the race detector and looking like a data race. 211 c.count = n.count 212 c.max = n.max 213 c.items = n.items 214 if !c.leaf { 215 // Copy children and increase each refcount. 216 c.children = n.children 217 for i := int16(0); i <= c.count; i++ { 218 c.children[i].incRef() 219 } 220 } 221 return c 222 } 223 224 func (n *node) insertAt(index int, item *example, nd *node) { 225 if index < int(n.count) { 226 copy(n.items[index+1:n.count+1], n.items[index:n.count]) 227 if !n.leaf { 228 copy(n.children[index+2:n.count+2], n.children[index+1:n.count+1]) 229 } 230 } 231 n.items[index] = item 232 if !n.leaf { 233 n.children[index+1] = nd 234 } 235 n.count++ 236 } 237 238 func (n *node) pushBack(item *example, nd *node) { 239 n.items[n.count] = item 240 if !n.leaf { 241 n.children[n.count+1] = nd 242 } 243 n.count++ 244 } 245 246 func (n *node) pushFront(item *example, nd *node) { 247 if !n.leaf { 248 copy(n.children[1:n.count+2], n.children[:n.count+1]) 249 n.children[0] = nd 250 } 251 copy(n.items[1:n.count+1], n.items[:n.count]) 252 n.items[0] = item 253 n.count++ 254 } 255 256 // removeAt removes a value at a given index, pulling all subsequent values 257 // back. 258 func (n *node) removeAt(index int) (*example, *node) { 259 var child *node 260 if !n.leaf { 261 child = n.children[index+1] 262 copy(n.children[index+1:n.count], n.children[index+2:n.count+1]) 263 n.children[n.count] = nil 264 } 265 n.count-- 266 out := n.items[index] 267 copy(n.items[index:n.count], n.items[index+1:n.count+1]) 268 n.items[n.count] = nilT 269 return out, child 270 } 271 272 // popBack removes and returns the last element in the list. 273 func (n *node) popBack() (*example, *node) { 274 n.count-- 275 out := n.items[n.count] 276 n.items[n.count] = nilT 277 if n.leaf { 278 return out, nil 279 } 280 child := n.children[n.count+1] 281 n.children[n.count+1] = nil 282 return out, child 283 } 284 285 // popFront removes and returns the first element in the list. 286 func (n *node) popFront() (*example, *node) { 287 n.count-- 288 var child *node 289 if !n.leaf { 290 child = n.children[0] 291 copy(n.children[:n.count+1], n.children[1:n.count+2]) 292 n.children[n.count+1] = nil 293 } 294 out := n.items[0] 295 copy(n.items[:n.count], n.items[1:n.count+1]) 296 n.items[n.count] = nilT 297 return out, child 298 } 299 300 // find returns the index where the given item should be inserted into this 301 // list. 'found' is true if the item already exists in the list at the given 302 // index. 303 func (n *node) find(item *example) (index int, found bool) { 304 // Logic copied from sort.Search. Inlining this gave 305 // an 11% speedup on BenchmarkBTreeDeleteInsert. 306 i, j := 0, int(n.count) 307 for i < j { 308 h := int(uint(i+j) >> 1) // avoid overflow when computing h 309 // i ≤ h < j 310 v := cmp(item, n.items[h]) 311 if v == 0 { 312 return h, true 313 } else if v > 0 { 314 i = h + 1 315 } else { 316 j = h 317 } 318 } 319 return i, false 320 } 321 322 // split splits the given node at the given index. The current node shrinks, 323 // and this function returns the item that existed at that index and a new 324 // node containing all items/children after it. 325 // 326 // Before: 327 // 328 // +-----------+ 329 // | x y z | 330 // +--/-/-\-\--+ 331 // 332 // After: 333 // 334 // +-----------+ 335 // | y | 336 // +----/-\----+ 337 // / \ 338 // v v 339 // +-----------+ +-----------+ 340 // | x | | z | 341 // +-----------+ +-----------+ 342 // 343 func (n *node) split(i int) (*example, *node) { 344 out := n.items[i] 345 var next *node 346 if n.leaf { 347 next = newLeafNode() 348 } else { 349 next = newNode() 350 } 351 next.count = n.count - int16(i+1) 352 copy(next.items[:], n.items[i+1:n.count]) 353 for j := int16(i); j < n.count; j++ { 354 n.items[j] = nilT 355 } 356 if !n.leaf { 357 copy(next.children[:], n.children[i+1:n.count+1]) 358 for j := int16(i + 1); j <= n.count; j++ { 359 n.children[j] = nil 360 } 361 } 362 n.count = int16(i) 363 364 next.max = next.findUpperBound() 365 if n.max.compare(next.max) != 0 && n.max.compare(upperBound(out)) != 0 { 366 // If upper bound wasn't from new node or item 367 // at index i, it must still be from old node. 368 } else { 369 n.max = n.findUpperBound() 370 } 371 return out, next 372 } 373 374 // insert inserts a item into the subtree rooted at this node, making sure no 375 // nodes in the subtree exceed maxItems items. Returns true if an existing item 376 // was replaced and false if a item was inserted. Also returns whether the 377 // node's upper bound changes. 378 func (n *node) insert(item *example) (replaced, newBound bool) { 379 i, found := n.find(item) 380 if found { 381 n.items[i] = item 382 return true, false 383 } 384 if n.leaf { 385 n.insertAt(i, item, nil) 386 return false, n.adjustUpperBoundOnInsertion(item, nil) 387 } 388 if n.children[i].count >= maxItems { 389 splitLa, splitNode := mut(&n.children[i]).split(maxItems / 2) 390 n.insertAt(i, splitLa, splitNode) 391 392 switch cmp := cmp(item, n.items[i]); { 393 case cmp < 0: 394 // no change, we want first split node 395 case cmp > 0: 396 i++ // we want second split node 397 default: 398 n.items[i] = item 399 return true, false 400 } 401 } 402 replaced, newBound = mut(&n.children[i]).insert(item) 403 if newBound { 404 newBound = n.adjustUpperBoundOnInsertion(item, nil) 405 } 406 return replaced, newBound 407 } 408 409 // removeMax removes and returns the maximum item from the subtree rooted 410 // at this node. 411 func (n *node) removeMax() *example { 412 if n.leaf { 413 n.count-- 414 out := n.items[n.count] 415 n.items[n.count] = nilT 416 n.adjustUpperBoundOnRemoval(out, nil) 417 return out 418 } 419 child := mut(&n.children[n.count]) 420 if child.count <= minItems { 421 n.rebalanceOrMerge(int(n.count)) 422 return n.removeMax() 423 } 424 return child.removeMax() 425 } 426 427 // remove removes a item from the subtree rooted at this node. Returns 428 // the item that was removed or nil if no matching item was found. 429 // Also returns whether the node's upper bound changes. 430 func (n *node) remove(item *example) (out *example, newBound bool) { 431 i, found := n.find(item) 432 if n.leaf { 433 if found { 434 out, _ = n.removeAt(i) 435 return out, n.adjustUpperBoundOnRemoval(out, nil) 436 } 437 return nilT, false 438 } 439 if n.children[i].count <= minItems { 440 // Child not large enough to remove from. 441 n.rebalanceOrMerge(i) 442 return n.remove(item) 443 } 444 child := mut(&n.children[i]) 445 if found { 446 // Replace the item being removed with the max item in our left child. 447 out = n.items[i] 448 n.items[i] = child.removeMax() 449 return out, n.adjustUpperBoundOnRemoval(out, nil) 450 } 451 // Latch is not in this node and child is large enough to remove from. 452 out, newBound = child.remove(item) 453 if newBound { 454 newBound = n.adjustUpperBoundOnRemoval(out, nil) 455 } 456 return out, newBound 457 } 458 459 // rebalanceOrMerge grows child 'i' to ensure it has sufficient room to remove 460 // a item from it while keeping it at or above minItems. 461 func (n *node) rebalanceOrMerge(i int) { 462 switch { 463 case i > 0 && n.children[i-1].count > minItems: 464 // Rebalance from left sibling. 465 // 466 // +-----------+ 467 // | y | 468 // +----/-\----+ 469 // / \ 470 // v v 471 // +-----------+ +-----------+ 472 // | x | | | 473 // +----------\+ +-----------+ 474 // \ 475 // v 476 // a 477 // 478 // After: 479 // 480 // +-----------+ 481 // | x | 482 // +----/-\----+ 483 // / \ 484 // v v 485 // +-----------+ +-----------+ 486 // | | | y | 487 // +-----------+ +/----------+ 488 // / 489 // v 490 // a 491 // 492 left := mut(&n.children[i-1]) 493 child := mut(&n.children[i]) 494 xLa, grandChild := left.popBack() 495 yLa := n.items[i-1] 496 child.pushFront(yLa, grandChild) 497 n.items[i-1] = xLa 498 499 left.adjustUpperBoundOnRemoval(xLa, grandChild) 500 child.adjustUpperBoundOnInsertion(yLa, grandChild) 501 502 case i < int(n.count) && n.children[i+1].count > minItems: 503 // Rebalance from right sibling. 504 // 505 // +-----------+ 506 // | y | 507 // +----/-\----+ 508 // / \ 509 // v v 510 // +-----------+ +-----------+ 511 // | | | x | 512 // +-----------+ +/----------+ 513 // / 514 // v 515 // a 516 // 517 // After: 518 // 519 // +-----------+ 520 // | x | 521 // +----/-\----+ 522 // / \ 523 // v v 524 // +-----------+ +-----------+ 525 // | y | | | 526 // +----------\+ +-----------+ 527 // \ 528 // v 529 // a 530 // 531 right := mut(&n.children[i+1]) 532 child := mut(&n.children[i]) 533 xLa, grandChild := right.popFront() 534 yLa := n.items[i] 535 child.pushBack(yLa, grandChild) 536 n.items[i] = xLa 537 538 right.adjustUpperBoundOnRemoval(xLa, grandChild) 539 child.adjustUpperBoundOnInsertion(yLa, grandChild) 540 541 default: 542 // Merge with either the left or right sibling. 543 // 544 // +-----------+ 545 // | u y v | 546 // +----/-\----+ 547 // / \ 548 // v v 549 // +-----------+ +-----------+ 550 // | x | | z | 551 // +-----------+ +-----------+ 552 // 553 // After: 554 // 555 // +-----------+ 556 // | u v | 557 // +-----|-----+ 558 // | 559 // v 560 // +-----------+ 561 // | x y z | 562 // +-----------+ 563 // 564 if i >= int(n.count) { 565 i = int(n.count - 1) 566 } 567 child := mut(&n.children[i]) 568 // Make mergeChild mutable, bumping the refcounts on its children if necessary. 569 _ = mut(&n.children[i+1]) 570 mergeLa, mergeChild := n.removeAt(i) 571 child.items[child.count] = mergeLa 572 copy(child.items[child.count+1:], mergeChild.items[:mergeChild.count]) 573 if !child.leaf { 574 copy(child.children[child.count+1:], mergeChild.children[:mergeChild.count+1]) 575 } 576 child.count += mergeChild.count + 1 577 578 child.adjustUpperBoundOnInsertion(mergeLa, mergeChild) 579 mergeChild.decRef(false /* recursive */) 580 } 581 } 582 583 // findUpperBound returns the largest end key node range, assuming that its 584 // children have correct upper bounds already set. 585 func (n *node) findUpperBound() keyBound { 586 var max keyBound 587 for i := int16(0); i < n.count; i++ { 588 up := upperBound(n.items[i]) 589 if max.compare(up) < 0 { 590 max = up 591 } 592 } 593 if !n.leaf { 594 for i := int16(0); i <= n.count; i++ { 595 up := n.children[i].max 596 if max.compare(up) < 0 { 597 max = up 598 } 599 } 600 } 601 return max 602 } 603 604 // adjustUpperBoundOnInsertion adjusts the upper key bound for this node 605 // given a item and an optional child node that was inserted. Returns 606 // true is the upper bound was changed and false if not. 607 func (n *node) adjustUpperBoundOnInsertion(item *example, child *node) bool { 608 up := upperBound(item) 609 if child != nil { 610 if up.compare(child.max) < 0 { 611 up = child.max 612 } 613 } 614 if n.max.compare(up) < 0 { 615 n.max = up 616 return true 617 } 618 return false 619 } 620 621 // adjustUpperBoundOnRemoval adjusts the upper key bound for this node 622 // given a item and an optional child node that were removed. Returns 623 // true is the upper bound was changed and false if not. 624 func (n *node) adjustUpperBoundOnRemoval(item *example, child *node) bool { 625 up := upperBound(item) 626 if child != nil { 627 if up.compare(child.max) < 0 { 628 up = child.max 629 } 630 } 631 if n.max.compare(up) == 0 { 632 // up was previous upper bound of n. 633 n.max = n.findUpperBound() 634 return n.max.compare(up) != 0 635 } 636 return false 637 } 638 639 // btree is an implementation of an augmented interval B-Tree. 640 // 641 // btree stores items in an ordered structure, allowing easy insertion, 642 // removal, and iteration. It represents intervals and permits an interval 643 // search operation following the approach laid out in CLRS, Chapter 14. 644 // The B-Tree stores items in order based on their start key and each 645 // B-Tree node maintains the upper-bound end key of all items in its 646 // subtree. 647 // 648 // Write operations are not safe for concurrent mutation by multiple 649 // goroutines, but Read operations are. 650 type btree struct { 651 root *node 652 length int 653 } 654 655 // Reset removes all items from the btree. In doing so, it allows memory 656 // held by the btree to be recycled. Failure to call this method before 657 // letting a btree be GCed is safe in that it won't cause a memory leak, 658 // but it will prevent btree nodes from being efficiently re-used. 659 func (t *btree) Reset() { 660 if t.root != nil { 661 t.root.decRef(true /* recursive */) 662 t.root = nil 663 } 664 t.length = 0 665 } 666 667 // Clone clones the btree, lazily. It does so in constant time. 668 func (t *btree) Clone() btree { 669 c := *t 670 if c.root != nil { 671 // Incrementing the reference count on the root node is sufficient to 672 // ensure that no node in the cloned tree can be mutated by an actor 673 // holding a reference to the original tree and vice versa. This 674 // property is upheld because the root node in the receiver btree and 675 // the returned btree will both necessarily have a reference count of at 676 // least 2 when this method returns. All tree mutations recursively 677 // acquire mutable node references (see mut) as they traverse down the 678 // tree. The act of acquiring a mutable node reference performs a clone 679 // if a node's reference count is greater than one. Cloning a node (see 680 // clone) increases the reference count on each of its children, 681 // ensuring that they have a reference count of at least 2. This, in 682 // turn, ensures that any of the child nodes that are modified will also 683 // be copied-on-write, recursively ensuring the immutability property 684 // over the entire tree. 685 c.root.incRef() 686 } 687 return c 688 } 689 690 // Delete removes a item equal to the passed in item from the tree. 691 func (t *btree) Delete(item *example) { 692 if t.root == nil || t.root.count == 0 { 693 return 694 } 695 if out, _ := mut(&t.root).remove(item); out != nilT { 696 t.length-- 697 } 698 if t.root.count == 0 { 699 old := t.root 700 if t.root.leaf { 701 t.root = nil 702 } else { 703 t.root = t.root.children[0] 704 } 705 old.decRef(false /* recursive */) 706 } 707 } 708 709 // Set adds the given item to the tree. If a item in the tree already 710 // equals the given one, it is replaced with the new item. 711 func (t *btree) Set(item *example) { 712 if t.root == nil { 713 t.root = newLeafNode() 714 } else if t.root.count >= maxItems { 715 splitLa, splitNode := mut(&t.root).split(maxItems / 2) 716 newRoot := newNode() 717 newRoot.count = 1 718 newRoot.items[0] = splitLa 719 newRoot.children[0] = t.root 720 newRoot.children[1] = splitNode 721 newRoot.max = newRoot.findUpperBound() 722 t.root = newRoot 723 } 724 if replaced, _ := mut(&t.root).insert(item); !replaced { 725 t.length++ 726 } 727 } 728 729 // MakeIter returns a new iterator object. It is not safe to continue using an 730 // iterator after modifications are made to the tree. If modifications are made, 731 // create a new iterator. 732 func (t *btree) MakeIter() iterator { 733 return iterator{r: t.root, pos: -1} 734 } 735 736 // Height returns the height of the tree. 737 func (t *btree) Height() int { 738 if t.root == nil { 739 return 0 740 } 741 h := 1 742 n := t.root 743 for !n.leaf { 744 n = n.children[0] 745 h++ 746 } 747 return h 748 } 749 750 // Len returns the number of items currently in the tree. 751 func (t *btree) Len() int { 752 return t.length 753 } 754 755 // String returns a string description of the tree. The format is 756 // similar to the https://en.wikipedia.org/wiki/Newick_format. 757 func (t *btree) String() string { 758 if t.length == 0 { 759 return ";" 760 } 761 var b strings.Builder 762 t.root.writeString(&b) 763 return b.String() 764 } 765 766 func (n *node) writeString(b *strings.Builder) { 767 if n.leaf { 768 for i := int16(0); i < n.count; i++ { 769 if i != 0 { 770 b.WriteString(",") 771 } 772 b.WriteString(n.items[i].String()) 773 } 774 return 775 } 776 for i := int16(0); i <= n.count; i++ { 777 b.WriteString("(") 778 n.children[i].writeString(b) 779 b.WriteString(")") 780 if i < n.count { 781 b.WriteString(n.items[i].String()) 782 } 783 } 784 } 785 786 // iterStack represents a stack of (node, pos) tuples, which captures 787 // iteration state as an iterator descends a btree. 788 type iterStack struct { 789 a iterStackArr 790 aLen int16 // -1 when using s 791 s []iterFrame 792 } 793 794 // Used to avoid allocations for stacks below a certain size. 795 type iterStackArr [3]iterFrame 796 797 type iterFrame struct { 798 n *node 799 pos int16 800 } 801 802 func (is *iterStack) push(f iterFrame) { 803 if is.aLen == -1 { 804 is.s = append(is.s, f) 805 } else if int(is.aLen) == len(is.a) { 806 is.s = make([]iterFrame, int(is.aLen)+1, 2*int(is.aLen)) 807 copy(is.s, is.a[:]) 808 is.s[int(is.aLen)] = f 809 is.aLen = -1 810 } else { 811 is.a[is.aLen] = f 812 is.aLen++ 813 } 814 } 815 816 func (is *iterStack) pop() iterFrame { 817 if is.aLen == -1 { 818 f := is.s[len(is.s)-1] 819 is.s = is.s[:len(is.s)-1] 820 return f 821 } 822 is.aLen-- 823 return is.a[is.aLen] 824 } 825 826 func (is *iterStack) len() int { 827 if is.aLen == -1 { 828 return len(is.s) 829 } 830 return int(is.aLen) 831 } 832 833 func (is *iterStack) reset() { 834 if is.aLen == -1 { 835 is.s = is.s[:0] 836 } else { 837 is.aLen = 0 838 } 839 } 840 841 // iterator is responsible for search and traversal within a btree. 842 type iterator struct { 843 r *node 844 n *node 845 pos int16 846 s iterStack 847 o overlapScan 848 } 849 850 func (i *iterator) reset() { 851 i.n = i.r 852 i.pos = -1 853 i.s.reset() 854 i.o = overlapScan{} 855 } 856 857 func (i *iterator) descend(n *node, pos int16) { 858 i.s.push(iterFrame{n: n, pos: pos}) 859 i.n = n.children[pos] 860 i.pos = 0 861 } 862 863 // ascend ascends up to the current node's parent and resets the position 864 // to the one previously set for this parent node. 865 func (i *iterator) ascend() { 866 f := i.s.pop() 867 i.n = f.n 868 i.pos = f.pos 869 } 870 871 // SeekGE seeks to the first item greater-than or equal to the provided 872 // item. 873 func (i *iterator) SeekGE(item *example) { 874 i.reset() 875 if i.n == nil { 876 return 877 } 878 for { 879 pos, found := i.n.find(item) 880 i.pos = int16(pos) 881 if found { 882 return 883 } 884 if i.n.leaf { 885 if i.pos == i.n.count { 886 i.Next() 887 } 888 return 889 } 890 i.descend(i.n, i.pos) 891 } 892 } 893 894 // SeekLT seeks to the first item less-than the provided item. 895 func (i *iterator) SeekLT(item *example) { 896 i.reset() 897 if i.n == nil { 898 return 899 } 900 for { 901 pos, found := i.n.find(item) 902 i.pos = int16(pos) 903 if found || i.n.leaf { 904 i.Prev() 905 return 906 } 907 i.descend(i.n, i.pos) 908 } 909 } 910 911 // First seeks to the first item in the btree. 912 func (i *iterator) First() { 913 i.reset() 914 if i.n == nil { 915 return 916 } 917 for !i.n.leaf { 918 i.descend(i.n, 0) 919 } 920 i.pos = 0 921 } 922 923 // Last seeks to the last item in the btree. 924 func (i *iterator) Last() { 925 i.reset() 926 if i.n == nil { 927 return 928 } 929 for !i.n.leaf { 930 i.descend(i.n, i.n.count) 931 } 932 i.pos = i.n.count - 1 933 } 934 935 // Next positions the iterator to the item immediately following 936 // its current position. 937 func (i *iterator) Next() { 938 if i.n == nil { 939 return 940 } 941 942 if i.n.leaf { 943 i.pos++ 944 if i.pos < i.n.count { 945 return 946 } 947 for i.s.len() > 0 && i.pos >= i.n.count { 948 i.ascend() 949 } 950 return 951 } 952 953 i.descend(i.n, i.pos+1) 954 for !i.n.leaf { 955 i.descend(i.n, 0) 956 } 957 i.pos = 0 958 } 959 960 // Prev positions the iterator to the item immediately preceding 961 // its current position. 962 func (i *iterator) Prev() { 963 if i.n == nil { 964 return 965 } 966 967 if i.n.leaf { 968 i.pos-- 969 if i.pos >= 0 { 970 return 971 } 972 for i.s.len() > 0 && i.pos < 0 { 973 i.ascend() 974 i.pos-- 975 } 976 return 977 } 978 979 i.descend(i.n, i.pos) 980 for !i.n.leaf { 981 i.descend(i.n, i.n.count) 982 } 983 i.pos = i.n.count - 1 984 } 985 986 // Valid returns whether the iterator is positioned at a valid position. 987 func (i *iterator) Valid() bool { 988 return i.pos >= 0 && i.pos < i.n.count 989 } 990 991 // Cur returns the item at the iterator's current position. It is illegal 992 // to call Cur if the iterator is not valid. 993 func (i *iterator) Cur() *example { 994 return i.n.items[i.pos] 995 } 996 997 // An overlap scan is a scan over all items that overlap with the provided 998 // item in order of the overlapping items' start keys. The goal of the scan 999 // is to minimize the number of key comparisons performed in total. The 1000 // algorithm operates based on the following two invariants maintained by 1001 // augmented interval btree: 1002 // 1. all items are sorted in the btree based on their start key. 1003 // 2. all btree nodes maintain the upper bound end key of all items 1004 // in their subtree. 1005 // 1006 // The scan algorithm starts in "unconstrained minimum" and "unconstrained 1007 // maximum" states. To enter a "constrained minimum" state, the scan must reach 1008 // items in the tree with start keys above the search range's start key. 1009 // Because items in the tree are sorted by start key, once the scan enters the 1010 // "constrained minimum" state it will remain there. To enter a "constrained 1011 // maximum" state, the scan must determine the first child btree node in a given 1012 // subtree that can have items with start keys above the search range's end 1013 // key. The scan then remains in the "constrained maximum" state until it 1014 // traverse into this child node, at which point it moves to the "unconstrained 1015 // maximum" state again. 1016 // 1017 // The scan algorithm works like a standard btree forward scan with the 1018 // following augmentations: 1019 // 1. before tranversing the tree, the scan performs a binary search on the 1020 // root node's items to determine a "soft" lower-bound constraint position 1021 // and a "hard" upper-bound constraint position in the root's children. 1022 // 2. when tranversing into a child node in the lower or upper bound constraint 1023 // position, the constraint is refined by searching the child's items. 1024 // 3. the initial traversal down the tree follows the left-most children 1025 // whose upper bound end keys are equal to or greater than the start key 1026 // of the search range. The children followed will be equal to or less 1027 // than the soft lower bound constraint. 1028 // 4. once the initial tranversal completes and the scan is in the left-most 1029 // btree node whose upper bound overlaps the search range, key comparisons 1030 // must be performed with each item in the tree. This is necessary because 1031 // any of these items may have end keys that cause them to overlap with the 1032 // search range. 1033 // 5. once the scan reaches the lower bound constraint position (the first item 1034 // with a start key equal to or greater than the search range's start key), 1035 // it can begin scaning without performing key comparisons. This is allowed 1036 // because all items from this point forward will have end keys that are 1037 // greater than the search range's start key. 1038 // 6. once the scan reaches the upper bound constraint position, it terminates. 1039 // It does so because the item at this position is the first item with a 1040 // start key larger than the search range's end key. 1041 type overlapScan struct { 1042 // The "soft" lower-bound constraint. 1043 constrMinN *node 1044 constrMinPos int16 1045 constrMinReached bool 1046 1047 // The "hard" upper-bound constraint. 1048 constrMaxN *node 1049 constrMaxPos int16 1050 } 1051 1052 // FirstOverlap seeks to the first item in the btree that overlaps with the 1053 // provided search item. 1054 func (i *iterator) FirstOverlap(item *example) { 1055 i.reset() 1056 if i.n == nil { 1057 return 1058 } 1059 i.pos = 0 1060 i.o = overlapScan{} 1061 i.constrainMinSearchBounds(item) 1062 i.constrainMaxSearchBounds(item) 1063 i.findNextOverlap(item) 1064 } 1065 1066 // NextOverlap positions the iterator to the item immediately following 1067 // its current position that overlaps with the search item. 1068 func (i *iterator) NextOverlap(item *example) { 1069 if i.n == nil { 1070 return 1071 } 1072 i.pos++ 1073 i.findNextOverlap(item) 1074 } 1075 1076 func (i *iterator) constrainMinSearchBounds(item *example) { 1077 k := item.Key() 1078 j := sort.Search(int(i.n.count), func(j int) bool { 1079 return bytes.Compare(k, i.n.items[j].Key()) <= 0 1080 }) 1081 i.o.constrMinN = i.n 1082 i.o.constrMinPos = int16(j) 1083 } 1084 1085 func (i *iterator) constrainMaxSearchBounds(item *example) { 1086 up := upperBound(item) 1087 j := sort.Search(int(i.n.count), func(j int) bool { 1088 return !up.contains(i.n.items[j]) 1089 }) 1090 i.o.constrMaxN = i.n 1091 i.o.constrMaxPos = int16(j) 1092 } 1093 1094 func (i *iterator) findNextOverlap(item *example) { 1095 for { 1096 if i.pos > i.n.count { 1097 // Iterate up tree. 1098 i.ascend() 1099 } else if !i.n.leaf { 1100 // Iterate down tree. 1101 if i.o.constrMinReached || i.n.children[i.pos].max.contains(item) { 1102 par := i.n 1103 pos := i.pos 1104 i.descend(par, pos) 1105 1106 // Refine the constraint bounds, if necessary. 1107 if par == i.o.constrMinN && pos == i.o.constrMinPos { 1108 i.constrainMinSearchBounds(item) 1109 } 1110 if par == i.o.constrMaxN && pos == i.o.constrMaxPos { 1111 i.constrainMaxSearchBounds(item) 1112 } 1113 continue 1114 } 1115 } 1116 1117 // Check search bounds. 1118 if i.n == i.o.constrMaxN && i.pos == i.o.constrMaxPos { 1119 // Invalid. Past possible overlaps. 1120 i.pos = i.n.count 1121 return 1122 } 1123 if i.n == i.o.constrMinN && i.pos == i.o.constrMinPos { 1124 // The scan reached the soft lower-bound constraint. 1125 i.o.constrMinReached = true 1126 } 1127 1128 // Iterate across node. 1129 if i.pos < i.n.count { 1130 // Check for overlapping item. 1131 if i.o.constrMinReached { 1132 // Fast-path to avoid span comparison. i.o.constrMinReached 1133 // tells us that all items have end keys above our search 1134 // span's start key. 1135 return 1136 } 1137 if upperBound(i.n.items[i.pos]).contains(item) { 1138 return 1139 } 1140 } 1141 i.pos++ 1142 } 1143 }