Source file src/github.com/google/btree/btree_generic.go

Documentation: github.com/google/btree

     1  // Copyright 2014-2022 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  // In Go 1.18 and beyond, a BTreeG generic is created, and BTree is a specific
    19  // instantiation of that generic for the Item interface, with a backwards-
    20  // compatible API.  Before go1.18, generics are not supported,
    21  // and BTree is just an implementation based around the Item interface.
    22  
    23  // Package btree implements in-memory B-Trees of arbitrary degree.
    24  //
    25  // btree implements an in-memory B-Tree for use as an ordered data structure.
    26  // It is not meant for persistent storage solutions.
    27  //
    28  // It has a flatter structure than an equivalent red-black or other binary tree,
    29  // which in some cases yields better memory usage and/or performance.
    30  // See some discussion on the matter here:
    31  //   http://google-opensource.blogspot.com/2013/01/c-containers-that-save-memory-and-time.html
    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  // These issues don't tend to matter, though, when working with strings or other
    47  // heap-allocated structures, since C++-equivalent structures also must store
    48  // pointers and also distribute their values across the heap.
    49  //
    50  // This implementation is designed to be a drop-in replacement to gollrb.LLRB
    51  // trees, (http://github.com/petar/gollrb), an excellent and probably the most
    52  // widely used ordered tree implementation in the Go ecosystem currently.
    53  // Its functions, therefore, exactly mirror those of
    54  // llrb.LLRB where possible.  Unlike gollrb, though, we currently don't
    55  // support storing multiple equivalent values.
    56  //
    57  // There are two implementations; those suffixed with 'G' are generics, usable
    58  // for any type, and require a passed-in "less" function to define their ordering.
    59  // Those without this prefix are specific to the 'Item' interface, and use
    60  // its 'Less' function for ordering.
    61  package btree
    62  
    63  import (
    64  	"fmt"
    65  	"io"
    66  	"sort"
    67  	"strings"
    68  	"sync"
    69  )
    70  
    71  // Item represents a single object in the tree.
    72  type Item interface {
    73  	// Less tests whether the current item is less than the given argument.
    74  	//
    75  	// This must provide a strict weak ordering.
    76  	// If !a.Less(b) && !b.Less(a), we treat this to mean a == b (i.e. we can only
    77  	// hold one of either a or b in the tree).
    78  	Less(than Item) bool
    79  }
    80  
    81  const (
    82  	DefaultFreeListSize = 32
    83  )
    84  
    85  // FreeListG represents a free list of btree nodes. By default each
    86  // BTree has its own FreeList, but multiple BTrees can share the same
    87  // FreeList, in particular when they're created with Clone.
    88  // Two Btrees using the same freelist are safe for concurrent write access.
    89  type FreeListG[T any] struct {
    90  	mu       sync.Mutex
    91  	freelist []*node[T]
    92  }
    93  
    94  // NewFreeListG creates a new free list.
    95  // size is the maximum size of the returned free list.
    96  func NewFreeListG[T any](size int) *FreeListG[T] {
    97  	return &FreeListG[T]{freelist: make([]*node[T], 0, size)}
    98  }
    99  
   100  func (f *FreeListG[T]) newNode() (n *node[T]) {
   101  	f.mu.Lock()
   102  	index := len(f.freelist) - 1
   103  	if index < 0 {
   104  		f.mu.Unlock()
   105  		return new(node[T])
   106  	}
   107  	n = f.freelist[index]
   108  	f.freelist[index] = nil
   109  	f.freelist = f.freelist[:index]
   110  	f.mu.Unlock()
   111  	return
   112  }
   113  
   114  func (f *FreeListG[T]) freeNode(n *node[T]) (out bool) {
   115  	f.mu.Lock()
   116  	if len(f.freelist) < cap(f.freelist) {
   117  		f.freelist = append(f.freelist, n)
   118  		out = true
   119  	}
   120  	f.mu.Unlock()
   121  	return
   122  }
   123  
   124  // ItemIteratorG allows callers of {A/De}scend* to iterate in-order over portions of
   125  // the tree.  When this function returns false, iteration will stop and the
   126  // associated Ascend* function will immediately return.
   127  type ItemIteratorG[T any] func(item T) bool
   128  
   129  // Ordered represents the set of types for which the '<' operator work.
   130  type Ordered interface {
   131  	~int | ~int8 | ~int16 | ~int32 | ~int64 | ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~float32 | ~float64 | ~string
   132  }
   133  
   134  // Less[T] returns a default LessFunc that uses the '<' operator for types that support it.
   135  func Less[T Ordered]() LessFunc[T] {
   136  	return func(a, b T) bool { return a < b }
   137  }
   138  
   139  // NewOrderedG creates a new B-Tree for ordered types.
   140  func NewOrderedG[T Ordered](degree int) *BTreeG[T] {
   141  	return NewG[T](degree, Less[T]())
   142  }
   143  
   144  // NewG creates a new B-Tree with the given degree.
   145  //
   146  // NewG(2), for example, will create a 2-3-4 tree (each node contains 1-3 items
   147  // and 2-4 children).
   148  //
   149  // The passed-in LessFunc determines how objects of type T are ordered.
   150  func NewG[T any](degree int, less LessFunc[T]) *BTreeG[T] {
   151  	return NewWithFreeListG(degree, less, NewFreeListG[T](DefaultFreeListSize))
   152  }
   153  
   154  // NewWithFreeListG creates a new B-Tree that uses the given node free list.
   155  func NewWithFreeListG[T any](degree int, less LessFunc[T], f *FreeListG[T]) *BTreeG[T] {
   156  	if degree <= 1 {
   157  		panic("bad degree")
   158  	}
   159  	return &BTreeG[T]{
   160  		degree: degree,
   161  		cow:    &copyOnWriteContext[T]{freelist: f, less: less},
   162  	}
   163  }
   164  
   165  // items stores items in a node.
   166  type items[T any] []T
   167  
   168  // insertAt inserts a value into the given index, pushing all subsequent values
   169  // forward.
   170  func (s *items[T]) insertAt(index int, item T) {
   171  	var zero T
   172  	*s = append(*s, zero)
   173  	if index < len(*s) {
   174  		copy((*s)[index+1:], (*s)[index:])
   175  	}
   176  	(*s)[index] = item
   177  }
   178  
   179  // removeAt removes a value at a given index, pulling all subsequent values
   180  // back.
   181  func (s *items[T]) removeAt(index int) T {
   182  	item := (*s)[index]
   183  	copy((*s)[index:], (*s)[index+1:])
   184  	var zero T
   185  	(*s)[len(*s)-1] = zero
   186  	*s = (*s)[:len(*s)-1]
   187  	return item
   188  }
   189  
   190  // pop removes and returns the last element in the list.
   191  func (s *items[T]) pop() (out T) {
   192  	index := len(*s) - 1
   193  	out = (*s)[index]
   194  	var zero T
   195  	(*s)[index] = zero
   196  	*s = (*s)[:index]
   197  	return
   198  }
   199  
   200  // truncate truncates this instance at index so that it contains only the
   201  // first index items. index must be less than or equal to length.
   202  func (s *items[T]) truncate(index int) {
   203  	var toClear items[T]
   204  	*s, toClear = (*s)[:index], (*s)[index:]
   205  	var zero T
   206  	for i := 0; i < len(toClear); i++ {
   207  		toClear[i] = zero
   208  	}
   209  }
   210  
   211  // find returns the index where the given item should be inserted into this
   212  // list.  'found' is true if the item already exists in the list at the given
   213  // index.
   214  func (s items[T]) find(item T, less func(T, T) bool) (index int, found bool) {
   215  	i := sort.Search(len(s), func(i int) bool {
   216  		return less(item, s[i])
   217  	})
   218  	if i > 0 && !less(s[i-1], item) {
   219  		return i - 1, true
   220  	}
   221  	return i, false
   222  }
   223  
   224  // node is an internal node in a tree.
   225  //
   226  // It must at all times maintain the invariant that either
   227  //   * len(children) == 0, len(items) unconstrained
   228  //   * len(children) == len(items) + 1
   229  type node[T any] struct {
   230  	items    items[T]
   231  	children items[*node[T]]
   232  	cow      *copyOnWriteContext[T]
   233  }
   234  
   235  func (n *node[T]) mutableFor(cow *copyOnWriteContext[T]) *node[T] {
   236  	if n.cow == cow {
   237  		return n
   238  	}
   239  	out := cow.newNode()
   240  	if cap(out.items) >= len(n.items) {
   241  		out.items = out.items[:len(n.items)]
   242  	} else {
   243  		out.items = make(items[T], len(n.items), cap(n.items))
   244  	}
   245  	copy(out.items, n.items)
   246  	// Copy children
   247  	if cap(out.children) >= len(n.children) {
   248  		out.children = out.children[:len(n.children)]
   249  	} else {
   250  		out.children = make(items[*node[T]], len(n.children), cap(n.children))
   251  	}
   252  	copy(out.children, n.children)
   253  	return out
   254  }
   255  
   256  func (n *node[T]) mutableChild(i int) *node[T] {
   257  	c := n.children[i].mutableFor(n.cow)
   258  	n.children[i] = c
   259  	return c
   260  }
   261  
   262  // split splits the given node at the given index.  The current node shrinks,
   263  // and this function returns the item that existed at that index and a new node
   264  // containing all items/children after it.
   265  func (n *node[T]) split(i int) (T, *node[T]) {
   266  	item := n.items[i]
   267  	next := n.cow.newNode()
   268  	next.items = append(next.items, n.items[i+1:]...)
   269  	n.items.truncate(i)
   270  	if len(n.children) > 0 {
   271  		next.children = append(next.children, n.children[i+1:]...)
   272  		n.children.truncate(i + 1)
   273  	}
   274  	return item, next
   275  }
   276  
   277  // maybeSplitChild checks if a child should be split, and if so splits it.
   278  // Returns whether or not a split occurred.
   279  func (n *node[T]) maybeSplitChild(i, maxItems int) bool {
   280  	if len(n.children[i].items) < maxItems {
   281  		return false
   282  	}
   283  	first := n.mutableChild(i)
   284  	item, second := first.split(maxItems / 2)
   285  	n.items.insertAt(i, item)
   286  	n.children.insertAt(i+1, second)
   287  	return true
   288  }
   289  
   290  // insert inserts an item into the subtree rooted at this node, making sure
   291  // no nodes in the subtree exceed maxItems items.  Should an equivalent item be
   292  // be found/replaced by insert, it will be returned.
   293  func (n *node[T]) insert(item T, maxItems int) (_ T, _ bool) {
   294  	i, found := n.items.find(item, n.cow.less)
   295  	if found {
   296  		out := n.items[i]
   297  		n.items[i] = item
   298  		return out, true
   299  	}
   300  	if len(n.children) == 0 {
   301  		n.items.insertAt(i, item)
   302  		return
   303  	}
   304  	if n.maybeSplitChild(i, maxItems) {
   305  		inTree := n.items[i]
   306  		switch {
   307  		case n.cow.less(item, inTree):
   308  			// no change, we want first split node
   309  		case n.cow.less(inTree, item):
   310  			i++ // we want second split node
   311  		default:
   312  			out := n.items[i]
   313  			n.items[i] = item
   314  			return out, true
   315  		}
   316  	}
   317  	return n.mutableChild(i).insert(item, maxItems)
   318  }
   319  
   320  // get finds the given key in the subtree and returns it.
   321  func (n *node[T]) get(key T) (_ T, _ bool) {
   322  	i, found := n.items.find(key, n.cow.less)
   323  	if found {
   324  		return n.items[i], true
   325  	} else if len(n.children) > 0 {
   326  		return n.children[i].get(key)
   327  	}
   328  	return
   329  }
   330  
   331  // min returns the first item in the subtree.
   332  func min[T any](n *node[T]) (_ T, found bool) {
   333  	if n == nil {
   334  		return
   335  	}
   336  	for len(n.children) > 0 {
   337  		n = n.children[0]
   338  	}
   339  	if len(n.items) == 0 {
   340  		return
   341  	}
   342  	return n.items[0], true
   343  }
   344  
   345  // max returns the last item in the subtree.
   346  func max[T any](n *node[T]) (_ T, found bool) {
   347  	if n == nil {
   348  		return
   349  	}
   350  	for len(n.children) > 0 {
   351  		n = n.children[len(n.children)-1]
   352  	}
   353  	if len(n.items) == 0 {
   354  		return
   355  	}
   356  	return n.items[len(n.items)-1], true
   357  }
   358  
   359  // toRemove details what item to remove in a node.remove call.
   360  type toRemove int
   361  
   362  const (
   363  	removeItem toRemove = iota // removes the given item
   364  	removeMin                  // removes smallest item in the subtree
   365  	removeMax                  // removes largest item in the subtree
   366  )
   367  
   368  // remove removes an item from the subtree rooted at this node.
   369  func (n *node[T]) remove(item T, minItems int, typ toRemove) (_ T, _ bool) {
   370  	var i int
   371  	var found bool
   372  	switch typ {
   373  	case removeMax:
   374  		if len(n.children) == 0 {
   375  			return n.items.pop(), true
   376  		}
   377  		i = len(n.items)
   378  	case removeMin:
   379  		if len(n.children) == 0 {
   380  			return n.items.removeAt(0), true
   381  		}
   382  		i = 0
   383  	case removeItem:
   384  		i, found = n.items.find(item, n.cow.less)
   385  		if len(n.children) == 0 {
   386  			if found {
   387  				return n.items.removeAt(i), true
   388  			}
   389  			return
   390  		}
   391  	default:
   392  		panic("invalid type")
   393  	}
   394  	// If we get to here, we have children.
   395  	if len(n.children[i].items) <= minItems {
   396  		return n.growChildAndRemove(i, item, minItems, typ)
   397  	}
   398  	child := n.mutableChild(i)
   399  	// Either we had enough items to begin with, or we've done some
   400  	// merging/stealing, because we've got enough now and we're ready to return
   401  	// stuff.
   402  	if found {
   403  		// The item exists at index 'i', and the child we've selected can give us a
   404  		// predecessor, since if we've gotten here it's got > minItems items in it.
   405  		out := n.items[i]
   406  		// We use our special-case 'remove' call with typ=maxItem to pull the
   407  		// predecessor of item i (the rightmost leaf of our immediate left child)
   408  		// and set it into where we pulled the item from.
   409  		var zero T
   410  		n.items[i], _ = child.remove(zero, minItems, removeMax)
   411  		return out, true
   412  	}
   413  	// Final recursive call.  Once we're here, we know that the item isn't in this
   414  	// node and that the child is big enough to remove from.
   415  	return child.remove(item, minItems, typ)
   416  }
   417  
   418  // growChildAndRemove grows child 'i' to make sure it's possible to remove an
   419  // item from it while keeping it at minItems, then calls remove to actually
   420  // remove it.
   421  //
   422  // Most documentation says we have to do two sets of special casing:
   423  //   1) item is in this node
   424  //   2) item is in child
   425  // In both cases, we need to handle the two subcases:
   426  //   A) node has enough values that it can spare one
   427  //   B) node doesn't have enough values
   428  // For the latter, we have to check:
   429  //   a) left sibling has node to spare
   430  //   b) right sibling has node to spare
   431  //   c) we must merge
   432  // To simplify our code here, we handle cases #1 and #2 the same:
   433  // If a node doesn't have enough items, we make sure it does (using a,b,c).
   434  // We then simply redo our remove call, and the second time (regardless of
   435  // whether we're in case 1 or 2), we'll have enough items and can guarantee
   436  // that we hit case A.
   437  func (n *node[T]) growChildAndRemove(i int, item T, minItems int, typ toRemove) (T, bool) {
   438  	if i > 0 && len(n.children[i-1].items) > minItems {
   439  		// Steal from left child
   440  		child := n.mutableChild(i)
   441  		stealFrom := n.mutableChild(i - 1)
   442  		stolenItem := stealFrom.items.pop()
   443  		child.items.insertAt(0, n.items[i-1])
   444  		n.items[i-1] = stolenItem
   445  		if len(stealFrom.children) > 0 {
   446  			child.children.insertAt(0, stealFrom.children.pop())
   447  		}
   448  	} else if i < len(n.items) && len(n.children[i+1].items) > minItems {
   449  		// steal from right child
   450  		child := n.mutableChild(i)
   451  		stealFrom := n.mutableChild(i + 1)
   452  		stolenItem := stealFrom.items.removeAt(0)
   453  		child.items = append(child.items, n.items[i])
   454  		n.items[i] = stolenItem
   455  		if len(stealFrom.children) > 0 {
   456  			child.children = append(child.children, stealFrom.children.removeAt(0))
   457  		}
   458  	} else {
   459  		if i >= len(n.items) {
   460  			i--
   461  		}
   462  		child := n.mutableChild(i)
   463  		// merge with right child
   464  		mergeItem := n.items.removeAt(i)
   465  		mergeChild := n.children.removeAt(i + 1)
   466  		child.items = append(child.items, mergeItem)
   467  		child.items = append(child.items, mergeChild.items...)
   468  		child.children = append(child.children, mergeChild.children...)
   469  		n.cow.freeNode(mergeChild)
   470  	}
   471  	return n.remove(item, minItems, typ)
   472  }
   473  
   474  type direction int
   475  
   476  const (
   477  	descend = direction(-1)
   478  	ascend  = direction(+1)
   479  )
   480  
   481  type optionalItem[T any] struct {
   482  	item  T
   483  	valid bool
   484  }
   485  
   486  func optional[T any](item T) optionalItem[T] {
   487  	return optionalItem[T]{item: item, valid: true}
   488  }
   489  func empty[T any]() optionalItem[T] {
   490  	return optionalItem[T]{}
   491  }
   492  
   493  // iterate provides a simple method for iterating over elements in the tree.
   494  //
   495  // When ascending, the 'start' should be less than 'stop' and when descending,
   496  // the 'start' should be greater than 'stop'. Setting 'includeStart' to true
   497  // will force the iterator to include the first item when it equals 'start',
   498  // thus creating a "greaterOrEqual" or "lessThanEqual" rather than just a
   499  // "greaterThan" or "lessThan" queries.
   500  func (n *node[T]) iterate(dir direction, start, stop optionalItem[T], includeStart bool, hit bool, iter ItemIteratorG[T]) (bool, bool) {
   501  	var ok, found bool
   502  	var index int
   503  	switch dir {
   504  	case ascend:
   505  		if start.valid {
   506  			index, _ = n.items.find(start.item, n.cow.less)
   507  		}
   508  		for i := index; i < len(n.items); i++ {
   509  			if len(n.children) > 0 {
   510  				if hit, ok = n.children[i].iterate(dir, start, stop, includeStart, hit, iter); !ok {
   511  					return hit, false
   512  				}
   513  			}
   514  			if !includeStart && !hit && start.valid && !n.cow.less(start.item, n.items[i]) {
   515  				hit = true
   516  				continue
   517  			}
   518  			hit = true
   519  			if stop.valid && !n.cow.less(n.items[i], stop.item) {
   520  				return hit, false
   521  			}
   522  			if !iter(n.items[i]) {
   523  				return hit, false
   524  			}
   525  		}
   526  		if len(n.children) > 0 {
   527  			if hit, ok = n.children[len(n.children)-1].iterate(dir, start, stop, includeStart, hit, iter); !ok {
   528  				return hit, false
   529  			}
   530  		}
   531  	case descend:
   532  		if start.valid {
   533  			index, found = n.items.find(start.item, n.cow.less)
   534  			if !found {
   535  				index = index - 1
   536  			}
   537  		} else {
   538  			index = len(n.items) - 1
   539  		}
   540  		for i := index; i >= 0; i-- {
   541  			if start.valid && !n.cow.less(n.items[i], start.item) {
   542  				if !includeStart || hit || n.cow.less(start.item, n.items[i]) {
   543  					continue
   544  				}
   545  			}
   546  			if len(n.children) > 0 {
   547  				if hit, ok = n.children[i+1].iterate(dir, start, stop, includeStart, hit, iter); !ok {
   548  					return hit, false
   549  				}
   550  			}
   551  			if stop.valid && !n.cow.less(stop.item, n.items[i]) {
   552  				return hit, false //	continue
   553  			}
   554  			hit = true
   555  			if !iter(n.items[i]) {
   556  				return hit, false
   557  			}
   558  		}
   559  		if len(n.children) > 0 {
   560  			if hit, ok = n.children[0].iterate(dir, start, stop, includeStart, hit, iter); !ok {
   561  				return hit, false
   562  			}
   563  		}
   564  	}
   565  	return hit, true
   566  }
   567  
   568  // print is used for testing/debugging purposes.
   569  func (n *node[T]) print(w io.Writer, level int) {
   570  	fmt.Fprintf(w, "%sNODE:%v\n", strings.Repeat("  ", level), n.items)
   571  	for _, c := range n.children {
   572  		c.print(w, level+1)
   573  	}
   574  }
   575  
   576  // BTreeG is a generic implementation of a B-Tree.
   577  //
   578  // BTreeG stores items of type T in an ordered structure, allowing easy insertion,
   579  // removal, and iteration.
   580  //
   581  // Write operations are not safe for concurrent mutation by multiple
   582  // goroutines, but Read operations are.
   583  type BTreeG[T any] struct {
   584  	degree int
   585  	length int
   586  	root   *node[T]
   587  	cow    *copyOnWriteContext[T]
   588  }
   589  
   590  // LessFunc[T] determines how to order a type 'T'.  It should implement a strict
   591  // ordering, and should return true if within that ordering, 'a' < 'b'.
   592  type LessFunc[T any] func(a, b T) bool
   593  
   594  // copyOnWriteContext pointers determine node ownership... a tree with a write
   595  // context equivalent to a node's write context is allowed to modify that node.
   596  // A tree whose write context does not match a node's is not allowed to modify
   597  // it, and must create a new, writable copy (IE: it's a Clone).
   598  //
   599  // When doing any write operation, we maintain the invariant that the current
   600  // node's context is equal to the context of the tree that requested the write.
   601  // We do this by, before we descend into any node, creating a copy with the
   602  // correct context if the contexts don't match.
   603  //
   604  // Since the node we're currently visiting on any write has the requesting
   605  // tree's context, that node is modifiable in place.  Children of that node may
   606  // not share context, but before we descend into them, we'll make a mutable
   607  // copy.
   608  type copyOnWriteContext[T any] struct {
   609  	freelist *FreeListG[T]
   610  	less     LessFunc[T]
   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 *BTreeG[T]) Clone() (t2 *BTreeG[T]) {
   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 *BTreeG[T]) 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 *BTreeG[T]) minItems() int {
   645  	return t.degree - 1
   646  }
   647  
   648  func (c *copyOnWriteContext[T]) newNode() (n *node[T]) {
   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[T]) freeNode(n *node[T]) 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  // and the second return value is true.  Otherwise, (zeroValue, false)
   684  //
   685  // nil cannot be added to the tree (will panic).
   686  func (t *BTreeG[T]) ReplaceOrInsert(item T) (_ T, _ bool) {
   687  	if t.root == nil {
   688  		t.root = t.cow.newNode()
   689  		t.root.items = append(t.root.items, item)
   690  		t.length++
   691  		return
   692  	} else {
   693  		t.root = t.root.mutableFor(t.cow)
   694  		if len(t.root.items) >= t.maxItems() {
   695  			item2, second := t.root.split(t.maxItems() / 2)
   696  			oldroot := t.root
   697  			t.root = t.cow.newNode()
   698  			t.root.items = append(t.root.items, item2)
   699  			t.root.children = append(t.root.children, oldroot, second)
   700  		}
   701  	}
   702  	out, outb := t.root.insert(item, t.maxItems())
   703  	if !outb {
   704  		t.length++
   705  	}
   706  	return out, outb
   707  }
   708  
   709  // Delete removes an item equal to the passed in item from the tree, returning
   710  // it.  If no such item exists, returns (zeroValue, false).
   711  func (t *BTreeG[T]) Delete(item T) (T, bool) {
   712  	return t.deleteItem(item, removeItem)
   713  }
   714  
   715  // DeleteMin removes the smallest item in the tree and returns it.
   716  // If no such item exists, returns (zeroValue, false).
   717  func (t *BTreeG[T]) DeleteMin() (T, bool) {
   718  	var zero T
   719  	return t.deleteItem(zero, removeMin)
   720  }
   721  
   722  // DeleteMax removes the largest item in the tree and returns it.
   723  // If no such item exists, returns (zeroValue, false).
   724  func (t *BTreeG[T]) DeleteMax() (T, bool) {
   725  	var zero T
   726  	return t.deleteItem(zero, removeMax)
   727  }
   728  
   729  func (t *BTreeG[T]) deleteItem(item T, typ toRemove) (_ T, _ bool) {
   730  	if t.root == nil || len(t.root.items) == 0 {
   731  		return
   732  	}
   733  	t.root = t.root.mutableFor(t.cow)
   734  	out, outb := t.root.remove(item, t.minItems(), typ)
   735  	if len(t.root.items) == 0 && len(t.root.children) > 0 {
   736  		oldroot := t.root
   737  		t.root = t.root.children[0]
   738  		t.cow.freeNode(oldroot)
   739  	}
   740  	if outb {
   741  		t.length--
   742  	}
   743  	return out, outb
   744  }
   745  
   746  // AscendRange calls the iterator for every value in the tree within the range
   747  // [greaterOrEqual, lessThan), until iterator returns false.
   748  func (t *BTreeG[T]) AscendRange(greaterOrEqual, lessThan T, iterator ItemIteratorG[T]) {
   749  	if t.root == nil {
   750  		return
   751  	}
   752  	t.root.iterate(ascend, optional[T](greaterOrEqual), optional[T](lessThan), true, false, iterator)
   753  }
   754  
   755  // AscendLessThan calls the iterator for every value in the tree within the range
   756  // [first, pivot), until iterator returns false.
   757  func (t *BTreeG[T]) AscendLessThan(pivot T, iterator ItemIteratorG[T]) {
   758  	if t.root == nil {
   759  		return
   760  	}
   761  	t.root.iterate(ascend, empty[T](), optional(pivot), false, false, iterator)
   762  }
   763  
   764  // AscendGreaterOrEqual calls the iterator for every value in the tree within
   765  // the range [pivot, last], until iterator returns false.
   766  func (t *BTreeG[T]) AscendGreaterOrEqual(pivot T, iterator ItemIteratorG[T]) {
   767  	if t.root == nil {
   768  		return
   769  	}
   770  	t.root.iterate(ascend, optional[T](pivot), empty[T](), true, false, iterator)
   771  }
   772  
   773  // Ascend calls the iterator for every value in the tree within the range
   774  // [first, last], until iterator returns false.
   775  func (t *BTreeG[T]) Ascend(iterator ItemIteratorG[T]) {
   776  	if t.root == nil {
   777  		return
   778  	}
   779  	t.root.iterate(ascend, empty[T](), empty[T](), false, false, iterator)
   780  }
   781  
   782  // DescendRange calls the iterator for every value in the tree within the range
   783  // [lessOrEqual, greaterThan), until iterator returns false.
   784  func (t *BTreeG[T]) DescendRange(lessOrEqual, greaterThan T, iterator ItemIteratorG[T]) {
   785  	if t.root == nil {
   786  		return
   787  	}
   788  	t.root.iterate(descend, optional[T](lessOrEqual), optional[T](greaterThan), true, false, iterator)
   789  }
   790  
   791  // DescendLessOrEqual calls the iterator for every value in the tree within the range
   792  // [pivot, first], until iterator returns false.
   793  func (t *BTreeG[T]) DescendLessOrEqual(pivot T, iterator ItemIteratorG[T]) {
   794  	if t.root == nil {
   795  		return
   796  	}
   797  	t.root.iterate(descend, optional[T](pivot), empty[T](), true, false, iterator)
   798  }
   799  
   800  // DescendGreaterThan calls the iterator for every value in the tree within
   801  // the range [last, pivot), until iterator returns false.
   802  func (t *BTreeG[T]) DescendGreaterThan(pivot T, iterator ItemIteratorG[T]) {
   803  	if t.root == nil {
   804  		return
   805  	}
   806  	t.root.iterate(descend, empty[T](), optional[T](pivot), false, false, iterator)
   807  }
   808  
   809  // Descend calls the iterator for every value in the tree within the range
   810  // [last, first], until iterator returns false.
   811  func (t *BTreeG[T]) Descend(iterator ItemIteratorG[T]) {
   812  	if t.root == nil {
   813  		return
   814  	}
   815  	t.root.iterate(descend, empty[T](), empty[T](), false, false, iterator)
   816  }
   817  
   818  // Get looks for the key item in the tree, returning it.  It returns
   819  // (zeroValue, false) if unable to find that item.
   820  func (t *BTreeG[T]) Get(key T) (_ T, _ bool) {
   821  	if t.root == nil {
   822  		return
   823  	}
   824  	return t.root.get(key)
   825  }
   826  
   827  // Min returns the smallest item in the tree, or (zeroValue, false) if the tree is empty.
   828  func (t *BTreeG[T]) Min() (_ T, _ bool) {
   829  	return min(t.root)
   830  }
   831  
   832  // Max returns the largest item in the tree, or (zeroValue, false) if the tree is empty.
   833  func (t *BTreeG[T]) Max() (_ T, _ bool) {
   834  	return max(t.root)
   835  }
   836  
   837  // Has returns true if the given key is in the tree.
   838  func (t *BTreeG[T]) Has(key T) bool {
   839  	_, ok := t.Get(key)
   840  	return ok
   841  }
   842  
   843  // Len returns the number of items currently in the tree.
   844  func (t *BTreeG[T]) 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 *BTreeG[T]) 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[T]) reset(c *copyOnWriteContext[T]) 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  
   895  // BTree is an implementation of a B-Tree.
   896  //
   897  // BTree stores Item instances in an ordered structure, allowing easy insertion,
   898  // removal, and iteration.
   899  //
   900  // Write operations are not safe for concurrent mutation by multiple
   901  // goroutines, but Read operations are.
   902  type BTree BTreeG[Item]
   903  
   904  var itemLess LessFunc[Item] = func(a, b Item) bool {
   905  	return a.Less(b)
   906  }
   907  
   908  // New creates a new B-Tree with the given degree.
   909  //
   910  // New(2), for example, will create a 2-3-4 tree (each node contains 1-3 items
   911  // and 2-4 children).
   912  func New(degree int) *BTree {
   913  	return (*BTree)(NewG[Item](degree, itemLess))
   914  }
   915  
   916  // FreeList represents a free list of btree nodes. By default each
   917  // BTree has its own FreeList, but multiple BTrees can share the same
   918  // FreeList.
   919  // Two Btrees using the same freelist are safe for concurrent write access.
   920  type FreeList FreeListG[Item]
   921  
   922  // NewFreeList creates a new free list.
   923  // size is the maximum size of the returned free list.
   924  func NewFreeList(size int) *FreeList {
   925  	return (*FreeList)(NewFreeListG[Item](size))
   926  }
   927  
   928  // NewWithFreeList creates a new B-Tree that uses the given node free list.
   929  func NewWithFreeList(degree int, f *FreeList) *BTree {
   930  	return (*BTree)(NewWithFreeListG[Item](degree, itemLess, (*FreeListG[Item])(f)))
   931  }
   932  
   933  // ItemIterator allows callers of Ascend* to iterate in-order over portions of
   934  // the tree.  When this function returns false, iteration will stop and the
   935  // associated Ascend* function will immediately return.
   936  type ItemIterator ItemIteratorG[Item]
   937  
   938  // Clone clones the btree, lazily.  Clone should not be called concurrently,
   939  // but the original tree (t) and the new tree (t2) can be used concurrently
   940  // once the Clone call completes.
   941  //
   942  // The internal tree structure of b is marked read-only and shared between t and
   943  // t2.  Writes to both t and t2 use copy-on-write logic, creating new nodes
   944  // whenever one of b's original nodes would have been modified.  Read operations
   945  // should have no performance degredation.  Write operations for both t and t2
   946  // will initially experience minor slow-downs caused by additional allocs and
   947  // copies due to the aforementioned copy-on-write logic, but should converge to
   948  // the original performance characteristics of the original tree.
   949  func (t *BTree) Clone() (t2 *BTree) {
   950  	return (*BTree)((*BTreeG[Item])(t).Clone())
   951  }
   952  
   953  // Delete removes an item equal to the passed in item from the tree, returning
   954  // it.  If no such item exists, returns nil.
   955  func (t *BTree) Delete(item Item) Item {
   956  	i, _ := (*BTreeG[Item])(t).Delete(item)
   957  	return i
   958  }
   959  
   960  // DeleteMax removes the largest item in the tree and returns it.
   961  // If no such item exists, returns nil.
   962  func (t *BTree) DeleteMax() Item {
   963  	i, _ := (*BTreeG[Item])(t).DeleteMax()
   964  	return i
   965  }
   966  
   967  // DeleteMin removes the smallest item in the tree and returns it.
   968  // If no such item exists, returns nil.
   969  func (t *BTree) DeleteMin() Item {
   970  	i, _ := (*BTreeG[Item])(t).DeleteMin()
   971  	return i
   972  }
   973  
   974  // Get looks for the key item in the tree, returning it.  It returns nil if
   975  // unable to find that item.
   976  func (t *BTree) Get(key Item) Item {
   977  	i, _ := (*BTreeG[Item])(t).Get(key)
   978  	return i
   979  }
   980  
   981  // Max returns the largest item in the tree, or nil if the tree is empty.
   982  func (t *BTree) Max() Item {
   983  	i, _ := (*BTreeG[Item])(t).Max()
   984  	return i
   985  }
   986  
   987  // Min returns the smallest item in the tree, or nil if the tree is empty.
   988  func (t *BTree) Min() Item {
   989  	i, _ := (*BTreeG[Item])(t).Min()
   990  	return i
   991  }
   992  
   993  // Has returns true if the given key is in the tree.
   994  func (t *BTree) Has(key Item) bool {
   995  	return (*BTreeG[Item])(t).Has(key)
   996  }
   997  
   998  // ReplaceOrInsert adds the given item to the tree.  If an item in the tree
   999  // already equals the given one, it is removed from the tree and returned.
  1000  // Otherwise, nil is returned.
  1001  //
  1002  // nil cannot be added to the tree (will panic).
  1003  func (t *BTree) ReplaceOrInsert(item Item) Item {
  1004  	i, _ := (*BTreeG[Item])(t).ReplaceOrInsert(item)
  1005  	return i
  1006  }
  1007  
  1008  // AscendRange calls the iterator for every value in the tree within the range
  1009  // [greaterOrEqual, lessThan), until iterator returns false.
  1010  func (t *BTree) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) {
  1011  	(*BTreeG[Item])(t).AscendRange(greaterOrEqual, lessThan, (ItemIteratorG[Item])(iterator))
  1012  }
  1013  
  1014  // AscendLessThan calls the iterator for every value in the tree within the range
  1015  // [first, pivot), until iterator returns false.
  1016  func (t *BTree) AscendLessThan(pivot Item, iterator ItemIterator) {
  1017  	(*BTreeG[Item])(t).AscendLessThan(pivot, (ItemIteratorG[Item])(iterator))
  1018  }
  1019  
  1020  // AscendGreaterOrEqual calls the iterator for every value in the tree within
  1021  // the range [pivot, last], until iterator returns false.
  1022  func (t *BTree) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {
  1023  	(*BTreeG[Item])(t).AscendGreaterOrEqual(pivot, (ItemIteratorG[Item])(iterator))
  1024  }
  1025  
  1026  // Ascend calls the iterator for every value in the tree within the range
  1027  // [first, last], until iterator returns false.
  1028  func (t *BTree) Ascend(iterator ItemIterator) {
  1029  	(*BTreeG[Item])(t).Ascend((ItemIteratorG[Item])(iterator))
  1030  }
  1031  
  1032  // DescendRange calls the iterator for every value in the tree within the range
  1033  // [lessOrEqual, greaterThan), until iterator returns false.
  1034  func (t *BTree) DescendRange(lessOrEqual, greaterThan Item, iterator ItemIterator) {
  1035  	(*BTreeG[Item])(t).DescendRange(lessOrEqual, greaterThan, (ItemIteratorG[Item])(iterator))
  1036  }
  1037  
  1038  // DescendLessOrEqual calls the iterator for every value in the tree within the range
  1039  // [pivot, first], until iterator returns false.
  1040  func (t *BTree) DescendLessOrEqual(pivot Item, iterator ItemIterator) {
  1041  	(*BTreeG[Item])(t).DescendLessOrEqual(pivot, (ItemIteratorG[Item])(iterator))
  1042  }
  1043  
  1044  // DescendGreaterThan calls the iterator for every value in the tree within
  1045  // the range [last, pivot), until iterator returns false.
  1046  func (t *BTree) DescendGreaterThan(pivot Item, iterator ItemIterator) {
  1047  	(*BTreeG[Item])(t).DescendGreaterThan(pivot, (ItemIteratorG[Item])(iterator))
  1048  }
  1049  
  1050  // Descend calls the iterator for every value in the tree within the range
  1051  // [last, first], until iterator returns false.
  1052  func (t *BTree) Descend(iterator ItemIterator) {
  1053  	(*BTreeG[Item])(t).Descend((ItemIteratorG[Item])(iterator))
  1054  }
  1055  
  1056  // Len returns the number of items currently in the tree.
  1057  func (t *BTree) Len() int {
  1058  	return (*BTreeG[Item])(t).Len()
  1059  }
  1060  
  1061  // Clear removes all items from the btree.  If addNodesToFreelist is true,
  1062  // t's nodes are added to its freelist as part of this call, until the freelist
  1063  // is full.  Otherwise, the root node is simply dereferenced and the subtree
  1064  // left to Go's normal GC processes.
  1065  //
  1066  // This can be much faster
  1067  // than calling Delete on all elements, because that requires finding/removing
  1068  // each element in the tree and updating the tree accordingly.  It also is
  1069  // somewhat faster than creating a new tree to replace the old one, because
  1070  // nodes from the old tree are reclaimed into the freelist for use by the new
  1071  // one, instead of being lost to the garbage collector.
  1072  //
  1073  // This call takes:
  1074  //   O(1): when addNodesToFreelist is false, this is a single operation.
  1075  //   O(1): when the freelist is already full, it breaks out immediately
  1076  //   O(freelist size):  when the freelist is empty and the nodes are all owned
  1077  //       by this tree, nodes are added to the freelist until full.
  1078  //   O(tree size):  when all nodes are owned by another tree, all nodes are
  1079  //       iterated over looking for nodes to add to the freelist, and due to
  1080  //       ownership, none are.
  1081  func (t *BTree) Clear(addNodesToFreelist bool) {
  1082  	(*BTreeG[Item])(t).Clear(addNodesToFreelist)
  1083  }
  1084  

View as plain text