btree_generic.go raw

   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