btree.go

Documentation: cloud.google.com/go/internal/btree

     1  // Copyright 2014 Google LLC
     2  // Modified 2018 by Jonathan Amsterdam (jbamsterdam@gmail.com)
     3  //
     4  // Licensed under the Apache License, Version 2.0 (the "License");
     5  // you may not use this file except in compliance with the License.
     6  // You may obtain a copy of the License at
     7  //
     8  //     http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  // Package btree implements in-memory B-Trees of arbitrary degree.
    17  //
    18  // This implementation is based on google/btree (http://github.com/google/btree), and
    19  // much of the code is taken from there. But the API has been changed significantly,
    20  // particularly around iteration, and support for indexing by position has been
    21  // added.
    22  //
    23  // btree implements an in-memory B-Tree for use as an ordered data structure.
    24  // It is not meant for persistent storage solutions.
    25  //
    26  // It has a flatter structure than an equivalent red-black or other binary tree,
    27  // which in some cases yields better memory usage and/or performance.
    28  // See some discussion on the matter here:
    29  //
    30  //	http://google-opensource.blogspot.com/2013/01/c-containers-that-save-memory-and-time.html
    31  //
    32  // Note, though, that this project is in no way related to the C++ B-Tree
    33  // implementation written about there.
    34  //
    35  // Within this tree, each node contains a slice of items and a (possibly nil)
    36  // slice of children.  For basic numeric values or raw structs, this can cause
    37  // efficiency differences when compared to equivalent C++ template code that
    38  // stores values in arrays within the node:
    39  //   - Due to the overhead of storing values as interfaces (each
    40  //     value needs to be stored as the value itself, then 2 words for the
    41  //     interface pointing to that value and its type), resulting in higher
    42  //     memory use.
    43  //   - Since interfaces can point to values anywhere in memory, values are
    44  //     most likely not stored in contiguous blocks, resulting in a higher
    45  //     number of cache misses.
    46  //
    47  // These issues don't tend to matter, though, when working with strings or other
    48  // heap-allocated structures, since C++-equivalent structures also must store
    49  // pointers and also distribute their values across the heap.
    50  package btree
    51  
    52  import (
    53  	"sort"
    54  	"sync"
    55  )
    56  
    57  // Key represents a key into the tree.
    58  type Key interface{}
    59  
    60  // Value represents a value in the tree.
    61  type Value interface{}
    62  
    63  // item is a key-value pair.
    64  type item struct {
    65  	key   Key
    66  	value Value
    67  }
    68  
    69  type lessFunc func(interface{}, interface{}) bool
    70  
    71  // New creates a new B-Tree with the given degree and comparison function.
    72  //
    73  // New(2, less), for example, will create a 2-3-4 tree (each node contains 1-3 items
    74  // and 2-4 children).
    75  //
    76  // The less function tests whether the current item is less than the given argument.
    77  // It must provide a strict weak ordering.
    78  // If !less(a, b) && !less(b, a), we treat this to mean a == b (i.e. the tree
    79  // can hold only one of a or b).
    80  func New(degree int, less func(interface{}, interface{}) bool) *BTree {
    81  	if degree <= 1 {
    82  		panic("bad degree")
    83  	}
    84  	return &BTree{
    85  		degree: degree,
    86  		less:   less,
    87  		cow:    &copyOnWriteContext{},
    88  	}
    89  }
    90  
    91  // items stores items in a node.
    92  type items []item
    93  
    94  // insertAt inserts a value into the given index, pushing all subsequent values
    95  // forward.
    96  func (s *items) insertAt(index int, m item) {
    97  	*s = append(*s, item{})
    98  	if index < len(*s) {
    99  		copy((*s)[index+1:], (*s)[index:])
   100  	}
   101  	(*s)[index] = m
   102  }
   103  
   104  // removeAt removes a value at a given index, pulling all subsequent values
   105  // back.
   106  func (s *items) removeAt(index int) item {
   107  	m := (*s)[index]
   108  	copy((*s)[index:], (*s)[index+1:])
   109  	(*s)[len(*s)-1] = item{}
   110  	*s = (*s)[:len(*s)-1]
   111  	return m
   112  }
   113  
   114  // pop removes and returns the last element in the list.
   115  func (s *items) pop() item {
   116  	index := len(*s) - 1
   117  	out := (*s)[index]
   118  	(*s)[index] = item{}
   119  	*s = (*s)[:index]
   120  	return out
   121  }
   122  
   123  var nilItems = make(items, 16)
   124  
   125  // truncate truncates this instance at index so that it contains only the
   126  // first index items. index must be less than or equal to length.
   127  func (s *items) truncate(index int) {
   128  	var toClear items
   129  	*s, toClear = (*s)[:index], (*s)[index:]
   130  	for len(toClear) > 0 {
   131  		toClear = toClear[copy(toClear, nilItems):]
   132  	}
   133  }
   134  
   135  // find returns the index where an item with key should be inserted into this
   136  // list.  'found' is true if the item already exists in the list at the given
   137  // index.
   138  func (s items) find(k Key, less lessFunc) (index int, found bool) {
   139  	i := sort.Search(len(s), func(i int) bool { return less(k, s[i].key) })
   140  	// i is the smallest index of s for which k.Less(s[i].Key), or len(s).
   141  	if i > 0 && !less(s[i-1].key, k) {
   142  		return i - 1, true
   143  	}
   144  	return i, false
   145  }
   146  
   147  // children stores child nodes in a node.
   148  type children []*node
   149  
   150  // insertAt inserts a value into the given index, pushing all subsequent values
   151  // forward.
   152  func (s *children) insertAt(index int, n *node) {
   153  	*s = append(*s, nil)
   154  	if index < len(*s) {
   155  		copy((*s)[index+1:], (*s)[index:])
   156  	}
   157  	(*s)[index] = n
   158  }
   159  
   160  // removeAt removes a value at a given index, pulling all subsequent values
   161  // back.
   162  func (s *children) removeAt(index int) *node {
   163  	n := (*s)[index]
   164  	copy((*s)[index:], (*s)[index+1:])
   165  	(*s)[len(*s)-1] = nil
   166  	*s = (*s)[:len(*s)-1]
   167  	return n
   168  }
   169  
   170  // pop removes and returns the last element in the list.
   171  func (s *children) pop() (out *node) {
   172  	index := len(*s) - 1
   173  	out = (*s)[index]
   174  	(*s)[index] = nil
   175  	*s = (*s)[:index]
   176  	return
   177  }
   178  
   179  var nilChildren = make(children, 16)
   180  
   181  // truncate truncates this instance at index so that it contains only the
   182  // first index children. index must be less than or equal to length.
   183  func (s *children) truncate(index int) {
   184  	var toClear children
   185  	*s, toClear = (*s)[:index], (*s)[index:]
   186  	for len(toClear) > 0 {
   187  		toClear = toClear[copy(toClear, nilChildren):]
   188  	}
   189  }
   190  
   191  // node is an internal node in a tree.
   192  //
   193  // It must at all times maintain the invariant that either
   194  //   - len(children) == 0, len(items) unconstrained
   195  //   - len(children) == len(items) + 1
   196  type node struct {
   197  	items    items
   198  	children children
   199  	size     int // number of items in the subtree: len(items) + sum over i of children[i].size
   200  	cow      *copyOnWriteContext
   201  }
   202  
   203  func (n *node) computeSize() int {
   204  	sz := len(n.items)
   205  	for _, c := range n.children {
   206  		sz += c.size
   207  	}
   208  	return sz
   209  }
   210  
   211  func (n *node) mutableFor(cow *copyOnWriteContext) *node {
   212  	if n.cow == cow {
   213  		return n
   214  	}
   215  	out := cow.newNode()
   216  	if cap(out.items) >= len(n.items) {
   217  		out.items = out.items[:len(n.items)]
   218  	} else {
   219  		out.items = make(items, len(n.items), cap(n.items))
   220  	}
   221  	copy(out.items, n.items)
   222  	// Copy children
   223  	if cap(out.children) >= len(n.children) {
   224  		out.children = out.children[:len(n.children)]
   225  	} else {
   226  		out.children = make(children, len(n.children), cap(n.children))
   227  	}
   228  	copy(out.children, n.children)
   229  	out.size = n.size
   230  	return out
   231  }
   232  
   233  func (n *node) mutableChild(i int) *node {
   234  	c := n.children[i].mutableFor(n.cow)
   235  	n.children[i] = c
   236  	return c
   237  }
   238  
   239  // split splits the given node at the given index.  The current node shrinks,
   240  // and this function returns the item that existed at that index and a new node
   241  // containing all items/children after it.
   242  func (n *node) split(i int) (item, *node) {
   243  	item := n.items[i]
   244  	next := n.cow.newNode()
   245  	next.items = append(next.items, n.items[i+1:]...)
   246  	n.items.truncate(i)
   247  	if len(n.children) > 0 {
   248  		next.children = append(next.children, n.children[i+1:]...)
   249  		n.children.truncate(i + 1)
   250  	}
   251  	n.size = n.computeSize()
   252  	next.size = next.computeSize()
   253  	return item, next
   254  }
   255  
   256  // maybeSplitChild checks if a child should be split, and if so splits it.
   257  // Returns whether or not a split occurred.
   258  func (n *node) maybeSplitChild(i, maxItems int) bool {
   259  	if len(n.children[i].items) < maxItems {
   260  		return false
   261  	}
   262  	first := n.mutableChild(i)
   263  	item, second := first.split(maxItems / 2)
   264  	n.items.insertAt(i, item)
   265  	n.children.insertAt(i+1, second)
   266  	// The size of n doesn't change.
   267  	return true
   268  }
   269  
   270  // insert inserts an item into the subtree rooted at this node, making sure
   271  // no nodes in the subtree exceed maxItems items.  Should an equivalent item be
   272  // be found/replaced by insert, its value will be returned.
   273  //
   274  // If computeIndex is true, the third return value is the index of the value with respect to n.
   275  func (n *node) insert(m item, maxItems int, less lessFunc, computeIndex bool) (old Value, present bool, idx int) {
   276  	i, found := n.items.find(m.key, less)
   277  	if found {
   278  		out := n.items[i]
   279  		n.items[i] = m
   280  		if computeIndex {
   281  			idx = n.itemIndex(i)
   282  		}
   283  		return out.value, true, idx
   284  	}
   285  	if len(n.children) == 0 {
   286  		n.items.insertAt(i, m)
   287  		n.size++
   288  		return old, false, i
   289  	}
   290  	if n.maybeSplitChild(i, maxItems) {
   291  		inTree := n.items[i]
   292  		switch {
   293  		case less(m.key, inTree.key):
   294  			// no change, we want first split node
   295  		case less(inTree.key, m.key):
   296  			i++ // we want second split node
   297  		default:
   298  			out := n.items[i]
   299  			n.items[i] = m
   300  			if computeIndex {
   301  				idx = n.itemIndex(i)
   302  			}
   303  			return out.value, true, idx
   304  		}
   305  	}
   306  	old, present, idx = n.mutableChild(i).insert(m, maxItems, less, computeIndex)
   307  	if !present {
   308  		n.size++
   309  	}
   310  	if computeIndex {
   311  		idx += n.partialSize(i)
   312  	}
   313  	return old, present, idx
   314  }
   315  
   316  // get finds the given key in the subtree and returns the corresponding item, along with a boolean reporting
   317  // whether it was found.
   318  // If computeIndex is true, it also returns the index of the key relative to the node's subtree.
   319  func (n *node) get(k Key, computeIndex bool, less lessFunc) (item, bool, int) {
   320  	i, found := n.items.find(k, less)
   321  	if found {
   322  		return n.items[i], true, n.itemIndex(i)
   323  	}
   324  	if len(n.children) > 0 {
   325  		m, found, idx := n.children[i].get(k, computeIndex, less)
   326  		if computeIndex && found {
   327  			idx += n.partialSize(i)
   328  		}
   329  		return m, found, idx
   330  	}
   331  	return item{}, false, -1
   332  }
   333  
   334  // itemIndex returns the index w.r.t. n of the ith item in n.
   335  func (n *node) itemIndex(i int) int {
   336  	if len(n.children) == 0 {
   337  		return i
   338  	}
   339  	// Get the size of the node up to but not including the child to the right of
   340  	// item i. Subtract 1 because the index is 0-based.
   341  	return n.partialSize(i+1) - 1
   342  }
   343  
   344  // Returns the size of the non-leaf node up to but not including child i.
   345  func (n *node) partialSize(i int) int {
   346  	var sz int
   347  	for j, c := range n.children {
   348  		if j == i {
   349  			break
   350  		}
   351  		sz += c.size + 1
   352  	}
   353  	return sz
   354  }
   355  
   356  // cursorStackForKey returns a stack of cursors for the key, along with whether the key was found and the index.
   357  func (n *node) cursorStackForKey(k Key, cs cursorStack, less lessFunc) (cursorStack, bool, int) {
   358  	i, found := n.items.find(k, less)
   359  	cs.push(cursor{n, i})
   360  	idx := i
   361  	if found {
   362  		if len(n.children) > 0 {
   363  			idx = n.partialSize(i+1) - 1
   364  		}
   365  		return cs, true, idx
   366  	}
   367  	if len(n.children) > 0 {
   368  		cs, found, idx := n.children[i].cursorStackForKey(k, cs, less)
   369  		return cs, found, idx + n.partialSize(i)
   370  	}
   371  	return cs, false, idx
   372  }
   373  
   374  // at returns the item at the i'th position in the subtree rooted at n.
   375  // It assumes i is in range.
   376  func (n *node) at(i int) item {
   377  	if len(n.children) == 0 {
   378  		return n.items[i]
   379  	}
   380  	for j, c := range n.children {
   381  		if i < c.size {
   382  			return c.at(i)
   383  		}
   384  		i -= c.size
   385  		if i == 0 {
   386  			return n.items[j]
   387  		}
   388  		i--
   389  	}
   390  	panic("impossible")
   391  }
   392  
   393  // cursorStackForIndex returns a stack of cursors for the index.
   394  // It assumes i is in range.
   395  func (n *node) cursorStackForIndex(i int, cs cursorStack) cursorStack {
   396  	if len(n.children) == 0 {
   397  		return cs.push(cursor{n, i})
   398  	}
   399  	for j, c := range n.children {
   400  		if i < c.size {
   401  			return c.cursorStackForIndex(i, cs.push(cursor{n, j}))
   402  		}
   403  		i -= c.size
   404  		if i == 0 {
   405  			return cs.push(cursor{n, j})
   406  		}
   407  		i--
   408  	}
   409  	panic("impossible")
   410  }
   411  
   412  // toRemove details what item to remove in a node.remove call.
   413  type toRemove int
   414  
   415  const (
   416  	removeItem toRemove = iota // removes the given item
   417  	removeMin                  // removes smallest item in the subtree
   418  	removeMax                  // removes largest item in the subtree
   419  )
   420  
   421  // remove removes an item from the subtree rooted at this node.
   422  func (n *node) remove(key Key, minItems int, typ toRemove, less lessFunc) (item, bool) {
   423  	var i int
   424  	var found bool
   425  	switch typ {
   426  	case removeMax:
   427  		if len(n.children) == 0 {
   428  			n.size--
   429  			return n.items.pop(), true
   430  
   431  		}
   432  		i = len(n.items)
   433  	case removeMin:
   434  		if len(n.children) == 0 {
   435  			n.size--
   436  			return n.items.removeAt(0), true
   437  		}
   438  		i = 0
   439  	case removeItem:
   440  		i, found = n.items.find(key, less)
   441  		if len(n.children) == 0 {
   442  			if found {
   443  				n.size--
   444  				return n.items.removeAt(i), true
   445  			}
   446  			return item{}, false
   447  		}
   448  	default:
   449  		panic("invalid type")
   450  	}
   451  	// If we get to here, we have children.
   452  	if len(n.children[i].items) <= minItems {
   453  		return n.growChildAndRemove(i, key, minItems, typ, less)
   454  	}
   455  	child := n.mutableChild(i)
   456  	// Either we had enough items to begin with, or we've done some
   457  	// merging/stealing, because we've got enough now and we're ready to return
   458  	// stuff.
   459  	if found {
   460  		// The item exists at index 'i', and the child we've selected can give us a
   461  		// predecessor, since if we've gotten here it's got > minItems items in it.
   462  		out := n.items[i]
   463  		// We use our special-case 'remove' call with typ=maxItem to pull the
   464  		// predecessor of item i (the rightmost leaf of our immediate left child)
   465  		// and set it into where we pulled the item from.
   466  		n.items[i], _ = child.remove(nil, minItems, removeMax, less)
   467  		n.size--
   468  		return out, true
   469  	}
   470  	// Final recursive call.  Once we're here, we know that the item isn't in this
   471  	// node and that the child is big enough to remove from.
   472  	m, removed := child.remove(key, minItems, typ, less)
   473  	if removed {
   474  		n.size--
   475  	}
   476  	return m, removed
   477  }
   478  
   479  // growChildAndRemove grows child 'i' to make sure it's possible to remove an
   480  // item from it while keeping it at minItems, then calls remove to actually
   481  // remove it.
   482  //
   483  // Most documentation says we have to do two sets of special casing:
   484  //  1. item is in this node
   485  //  2. item is in child
   486  //
   487  // In both cases, we need to handle the two subcases:
   488  //
   489  //	A) node has enough values that it can spare one
   490  //	B) node doesn't have enough values
   491  //
   492  // For the latter, we have to check:
   493  //
   494  //	a) left sibling has node to spare
   495  //	b) right sibling has node to spare
   496  //	c) we must merge
   497  //
   498  // To simplify our code here, we handle cases #1 and #2 the same:
   499  // If a node doesn't have enough items, we make sure it does (using a,b,c).
   500  // We then simply redo our remove call, and the second time (regardless of
   501  // whether we're in case 1 or 2), we'll have enough items and can guarantee
   502  // that we hit case A.
   503  func (n *node) growChildAndRemove(i int, key Key, minItems int, typ toRemove, less lessFunc) (item, bool) {
   504  	if i > 0 && len(n.children[i-1].items) > minItems {
   505  		// Steal from left child
   506  		child := n.mutableChild(i)
   507  		stealFrom := n.mutableChild(i - 1)
   508  		stolenItem := stealFrom.items.pop()
   509  		stealFrom.size--
   510  		child.items.insertAt(0, n.items[i-1])
   511  		child.size++
   512  		n.items[i-1] = stolenItem
   513  		if len(stealFrom.children) > 0 {
   514  			c := stealFrom.children.pop()
   515  			stealFrom.size -= c.size
   516  			child.children.insertAt(0, c)
   517  			child.size += c.size
   518  		}
   519  	} else if i < len(n.items) && len(n.children[i+1].items) > minItems {
   520  		// steal from right child
   521  		child := n.mutableChild(i)
   522  		stealFrom := n.mutableChild(i + 1)
   523  		stolenItem := stealFrom.items.removeAt(0)
   524  		stealFrom.size--
   525  		child.items = append(child.items, n.items[i])
   526  		child.size++
   527  		n.items[i] = stolenItem
   528  		if len(stealFrom.children) > 0 {
   529  			c := stealFrom.children.removeAt(0)
   530  			stealFrom.size -= c.size
   531  			child.children = append(child.children, c)
   532  			child.size += c.size
   533  		}
   534  	} else {
   535  		if i >= len(n.items) {
   536  			i--
   537  		}
   538  		child := n.mutableChild(i)
   539  		// merge with right child
   540  		mergeItem := n.items.removeAt(i)
   541  		mergeChild := n.children.removeAt(i + 1)
   542  		child.items = append(child.items, mergeItem)
   543  		child.items = append(child.items, mergeChild.items...)
   544  		child.children = append(child.children, mergeChild.children...)
   545  		child.size = child.computeSize()
   546  		n.cow.freeNode(mergeChild)
   547  	}
   548  	return n.remove(key, minItems, typ, less)
   549  }
   550  
   551  // BTree is an implementation of a B-Tree.
   552  //
   553  // BTree stores item instances in an ordered structure, allowing easy insertion,
   554  // removal, and iteration.
   555  //
   556  // Write operations are not safe for concurrent mutation by multiple
   557  // goroutines, but Read operations are.
   558  type BTree struct {
   559  	degree int
   560  	less   lessFunc
   561  	root   *node
   562  	cow    *copyOnWriteContext
   563  }
   564  
   565  // copyOnWriteContext pointers determine node ownership. A tree with a cow
   566  // context equivalent to a node's cow context is allowed to modify that node.
   567  // A tree whose write context does not match a node's is not allowed to modify
   568  // it, and must create a new, writable copy (IE: it's a Clone).
   569  //
   570  // When doing any write operation, we maintain the invariant that the current
   571  // node's context is equal to the context of the tree that requested the write.
   572  // We do this by, before we descend into any node, creating a copy with the
   573  // correct context if the contexts don't match.
   574  //
   575  // Since the node we're currently visiting on any write has the requesting
   576  // tree's context, that node is modifiable in place.  Children of that node may
   577  // not share context, but before we descend into them, we'll make a mutable
   578  // copy.
   579  type copyOnWriteContext struct{ byte } // non-empty, because empty structs may have same addr
   580  
   581  // Clone clones the btree, lazily.  Clone should not be called concurrently,
   582  // but the original tree (t) and the new tree (t2) can be used concurrently
   583  // once the Clone call completes.
   584  //
   585  // The internal tree structure of b is marked read-only and shared between t and
   586  // t2.  Writes to both t and t2 use copy-on-write logic, creating new nodes
   587  // whenever one of b's original nodes would have been modified.  Read operations
   588  // should have no performance degredation.  Write operations for both t and t2
   589  // will initially experience minor slow-downs caused by additional allocs and
   590  // copies due to the aforementioned copy-on-write logic, but should converge to
   591  // the original performance characteristics of the original tree.
   592  func (t *BTree) Clone() *BTree {
   593  	// Create two entirely new copy-on-write contexts.
   594  	// This operation effectively creates three trees:
   595  	//   the original, shared nodes (old b.cow)
   596  	//   the new b.cow nodes
   597  	//   the new out.cow nodes
   598  	cow1, cow2 := *t.cow, *t.cow
   599  	out := *t
   600  	t.cow = &cow1
   601  	out.cow = &cow2
   602  	return &out
   603  }
   604  
   605  // maxItems returns the max number of items to allow per node.
   606  func (t *BTree) maxItems() int {
   607  	return t.degree*2 - 1
   608  }
   609  
   610  // minItems returns the min number of items to allow per node (ignored for the
   611  // root node).
   612  func (t *BTree) minItems() int {
   613  	return t.degree - 1
   614  }
   615  
   616  var nodePool = sync.Pool{New: func() interface{} { return new(node) }}
   617  
   618  func (c *copyOnWriteContext) newNode() *node {
   619  	n := nodePool.Get().(*node)
   620  	n.cow = c
   621  	return n
   622  }
   623  
   624  func (c *copyOnWriteContext) freeNode(n *node) {
   625  	if n.cow == c {
   626  		// clear to allow GC
   627  		n.items.truncate(0)
   628  		n.children.truncate(0)
   629  		n.cow = nil
   630  		nodePool.Put(n)
   631  	}
   632  }
   633  
   634  // Set sets the given key to the given value in the tree. If the key is present in
   635  // the tree, its value is changed and the old value is returned along with a second
   636  // return value of true. If the key is not in the tree, it is added, and the second
   637  // return value is false.
   638  func (t *BTree) Set(k Key, v Value) (old Value, present bool) {
   639  	old, present, _ = t.set(k, v, false)
   640  	return old, present
   641  }
   642  
   643  // SetWithIndex sets the given key to the given value in the tree, and returns the
   644  // index at which it was inserted.
   645  func (t *BTree) SetWithIndex(k Key, v Value) (old Value, present bool, index int) {
   646  	return t.set(k, v, true)
   647  }
   648  
   649  func (t *BTree) set(k Key, v Value, computeIndex bool) (old Value, present bool, idx int) {
   650  	if t.root == nil {
   651  		t.root = t.cow.newNode()
   652  		t.root.items = append(t.root.items, item{k, v})
   653  		t.root.size = 1
   654  		return old, false, 0
   655  	}
   656  	t.root = t.root.mutableFor(t.cow)
   657  	if len(t.root.items) >= t.maxItems() {
   658  		sz := t.root.size
   659  		item2, second := t.root.split(t.maxItems() / 2)
   660  		oldroot := t.root
   661  		t.root = t.cow.newNode()
   662  		t.root.items = append(t.root.items, item2)
   663  		t.root.children = append(t.root.children, oldroot, second)
   664  		t.root.size = sz
   665  	}
   666  
   667  	return t.root.insert(item{k, v}, t.maxItems(), t.less, computeIndex)
   668  }
   669  
   670  // Delete removes the item with the given key, returning its value. The second return value
   671  // reports whether the key was found.
   672  func (t *BTree) Delete(k Key) (Value, bool) {
   673  	m, removed := t.deleteItem(k, removeItem)
   674  	return m.value, removed
   675  }
   676  
   677  // DeleteMin removes the smallest item in the tree and returns its key and value.
   678  // If the tree is empty, it returns zero values.
   679  func (t *BTree) DeleteMin() (Key, Value) {
   680  	item, _ := t.deleteItem(nil, removeMin)
   681  	return item.key, item.value
   682  }
   683  
   684  // DeleteMax removes the largest item in the tree and returns its key and value.
   685  // If the tree is empty, it returns zero values.
   686  func (t *BTree) DeleteMax() (Key, Value) {
   687  	item, _ := t.deleteItem(nil, removeMax)
   688  	return item.key, item.value
   689  }
   690  
   691  func (t *BTree) deleteItem(key Key, typ toRemove) (item, bool) {
   692  	if t.root == nil || len(t.root.items) == 0 {
   693  		return item{}, false
   694  	}
   695  	t.root = t.root.mutableFor(t.cow)
   696  	out, removed := t.root.remove(key, t.minItems(), typ, t.less)
   697  	if len(t.root.items) == 0 && len(t.root.children) > 0 {
   698  		oldroot := t.root
   699  		t.root = t.root.children[0]
   700  		t.cow.freeNode(oldroot)
   701  	}
   702  	return out, removed
   703  }
   704  
   705  // Get returns the value for the given key in the tree, or the zero value if the
   706  // key is not in the tree.
   707  //
   708  // To distinguish a zero value from a key that is not present, use GetWithIndex.
   709  func (t *BTree) Get(k Key) Value {
   710  	var z Value
   711  	if t.root == nil {
   712  		return z
   713  	}
   714  	item, ok, _ := t.root.get(k, false, t.less)
   715  	if !ok {
   716  		return z
   717  	}
   718  	return item.value
   719  }
   720  
   721  // GetWithIndex returns the value and index for the given key in the tree, or the
   722  // zero value and -1 if the key is not in the tree.
   723  func (t *BTree) GetWithIndex(k Key) (Value, int) {
   724  	var z Value
   725  	if t.root == nil {
   726  		return z, -1
   727  	}
   728  	item, _, index := t.root.get(k, true, t.less)
   729  	return item.value, index
   730  }
   731  
   732  // At returns the key and value at index i. The minimum item has index 0.
   733  // If i is outside the range [0, t.Len()), At panics.
   734  func (t *BTree) At(i int) (Key, Value) {
   735  	if i < 0 || i >= t.Len() {
   736  		panic("btree: index out of range")
   737  	}
   738  	item := t.root.at(i)
   739  	return item.key, item.value
   740  }
   741  
   742  // Has reports whether the given key is in the tree.
   743  func (t *BTree) Has(k Key) bool {
   744  	if t.root == nil {
   745  		return false
   746  	}
   747  	_, ok, _ := t.root.get(k, false, t.less)
   748  	return ok
   749  }
   750  
   751  // Min returns the smallest key in the tree and its value. If the tree is empty, it
   752  // returns zero values.
   753  func (t *BTree) Min() (Key, Value) {
   754  	var k Key
   755  	var v Value
   756  	if t.root == nil {
   757  		return k, v
   758  	}
   759  	n := t.root
   760  	for len(n.children) > 0 {
   761  		n = n.children[0]
   762  	}
   763  	if len(n.items) == 0 {
   764  		return k, v
   765  	}
   766  	return n.items[0].key, n.items[0].value
   767  }
   768  
   769  // Max returns the largest key in the tree and its value. If the tree is empty, both
   770  // return values are zero values.
   771  func (t *BTree) Max() (Key, Value) {
   772  	var k Key
   773  	var v Value
   774  	if t.root == nil {
   775  		return k, v
   776  	}
   777  	n := t.root
   778  	for len(n.children) > 0 {
   779  		n = n.children[len(n.children)-1]
   780  	}
   781  	if len(n.items) == 0 {
   782  		return k, v
   783  	}
   784  	m := n.items[len(n.items)-1]
   785  	return m.key, m.value
   786  }
   787  
   788  // Len returns the number of items currently in the tree.
   789  func (t *BTree) Len() int {
   790  	if t.root == nil {
   791  		return 0
   792  	}
   793  	return t.root.size
   794  }
   795  
   796  // Before returns an iterator positioned just before k. After the first call to Next,
   797  // the Iterator will be at k, or at the key just greater than k if k is not in the tree.
   798  // Subsequent calls to Next will traverse the tree's items in ascending order.
   799  func (t *BTree) Before(k Key) *Iterator {
   800  	if t.root == nil {
   801  		return &Iterator{}
   802  	}
   803  	var cs cursorStack
   804  	cs, found, idx := t.root.cursorStackForKey(k, cs, t.less)
   805  	// If we found the key, the cursor stack is pointing to it. Since that is
   806  	// the first element we want, don't advance the iterator on the initial call to Next.
   807  	// If we haven't found the key, then the top of the cursor stack is either pointing at the
   808  	// item just after k, in which case we do not want to move the iterator; or the index
   809  	// is past the end of the items slice, in which case we do.
   810  	var stay bool
   811  	top := cs[len(cs)-1]
   812  	if found {
   813  		stay = true
   814  	} else if top.index < len(top.node.items) {
   815  		stay = true
   816  	} else {
   817  		idx--
   818  	}
   819  	return &Iterator{
   820  		cursors:    cs,
   821  		stay:       stay,
   822  		descending: false,
   823  		Index:      idx,
   824  	}
   825  }
   826  
   827  // After returns an iterator positioned just after k. After the first call to Next,
   828  // the Iterator will be at k, or at the key just less than k if k is not in the tree.
   829  // Subsequent calls to Next will traverse the tree's items in descending order.
   830  func (t *BTree) After(k Key) *Iterator {
   831  	if t.root == nil {
   832  		return &Iterator{}
   833  	}
   834  	var cs cursorStack
   835  	cs, found, idx := t.root.cursorStackForKey(k, cs, t.less)
   836  	// If we found the key, the cursor stack is pointing to it. Since that is
   837  	// the first element we want, don't advance the iterator on the initial call to Next.
   838  	// If we haven't found the key, the cursor stack is pointing just after the first item,
   839  	// so we do want to advance.
   840  	return &Iterator{
   841  		cursors:    cs,
   842  		stay:       found,
   843  		descending: true,
   844  		Index:      idx,
   845  	}
   846  }
   847  
   848  // BeforeIndex returns an iterator positioned just before the item with the given index.
   849  // The iterator will traverse the tree's items in ascending order.
   850  // If i is not in the range [0, tr.Len()], BeforeIndex panics.
   851  // Note that it is not an error to provide an index of tr.Len().
   852  func (t *BTree) BeforeIndex(i int) *Iterator {
   853  	return t.indexIterator(i, false)
   854  }
   855  
   856  // AfterIndex returns an iterator positioned just after the item with the given index.
   857  // The iterator will traverse the tree's items in descending order.
   858  // If i is not in the range [0, tr.Len()], AfterIndex panics.
   859  // Note that it is not an error to provide an index of tr.Len().
   860  func (t *BTree) AfterIndex(i int) *Iterator {
   861  	return t.indexIterator(i, true)
   862  }
   863  
   864  func (t *BTree) indexIterator(i int, descending bool) *Iterator {
   865  	if i < 0 || i > t.Len() {
   866  		panic("btree: index out of range")
   867  	}
   868  	if i == t.Len() {
   869  		return &Iterator{}
   870  	}
   871  	var cs cursorStack
   872  	return &Iterator{
   873  		cursors:    t.root.cursorStackForIndex(i, cs),
   874  		stay:       true,
   875  		descending: descending,
   876  		Index:      i,
   877  	}
   878  }
   879  
   880  // An Iterator supports traversing the items in the tree.
   881  type Iterator struct {
   882  	Key   Key
   883  	Value Value
   884  	// Index is the position of the item in the tree viewed as a sequence.
   885  	// The minimum item has index zero.
   886  	Index int
   887  
   888  	cursors    cursorStack // stack of nodes with indices; last element is the top
   889  	stay       bool        // don't do anything on the first call to Next.
   890  	descending bool        // traverse the items in descending order
   891  }
   892  
   893  // Next advances the Iterator to the next item in the tree. If Next returns true,
   894  // the Iterator's Key, Value and Index fields refer to the next item. If Next returns
   895  // false, there are no more items and the values of Key, Value and Index are undefined.
   896  //
   897  // If the tree is modified during iteration, the behavior is undefined.
   898  func (it *Iterator) Next() bool {
   899  	var more bool
   900  	switch {
   901  	case len(it.cursors) == 0:
   902  		more = false
   903  	case it.stay:
   904  		it.stay = false
   905  		more = true
   906  	case it.descending:
   907  		more = it.dec()
   908  	default:
   909  		more = it.inc()
   910  	}
   911  	if !more {
   912  		return false
   913  	}
   914  	top := it.cursors[len(it.cursors)-1]
   915  	item := top.node.items[top.index]
   916  	it.Key = item.key
   917  	it.Value = item.value
   918  	return true
   919  }
   920  
   921  // When inc returns true, the top cursor on the stack refers to the new current item.
   922  func (it *Iterator) inc() bool {
   923  	// Useful invariants for understanding this function:
   924  	// - Leaf nodes have zero children, and zero or more items.
   925  	// - Nonleaf nodes have one more child than item, and children[i] < items[i] < children[i+1].
   926  	// - The current item in the iterator is top.node.items[top.index].
   927  
   928  	it.Index++
   929  	// If we are at a non-leaf node, the current item is items[i], so
   930  	// now we want to continue with children[i+1], which must exist
   931  	// by the node invariant. We want the minimum item in that child's subtree.
   932  	top := it.cursors.incTop(1)
   933  	for len(top.node.children) > 0 {
   934  		top = cursor{top.node.children[top.index], 0}
   935  		it.cursors.push(top)
   936  	}
   937  	// Here, we are at a leaf node. top.index points to
   938  	// the new current item, if it's within the items slice.
   939  	for top.index >= len(top.node.items) {
   940  		// We've gone through everything in this node. Pop it off the stack.
   941  		it.cursors.pop()
   942  		// If the stack is now empty,we're past the last item in the tree.
   943  		if it.cursors.empty() {
   944  			return false
   945  		}
   946  		top = it.cursors.top()
   947  		// The new top's index points to a child, which we've just finished
   948  		// exploring. The next item is the one at the same index in the items slice.
   949  	}
   950  	// Here, the top cursor on the stack points to the new current item.
   951  	return true
   952  }
   953  
   954  func (it *Iterator) dec() bool {
   955  	// See the invariants for inc, above.
   956  	it.Index--
   957  	top := it.cursors.top()
   958  	// If we are at a non-leaf node, the current item is items[i], so
   959  	// now we want to continue with children[i]. We want the maximum item in that child's subtree.
   960  	for len(top.node.children) > 0 {
   961  		c := top.node.children[top.index]
   962  		top = cursor{c, len(c.items)}
   963  		it.cursors.push(top)
   964  	}
   965  	top = it.cursors.incTop(-1)
   966  	// Here, we are at a leaf node. top.index points to
   967  	// the new current item, if it's within the items slice.
   968  	for top.index < 0 {
   969  		// We've gone through everything in this node. Pop it off the stack.
   970  		it.cursors.pop()
   971  		// If the stack is now empty,we're past the last item in the tree.
   972  		if it.cursors.empty() {
   973  			return false
   974  		}
   975  		// The new top's index points to a child, which we've just finished
   976  		// exploring. That child is to the right of the item we want to advance to,
   977  		// so decrement the index.
   978  		top = it.cursors.incTop(-1)
   979  	}
   980  	return true
   981  }
   982  
   983  // A cursor is effectively a pointer into a node. A stack of cursors identifies an item in the tree,
   984  // and makes it possible to move to the next or previous item efficiently.
   985  //
   986  // If the cursor is on the top of the stack, its index points into the node's items slice, selecting
   987  // the current item. Otherwise, the index points into the children slice and identifies the child
   988  // that is next in the stack.
   989  type cursor struct {
   990  	node  *node
   991  	index int
   992  }
   993  
   994  // A cursorStack is a stack of cursors, representing a path of nodes from the root of the tree.
   995  type cursorStack []cursor
   996  
   997  func (s *cursorStack) push(c cursor) cursorStack {
   998  	*s = append(*s, c)
   999  	return *s
  1000  }
  1001  
  1002  func (s *cursorStack) pop() cursor {
  1003  	last := len(*s) - 1
  1004  	t := (*s)[last]
  1005  	*s = (*s)[:last]
  1006  	return t
  1007  }
  1008  
  1009  func (s *cursorStack) top() cursor {
  1010  	return (*s)[len(*s)-1]
  1011  }
  1012  
  1013  func (s *cursorStack) empty() bool {
  1014  	return len(*s) == 0
  1015  }
  1016  
  1017  // incTop increments top's index by n and returns it.
  1018  func (s *cursorStack) incTop(n int) cursor {
  1019  	(*s)[len(*s)-1].index += n // Don't call top: modify the original, not a copy.
  1020  	return s.top()
  1021  }
  1022
View as plain text