blob: e44a0f48804a76d97c00adacc5d20b5107ca54c5 [file] [log] [blame]
Abhay Kumara2ae5992025-11-10 14:02:24 +00001// 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.
61package btree
62
63import (
64 "fmt"
65 "io"
66 "sort"
67 "strings"
68 "sync"
69)
70
71// Item represents a single object in the tree.
72type 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
81const (
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.
89type 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.
96func NewFreeListG[T any](size int) *FreeListG[T] {
97 return &FreeListG[T]{freelist: make([]*node[T], 0, size)}
98}
99
100func (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
114func (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.
127type ItemIteratorG[T any] func(item T) bool
128
129// Ordered represents the set of types for which the '<' operator work.
130type 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.
135func 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.
140func 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.
150func 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.
155func 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.
166type items[T any] []T
167
168// insertAt inserts a value into the given index, pushing all subsequent values
169// forward.
170func (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.
181func (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.
191func (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.
202func (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.
214func (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
229type node[T any] struct {
230 items items[T]
231 children items[*node[T]]
232 cow *copyOnWriteContext[T]
233}
234
235func (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
256func (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.
265func (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.
279func (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.
293func (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.
321func (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.
332func 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.
346func 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.
360type toRemove int
361
362const (
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.
369func (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.
437func (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
474type direction int
475
476const (
477 descend = direction(-1)
478 ascend = direction(+1)
479)
480
481type optionalItem[T any] struct {
482 item T
483 valid bool
484}
485
486func optional[T any](item T) optionalItem[T] {
487 return optionalItem[T]{item: item, valid: true}
488}
489func 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.
500func (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.
569func (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.
583type 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'.
592type 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.
608type 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.
624func (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.
638func (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).
644func (t *BTreeG[T]) minItems() int {
645 return t.degree - 1
646}
647
648func (c *copyOnWriteContext[T]) newNode() (n *node[T]) {
649 n = c.freelist.newNode()
650 n.cow = c
651 return
652}
653
654type freeType int
655
656const (
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).
665func (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).
686func (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).
711func (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).
717func (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).
724func (t *BTreeG[T]) DeleteMax() (T, bool) {
725 var zero T
726 return t.deleteItem(zero, removeMax)
727}
728
729func (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.
748func (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.
757func (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.
766func (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.
775func (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.
784func (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.
793func (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.
802func (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.
811func (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.
820func (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.
828func (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.
833func (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.
838func (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.
844func (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.
868func (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.
878func (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.
888type Int int
889
890// Less returns true if int(a) < int(b).
891func (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.
902type BTree BTreeG[Item]
903
904var 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).
912func 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.
920type FreeList FreeListG[Item]
921
922// NewFreeList creates a new free list.
923// size is the maximum size of the returned free list.
924func 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.
929func 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.
936type 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.
949func (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.
955func (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.
962func (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.
969func (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.
976func (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.
982func (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.
988func (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.
994func (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).
1003func (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.
1010func (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.
1016func (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.
1022func (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.
1028func (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.
1034func (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.
1040func (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.
1046func (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.
1052func (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.
1057func (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.
1081func (t *BTree) Clear(addNodesToFreelist bool) {
1082 (*BTreeG[Item])(t).Clear(addNodesToFreelist)
1083}