btree.go

Documentation: github.com/google/btree

     1  // Copyright 2014 Google Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  //go:build !go1.18
    16  // +build !go1.18
    17  
    18  // Package btree implements in-memory B-Trees of arbitrary degree.
    19  //
    20  // btree implements an in-memory B-Tree for use as an ordered data structure.
    21  // It is not meant for persistent storage solutions.
    22  //
    23  // It has a flatter structure than an equivalent red-black or other binary tree,
    24  // which in some cases yields better memory usage and/or performance.
    25  // See some discussion on the matter here:
    26  //   http://google-opensource.blogspot.com/2013/01/c-containers-that-save-memory-and-time.html
    27  // Note, though, that this project is in no way related to the C++ B-Tree
    28  // implementation written about there.
    29  //
    30  // Within this tree, each node contains a slice of items and a (possibly nil)
    31  // slice of children.  For basic numeric values or raw structs, this can cause
    32  // efficiency differences when compared to equivalent C++ template code that
    33  // stores values in arrays within the node:
    34  //   * Due to the overhead of storing values as interfaces (each
    35  //     value needs to be stored as the value itself, then 2 words for the
    36  //     interface pointing to that value and its type), resulting in higher
    37  //     memory use.
    38  //   * Since interfaces can point to values anywhere in memory, values are
    39  //     most likely not stored in contiguous blocks, resulting in a higher
    40  //     number of cache misses.
    41  // These issues don't tend to matter, though, when working with strings or other
    42  // heap-allocated structures, since C++-equivalent structures also must store
    43  // pointers and also distribute their values across the heap.
    44  //
    45  // This implementation is designed to be a drop-in replacement to gollrb.LLRB
    46  // trees, (http://github.com/petar/gollrb), an excellent and probably the most
    47  // widely used ordered tree implementation in the Go ecosystem currently.
    48  // Its functions, therefore, exactly mirror those of
    49  // llrb.LLRB where possible.  Unlike gollrb, though, we currently don't
    50  // support storing multiple equivalent values.
    51  package btree
    52  
    53  import (
    54  	"fmt"
    55  	"io"
    56  	"sort"
    57  	"strings"
    58  	"sync"
    59  )
    60  
    61  // Item represents a single object in the tree.
    62  type Item interface {
    63  	// Less tests whether the current item is less than the given argument.
    64  	//
    65  	// This must provide a strict weak ordering.
    66  	// If !a.Less(b) && !b.Less(a), we treat this to mean a == b (i.e. we can only
    67  	// hold one of either a or b in the tree).
    68  	Less(than Item) bool
    69  }
    70  
    71  const (
    72  	DefaultFreeListSize = 32
    73  )
    74  
    75  var (
    76  	nilItems    = make(items, 16)
    77  	nilChildren = make(children, 16)
    78  )
    79  
    80  // FreeList represents a free list of btree nodes. By default each
    81  // BTree has its own FreeList, but multiple BTrees can share the same
    82  // FreeList.
    83  // Two Btrees using the same freelist are safe for concurrent write access.
    84  type FreeList struct {
    85  	mu       sync.Mutex
    86  	freelist []*node
    87  }
    88  
    89  // NewFreeList creates a new free list.
    90  // size is the maximum size of the returned free list.
    91  func NewFreeList(size int) *FreeList {
    92  	return &FreeList{freelist: make([]*node, 0, size)}
    93  }
    94  
    95  func (f *FreeList) newNode() (n *node) {
    96  	f.mu.Lock()
    97  	index := len(f.freelist) - 1
    98  	if index < 0 {
    99  		f.mu.Unlock()
   100  		return new(node)
   101  	}
   102  	n = f.freelist[index]
   103  	f.freelist[index] = nil
   104  	f.freelist = f.freelist[:index]
   105  	f.mu.Unlock()
   106  	return
   107  }
   108  
   109  // freeNode adds the given node to the list, returning true if it was added
   110  // and false if it was discarded.
   111  func (f *FreeList) freeNode(n *node) (out bool) {
   112  	f.mu.Lock()
   113  	if len(f.freelist) < cap(f.freelist) {
   114  		f.freelist = append(f.freelist, n)
   115  		out = true
   116  	}
   117  	f.mu.Unlock()
   118  	return
   119  }
   120  
   121  // ItemIterator allows callers of Ascend* to iterate in-order over portions of
   122  // the tree.  When this function returns false, iteration will stop and the
   123  // associated Ascend* function will immediately return.
   124  type ItemIterator func(i Item) bool
   125  
   126  // New creates a new B-Tree with the given degree.
   127  //
   128  // New(2), for example, will create a 2-3-4 tree (each node contains 1-3 items
   129  // and 2-4 children).
   130  func New(degree int) *BTree {
   131  	return NewWithFreeList(degree, NewFreeList(DefaultFreeListSize))
   132  }
   133  
   134  // NewWithFreeList creates a new B-Tree that uses the given node free list.
   135  func NewWithFreeList(degree int, f *FreeList) *BTree {
   136  	if degree <= 1 {
   137  		panic("bad degree")
   138  	}
   139  	return &BTree{
   140  		degree: degree,
   141  		cow:    &copyOnWriteContext{freelist: f},
   142  	}
   143  }
   144  
   145  // items stores items in a node.
   146  type items []Item
   147  
   148  // insertAt inserts a value into the given index, pushing all subsequent values
   149  // forward.
   150  func (s *items) insertAt(index int, item Item) {
   151  	*s = append(*s, nil)
   152  	if index < len(*s) {
   153  		copy((*s)[index+1:], (*s)[index:])
   154  	}
   155  	(*s)[index] = item
   156  }
   157  
   158  // removeAt removes a value at a given index, pulling all subsequent values
   159  // back.
   160  func (s *items) removeAt(index int) Item {
   161  	item := (*s)[index]
   162  	copy((*s)[index:], (*s)[index+1:])
   163  	(*s)[len(*s)-1] = nil
   164  	*s = (*s)[:len(*s)-1]
   165  	return item
   166  }
   167  
   168  // pop removes and returns the last element in the list.
   169  func (s *items) pop() (out Item) {
   170  	index := len(*s) - 1
   171  	out = (*s)[index]
   172  	(*s)[index] = nil
   173  	*s = (*s)[:index]
   174  	return
   175  }
   176  
   177  // truncate truncates this instance at index so that it contains only the
   178  // first index items. index must be less than or equal to length.
   179  func (s *items) truncate(index int) {
   180  	var toClear items
   181  	*s, toClear = (*s)[:index], (*s)[index:]
   182  	for len(toClear) > 0 {
   183  		toClear = toClear[copy(toClear, nilItems):]
   184  	}
   185  }
   186  
   187  // find returns the index where the given item should be inserted into this
   188  // list.  'found' is true if the item already exists in the list at the given
   189  // index.
   190  func (s items) find(item Item) (index int, found bool) {
   191  	i := sort.Search(len(s), func(i int) bool {
   192  		return item.Less(s[i])
   193  	})
   194  	if i > 0 && !s[i-1].Less(item) {
   195  		return i - 1, true
   196  	}
   197  	return i, false
   198  }
   199  
   200  // children stores child nodes in a node.
   201  type children []*node
   202  
   203  // insertAt inserts a value into the given index, pushing all subsequent values
   204  // forward.
   205  func (s *children) insertAt(index int, n *node) {
   206  	*s = append(*s, nil)
   207  	if index < len(*s) {
   208  		copy((*s)[index+1:], (*s)[index:])
   209  	}
   210  	(*s)[index] = n
   211  }
   212  
   213  // removeAt removes a value at a given index, pulling all subsequent values
   214  // back.
   215  func (s *children) removeAt(index int) *node {
   216  	n := (*s)[index]
   217  	copy((*s)[index:], (*s)[index+1:])
   218  	(*s)[len(*s)-1] = nil
   219  	*s = (*s)[:len(*s)-1]
   220  	return n
   221  }
   222  
   223  // pop removes and returns the last element in the list.
   224  func (s *children) pop() (out *node) {
   225  	index := len(*s) - 1
   226  	out = (*s)[index]
   227  	(*s)[index] = nil
   228  	*s = (*s)[:index]
   229  	return
   230  }
   231  
   232  // truncate truncates this instance at index so that it contains only the
   233  // first index children. index must be less than or equal to length.
   234  func (s *children) truncate(index int) {
   235  	var toClear children
   236  	*s, toClear = (*s)[:index], (*s)[index:]
   237  	for len(toClear) > 0 {
   238  		toClear = toClear[copy(toClear, nilChildren):]
   239  	}
   240  }
   241  
   242  // node is an internal node in a tree.
   243  //
   244  // It must at all times maintain the invariant that either
   245  //   * len(children) == 0, len(items) unconstrained
   246  //   * len(children) == len(items) + 1
   247  type node struct {
   248  	items    items
   249  	children children
   250  	cow      *copyOnWriteContext
   251  }
   252  
   253  func (n *node) mutableFor(cow *copyOnWriteContext) *node {
   254  	if n.cow == cow {
   255  		return n
   256  	}
   257  	out := cow.newNode()
   258  	if cap(out.items) >= len(n.items) {
   259  		out.items = out.items[:len(n.items)]
   260  	} else {
   261  		out.items = make(items, len(n.items), cap(n.items))
   262  	}
   263  	copy(out.items, n.items)
   264  	// Copy children
   265  	if cap(out.children) >= len(n.children) {
   266  		out.children = out.children[:len(n.children)]
   267  	} else {
   268  		out.children = make(children, len(n.children), cap(n.children))
   269  	}
   270  	copy(out.children, n.children)
   271  	return out
   272  }
   273  
   274  func (n *node) mutableChild(i int) *node {
   275  	c := n.children[i].mutableFor(n.cow)
   276  	n.children[i] = c
   277  	return c
   278  }
   279  
   280  // split splits the given node at the given index.  The current node shrinks,
   281  // and this function returns the item that existed at that index and a new node
   282  // containing all items/children after it.
   283  func (n *node) split(i int) (Item, *node) {
   284  	item := n.items[i]
   285  	next := n.cow.newNode()
   286  	next.items = append(next.items, n.items[i+1:]...)
   287  	n.items.truncate(i)
   288  	if len(n.children) > 0 {
   289  		next.children = append(next.children, n.children[i+1:]...)
   290  		n.children.truncate(i + 1)
   291  	}
   292  	return item, next
   293  }
   294  
   295  // maybeSplitChild checks if a child should be split, and if so splits it.
   296  // Returns whether or not a split occurred.
   297  func (n *node) maybeSplitChild(i, maxItems int) bool {
   298  	if len(n.children[i].items) < maxItems {
   299  		return false
   300  	}
   301  	first := n.mutableChild(i)
   302  	item, second := first.split(maxItems / 2)
   303  	n.items.insertAt(i, item)
   304  	n.children.insertAt(i+1, second)
   305  	return true
   306  }
   307  
   308  // insert inserts an item into the subtree rooted at this node, making sure
   309  // no nodes in the subtree exceed maxItems items.  Should an equivalent item be
   310  // be found/replaced by insert, it will be returned.
   311  func (n *node) insert(item Item, maxItems int) Item {
   312  	i, found := n.items.find(item)
   313  	if found {
   314  		out := n.items[i]
   315  		n.items[i] = item
   316  		return out
   317  	}
   318  	if len(n.children) == 0 {
   319  		n.items.insertAt(i, item)
   320  		return nil
   321  	}
   322  	if n.maybeSplitChild(i, maxItems) {
   323  		inTree := n.items[i]
   324  		switch {
   325  		case item.Less(inTree):
   326  			// no change, we want first split node
   327  		case inTree.Less(item):
   328  			i++ // we want second split node
   329  		default:
   330  			out := n.items[i]
   331  			n.items[i] = item
   332  			return out
   333  		}
   334  	}
   335  	return n.mutableChild(i).insert(item, maxItems)
   336  }
   337  
   338  // get finds the given key in the subtree and returns it.
   339  func (n *node) get(key Item) Item {
   340  	i, found := n.items.find(key)
   341  	if found {
   342  		return n.items[i]
   343  	} else if len(n.children) > 0 {
   344  		return n.children[i].get(key)
   345  	}
   346  	return nil
   347  }
   348  
   349  // min returns the first item in the subtree.
   350  func min(n *node) Item {
   351  	if n == nil {
   352  		return nil
   353  	}
   354  	for len(n.children) > 0 {
   355  		n = n.children[0]
   356  	}
   357  	if len(n.items) == 0 {
   358  		return nil
   359  	}
   360  	return n.items[0]
   361  }
   362  
   363  // max returns the last item in the subtree.
   364  func max(n *node) Item {
   365  	if n == nil {
   366  		return nil
   367  	}
   368  	for len(n.children) > 0 {
   369  		n = n.children[len(n.children)-1]
   370  	}
   371  	if len(n.items) == 0 {
   372  		return nil
   373  	}
   374  	return n.items[len(n.items)-1]
   375  }
   376  
   377  // toRemove details what item to remove in a node.remove call.
   378  type toRemove int
   379  
   380  const (
   381  	removeItem toRemove = iota // removes the given item
   382  	removeMin                  // removes smallest item in the subtree
   383  	removeMax                  // removes largest item in the subtree
   384  )
   385  
   386  // remove removes an item from the subtree rooted at this node.
   387  func (n *node) remove(item Item, minItems int, typ toRemove) Item {
   388  	var i int
   389  	var found bool
   390  	switch typ {
   391  	case removeMax:
   392  		if len(n.children) == 0 {
   393  			return n.items.pop()
   394  		}
   395  		i = len(n.items)
   396  	case removeMin:
   397  		if len(n.children) == 0 {
   398  			return n.items.removeAt(0)
   399  		}
   400  		i = 0
   401  	case removeItem:
   402  		i, found = n.items.find(item)
   403  		if len(n.children) == 0 {
   404  			if found {
   405  				return n.items.removeAt(i)
   406  			}
   407  			return nil
   408  		}
   409  	default:
   410  		panic("invalid type")
   411  	}
   412  	// If we get to here, we have children.
   413  	if len(n.children[i].items) <= minItems {
   414  		return n.growChildAndRemove(i, item, minItems, typ)
   415  	}
   416  	child := n.mutableChild(i)
   417  	// Either we had enough items to begin with, or we've done some
   418  	// merging/stealing, because we've got enough now and we're ready to return
   419  	// stuff.
   420  	if found {
   421  		// The item exists at index 'i', and the child we've selected can give us a
   422  		// predecessor, since if we've gotten here it's got > minItems items in it.
   423  		out := n.items[i]
   424  		// We use our special-case 'remove' call with typ=maxItem to pull the
   425  		// predecessor of item i (the rightmost leaf of our immediate left child)
   426  		// and set it into where we pulled the item from.
   427  		n.items[i] = child.remove(nil, minItems, removeMax)
   428  		return out
   429  	}
   430  	// Final recursive call.  Once we're here, we know that the item isn't in this
   431  	// node and that the child is big enough to remove from.
   432  	return child.remove(item, minItems, typ)
   433  }
   434  
   435  // growChildAndRemove grows child 'i' to make sure it's possible to remove an
   436  // item from it while keeping it at minItems, then calls remove to actually
   437  // remove it.
   438  //
   439  // Most documentation says we have to do two sets of special casing:
   440  //   1) item is in this node
   441  //   2) item is in child
   442  // In both cases, we need to handle the two subcases:
   443  //   A) node has enough values that it can spare one
   444  //   B) node doesn't have enough values
   445  // For the latter, we have to check:
   446  //   a) left sibling has node to spare
   447  //   b) right sibling has node to spare
   448  //   c) we must merge
   449  // To simplify our code here, we handle cases #1 and #2 the same:
   450  // If a node doesn't have enough items, we make sure it does (using a,b,c).
   451  // We then simply redo our remove call, and the second time (regardless of
   452  // whether we're in case 1 or 2), we'll have enough items and can guarantee
   453  // that we hit case A.
   454  func (n *node) growChildAndRemove(i int, item Item, minItems int, typ toRemove) Item {
   455  	if i > 0 && len(n.children[i-1].items) > minItems {
   456  		// Steal from left child
   457  		child := n.mutableChild(i)
   458  		stealFrom := n.mutableChild(i - 1)
   459  		stolenItem := stealFrom.items.pop()
   460  		child.items.insertAt(0, n.items[i-1])
   461  		n.items[i-1] = stolenItem
   462  		if len(stealFrom.children) > 0 {
   463  			child.children.insertAt(0, stealFrom.children.pop())
   464  		}
   465  	} else if i < len(n.items) && len(n.children[i+1].items) > minItems {
   466  		// steal from right child
   467  		child := n.mutableChild(i)
   468  		stealFrom := n.mutableChild(i + 1)
   469  		stolenItem := stealFrom.items.removeAt(0)
   470  		child.items = append(child.items, n.items[i])
   471  		n.items[i] = stolenItem
   472  		if len(stealFrom.children) > 0 {
   473  			child.children = append(child.children, stealFrom.children.removeAt(0))
   474  		}
   475  	} else {
   476  		if i >= len(n.items) {
   477  			i--
   478  		}
   479  		child := n.mutableChild(i)
   480  		// merge with right child
   481  		mergeItem := n.items.removeAt(i)
   482  		mergeChild := n.children.removeAt(i + 1)
   483  		child.items = append(child.items, mergeItem)
   484  		child.items = append(child.items, mergeChild.items...)
   485  		child.children = append(child.children, mergeChild.children...)
   486  		n.cow.freeNode(mergeChild)
   487  	}
   488  	return n.remove(item, minItems, typ)
   489  }
   490  
   491  type direction int
   492  
   493  const (
   494  	descend = direction(-1)
   495  	ascend  = direction(+1)
   496  )
   497  
   498  // iterate provides a simple method for iterating over elements in the tree.
   499  //
   500  // When ascending, the 'start' should be less than 'stop' and when descending,
   501  // the 'start' should be greater than 'stop'. Setting 'includeStart' to true
   502  // will force the iterator to include the first item when it equals 'start',
   503  // thus creating a "greaterOrEqual" or "lessThanEqual" rather than just a
   504  // "greaterThan" or "lessThan" queries.
   505  func (n *node) iterate(dir direction, start, stop Item, includeStart bool, hit bool, iter ItemIterator) (bool, bool) {
   506  	var ok, found bool
   507  	var index int
   508  	switch dir {
   509  	case ascend:
   510  		if start != nil {
   511  			index, _ = n.items.find(start)
   512  		}
   513  		for i := index; i < len(n.items); i++ {
   514  			if len(n.children) > 0 {
   515  				if hit, ok = n.children[i].iterate(dir, start, stop, includeStart, hit, iter); !ok {
   516  					return hit, false
   517  				}
   518  			}
   519  			if !includeStart && !hit && start != nil && !start.Less(n.items[i]) {
   520  				hit = true
   521  				continue
   522  			}
   523  			hit = true
   524  			if stop != nil && !n.items[i].Less(stop) {
   525  				return hit, false
   526  			}
   527  			if !iter(n.items[i]) {
   528  				return hit, false
   529  			}
   530  		}
   531  		if len(n.children) > 0 {
   532  			if hit, ok = n.children[len(n.children)-1].iterate(dir, start, stop, includeStart, hit, iter); !ok {
   533  				return hit, false
   534  			}
   535  		}
   536  	case descend:
   537  		if start != nil {
   538  			index, found = n.items.find(start)
   539  			if !found {
   540  				index = index - 1
   541  			}
   542  		} else {
   543  			index = len(n.items) - 1
   544  		}
   545  		for i := index; i >= 0; i-- {
   546  			if start != nil && !n.items[i].Less(start) {
   547  				if !includeStart || hit || start.Less(n.items[i]) {
   548  					continue
   549  				}
   550  			}
   551  			if len(n.children) > 0 {
   552  				if hit, ok = n.children[i+1].iterate(dir, start, stop, includeStart, hit, iter); !ok {
   553  					return hit, false
   554  				}
   555  			}
   556  			if stop != nil && !stop.Less(n.items[i]) {
   557  				return hit, false //	continue
   558  			}
   559  			hit = true
   560  			if !iter(n.items[i]) {
   561  				return hit, false
   562  			}
   563  		}
   564  		if len(n.children) > 0 {
   565  			if hit, ok = n.children[0].iterate(dir, start, stop, includeStart, hit, iter); !ok {
   566  				return hit, false
   567  			}
   568  		}
   569  	}
   570  	return hit, true
   571  }
   572  
   573  // Used for testing/debugging purposes.
   574  func (n *node) print(w io.Writer, level int) {
   575  	fmt.Fprintf(w, "%sNODE:%v\n", strings.Repeat("  ", level), n.items)
   576  	for _, c := range n.children {
   577  		c.print(w, level+1)
   578  	}
   579  }
   580  
   581  // BTree is an implementation of a B-Tree.
   582  //
   583  // BTree stores Item instances in an ordered structure, allowing easy insertion,
   584  // removal, and iteration.
   585  //
   586  // Write operations are not safe for concurrent mutation by multiple
   587  // goroutines, but Read operations are.
   588  type BTree struct {
   589  	degree int
   590  	length int
   591  	root   *node
   592  	cow    *copyOnWriteContext
   593  }
   594  
   595  // copyOnWriteContext pointers determine node ownership... a tree with a write
   596  // context equivalent to a node's write context is allowed to modify that node.
   597  // A tree whose write context does not match a node's is not allowed to modify
   598  // it, and must create a new, writable copy (IE: it's a Clone).
   599  //
   600  // When doing any write operation, we maintain the invariant that the current
   601  // node's context is equal to the context of the tree that requested the write.
   602  // We do this by, before we descend into any node, creating a copy with the
   603  // correct context if the contexts don't match.
   604  //
   605  // Since the node we're currently visiting on any write has the requesting
   606  // tree's context, that node is modifiable in place.  Children of that node may
   607  // not share context, but before we descend into them, we'll make a mutable
   608  // copy.
   609  type copyOnWriteContext struct {
   610  	freelist *FreeList
   611  }
   612  
   613  // Clone clones the btree, lazily.  Clone should not be called concurrently,
   614  // but the original tree (t) and the new tree (t2) can be used concurrently
   615  // once the Clone call completes.
   616  //
   617  // The internal tree structure of b is marked read-only and shared between t and
   618  // t2.  Writes to both t and t2 use copy-on-write logic, creating new nodes
   619  // whenever one of b's original nodes would have been modified.  Read operations
   620  // should have no performance degredation.  Write operations for both t and t2
   621  // will initially experience minor slow-downs caused by additional allocs and
   622  // copies due to the aforementioned copy-on-write logic, but should converge to
   623  // the original performance characteristics of the original tree.
   624  func (t *BTree) Clone() (t2 *BTree) {
   625  	// Create two entirely new copy-on-write contexts.
   626  	// This operation effectively creates three trees:
   627  	//   the original, shared nodes (old b.cow)
   628  	//   the new b.cow nodes
   629  	//   the new out.cow nodes
   630  	cow1, cow2 := *t.cow, *t.cow
   631  	out := *t
   632  	t.cow = &cow1
   633  	out.cow = &cow2
   634  	return &out
   635  }
   636  
   637  // maxItems returns the max number of items to allow per node.
   638  func (t *BTree) maxItems() int {
   639  	return t.degree*2 - 1
   640  }
   641  
   642  // minItems returns the min number of items to allow per node (ignored for the
   643  // root node).
   644  func (t *BTree) minItems() int {
   645  	return t.degree - 1
   646  }
   647  
   648  func (c *copyOnWriteContext) newNode() (n *node) {
   649  	n = c.freelist.newNode()
   650  	n.cow = c
   651  	return
   652  }
   653  
   654  type freeType int
   655  
   656  const (
   657  	ftFreelistFull freeType = iota // node was freed (available for GC, not stored in freelist)
   658  	ftStored                       // node was stored in the freelist for later use
   659  	ftNotOwned                     // node was ignored by COW, since it's owned by another one
   660  )
   661  
   662  // freeNode frees a node within a given COW context, if it's owned by that
   663  // context.  It returns what happened to the node (see freeType const
   664  // documentation).
   665  func (c *copyOnWriteContext) freeNode(n *node) freeType {
   666  	if n.cow == c {
   667  		// clear to allow GC
   668  		n.items.truncate(0)
   669  		n.children.truncate(0)
   670  		n.cow = nil
   671  		if c.freelist.freeNode(n) {
   672  			return ftStored
   673  		} else {
   674  			return ftFreelistFull
   675  		}
   676  	} else {
   677  		return ftNotOwned
   678  	}
   679  }
   680  
   681  // ReplaceOrInsert adds the given item to the tree.  If an item in the tree
   682  // already equals the given one, it is removed from the tree and returned.
   683  // Otherwise, nil is returned.
   684  //
   685  // nil cannot be added to the tree (will panic).
   686  func (t *BTree) ReplaceOrInsert(item Item) Item {
   687  	if item == nil {
   688  		panic("nil item being added to BTree")
   689  	}
   690  	if t.root == nil {
   691  		t.root = t.cow.newNode()
   692  		t.root.items = append(t.root.items, item)
   693  		t.length++
   694  		return nil
   695  	} else {
   696  		t.root = t.root.mutableFor(t.cow)
   697  		if len(t.root.items) >= t.maxItems() {
   698  			item2, second := t.root.split(t.maxItems() / 2)
   699  			oldroot := t.root
   700  			t.root = t.cow.newNode()
   701  			t.root.items = append(t.root.items, item2)
   702  			t.root.children = append(t.root.children, oldroot, second)
   703  		}
   704  	}
   705  	out := t.root.insert(item, t.maxItems())
   706  	if out == nil {
   707  		t.length++
   708  	}
   709  	return out
   710  }
   711  
   712  // Delete removes an item equal to the passed in item from the tree, returning
   713  // it.  If no such item exists, returns nil.
   714  func (t *BTree) Delete(item Item) Item {
   715  	return t.deleteItem(item, removeItem)
   716  }
   717  
   718  // DeleteMin removes the smallest item in the tree and returns it.
   719  // If no such item exists, returns nil.
   720  func (t *BTree) DeleteMin() Item {
   721  	return t.deleteItem(nil, removeMin)
   722  }
   723  
   724  // DeleteMax removes the largest item in the tree and returns it.
   725  // If no such item exists, returns nil.
   726  func (t *BTree) DeleteMax() Item {
   727  	return t.deleteItem(nil, removeMax)
   728  }
   729  
   730  func (t *BTree) deleteItem(item Item, typ toRemove) Item {
   731  	if t.root == nil || len(t.root.items) == 0 {
   732  		return nil
   733  	}
   734  	t.root = t.root.mutableFor(t.cow)
   735  	out := t.root.remove(item, t.minItems(), typ)
   736  	if len(t.root.items) == 0 && len(t.root.children) > 0 {
   737  		oldroot := t.root
   738  		t.root = t.root.children[0]
   739  		t.cow.freeNode(oldroot)
   740  	}
   741  	if out != nil {
   742  		t.length--
   743  	}
   744  	return out
   745  }
   746  
   747  // AscendRange calls the iterator for every value in the tree within the range
   748  // [greaterOrEqual, lessThan), until iterator returns false.
   749  func (t *BTree) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) {
   750  	if t.root == nil {
   751  		return
   752  	}
   753  	t.root.iterate(ascend, greaterOrEqual, lessThan, true, false, iterator)
   754  }
   755  
   756  // AscendLessThan calls the iterator for every value in the tree within the range
   757  // [first, pivot), until iterator returns false.
   758  func (t *BTree) AscendLessThan(pivot Item, iterator ItemIterator) {
   759  	if t.root == nil {
   760  		return
   761  	}
   762  	t.root.iterate(ascend, nil, pivot, false, false, iterator)
   763  }
   764  
   765  // AscendGreaterOrEqual calls the iterator for every value in the tree within
   766  // the range [pivot, last], until iterator returns false.
   767  func (t *BTree) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {
   768  	if t.root == nil {
   769  		return
   770  	}
   771  	t.root.iterate(ascend, pivot, nil, true, false, iterator)
   772  }
   773  
   774  // Ascend calls the iterator for every value in the tree within the range
   775  // [first, last], until iterator returns false.
   776  func (t *BTree) Ascend(iterator ItemIterator) {
   777  	if t.root == nil {
   778  		return
   779  	}
   780  	t.root.iterate(ascend, nil, nil, false, false, iterator)
   781  }
   782  
   783  // DescendRange calls the iterator for every value in the tree within the range
   784  // [lessOrEqual, greaterThan), until iterator returns false.
   785  func (t *BTree) DescendRange(lessOrEqual, greaterThan Item, iterator ItemIterator) {
   786  	if t.root == nil {
   787  		return
   788  	}
   789  	t.root.iterate(descend, lessOrEqual, greaterThan, true, false, iterator)
   790  }
   791  
   792  // DescendLessOrEqual calls the iterator for every value in the tree within the range
   793  // [pivot, first], until iterator returns false.
   794  func (t *BTree) DescendLessOrEqual(pivot Item, iterator ItemIterator) {
   795  	if t.root == nil {
   796  		return
   797  	}
   798  	t.root.iterate(descend, pivot, nil, true, false, iterator)
   799  }
   800  
   801  // DescendGreaterThan calls the iterator for every value in the tree within
   802  // the range [last, pivot), until iterator returns false.
   803  func (t *BTree) DescendGreaterThan(pivot Item, iterator ItemIterator) {
   804  	if t.root == nil {
   805  		return
   806  	}
   807  	t.root.iterate(descend, nil, pivot, false, false, iterator)
   808  }
   809  
   810  // Descend calls the iterator for every value in the tree within the range
   811  // [last, first], until iterator returns false.
   812  func (t *BTree) Descend(iterator ItemIterator) {
   813  	if t.root == nil {
   814  		return
   815  	}
   816  	t.root.iterate(descend, nil, nil, false, false, iterator)
   817  }
   818  
   819  // Get looks for the key item in the tree, returning it.  It returns nil if
   820  // unable to find that item.
   821  func (t *BTree) Get(key Item) Item {
   822  	if t.root == nil {
   823  		return nil
   824  	}
   825  	return t.root.get(key)
   826  }
   827  
   828  // Min returns the smallest item in the tree, or nil if the tree is empty.
   829  func (t *BTree) Min() Item {
   830  	return min(t.root)
   831  }
   832  
   833  // Max returns the largest item in the tree, or nil if the tree is empty.
   834  func (t *BTree) Max() Item {
   835  	return max(t.root)
   836  }
   837  
   838  // Has returns true if the given key is in the tree.
   839  func (t *BTree) Has(key Item) bool {
   840  	return t.Get(key) != nil
   841  }
   842  
   843  // Len returns the number of items currently in the tree.
   844  func (t *BTree) Len() int {
   845  	return t.length
   846  }
   847  
   848  // Clear removes all items from the btree.  If addNodesToFreelist is true,
   849  // t's nodes are added to its freelist as part of this call, until the freelist
   850  // is full.  Otherwise, the root node is simply dereferenced and the subtree
   851  // left to Go's normal GC processes.
   852  //
   853  // This can be much faster
   854  // than calling Delete on all elements, because that requires finding/removing
   855  // each element in the tree and updating the tree accordingly.  It also is
   856  // somewhat faster than creating a new tree to replace the old one, because
   857  // nodes from the old tree are reclaimed into the freelist for use by the new
   858  // one, instead of being lost to the garbage collector.
   859  //
   860  // This call takes:
   861  //   O(1): when addNodesToFreelist is false, this is a single operation.
   862  //   O(1): when the freelist is already full, it breaks out immediately
   863  //   O(freelist size):  when the freelist is empty and the nodes are all owned
   864  //       by this tree, nodes are added to the freelist until full.
   865  //   O(tree size):  when all nodes are owned by another tree, all nodes are
   866  //       iterated over looking for nodes to add to the freelist, and due to
   867  //       ownership, none are.
   868  func (t *BTree) Clear(addNodesToFreelist bool) {
   869  	if t.root != nil && addNodesToFreelist {
   870  		t.root.reset(t.cow)
   871  	}
   872  	t.root, t.length = nil, 0
   873  }
   874  
   875  // reset returns a subtree to the freelist.  It breaks out immediately if the
   876  // freelist is full, since the only benefit of iterating is to fill that
   877  // freelist up.  Returns true if parent reset call should continue.
   878  func (n *node) reset(c *copyOnWriteContext) bool {
   879  	for _, child := range n.children {
   880  		if !child.reset(c) {
   881  			return false
   882  		}
   883  	}
   884  	return c.freeNode(n) != ftFreelistFull
   885  }
   886  
   887  // Int implements the Item interface for integers.
   888  type Int int
   889  
   890  // Less returns true if int(a) < int(b).
   891  func (a Int) Less(b Item) bool {
   892  	return a < b.(Int)
   893  }
   894
View as plain text