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  }