123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083 |
- // Copyright 2014-2022 Google Inc.
- //
- // Licensed under the Apache License, Version 2.0 (the "License");
- // you may not use this file except in compliance with the License.
- // You may obtain a copy of the License at
- //
- // http://www.apache.org/licenses/LICENSE-2.0
- //
- // Unless required by applicable law or agreed to in writing, software
- // distributed under the License is distributed on an "AS IS" BASIS,
- // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- // See the License for the specific language governing permissions and
- // limitations under the License.
- //go:build go1.18
- // +build go1.18
- // In Go 1.18 and beyond, a BTreeG generic is created, and BTree is a specific
- // instantiation of that generic for the Item interface, with a backwards-
- // compatible API. Before go1.18, generics are not supported,
- // and BTree is just an implementation based around the Item interface.
- // Package btree implements in-memory B-Trees of arbitrary degree.
- //
- // btree implements an in-memory B-Tree for use as an ordered data structure.
- // It is not meant for persistent storage solutions.
- //
- // It has a flatter structure than an equivalent red-black or other binary tree,
- // which in some cases yields better memory usage and/or performance.
- // See some discussion on the matter here:
- // http://google-opensource.blogspot.com/2013/01/c-containers-that-save-memory-and-time.html
- // Note, though, that this project is in no way related to the C++ B-Tree
- // implementation written about there.
- //
- // Within this tree, each node contains a slice of items and a (possibly nil)
- // slice of children. For basic numeric values or raw structs, this can cause
- // efficiency differences when compared to equivalent C++ template code that
- // stores values in arrays within the node:
- // * Due to the overhead of storing values as interfaces (each
- // value needs to be stored as the value itself, then 2 words for the
- // interface pointing to that value and its type), resulting in higher
- // memory use.
- // * Since interfaces can point to values anywhere in memory, values are
- // most likely not stored in contiguous blocks, resulting in a higher
- // number of cache misses.
- // These issues don't tend to matter, though, when working with strings or other
- // heap-allocated structures, since C++-equivalent structures also must store
- // pointers and also distribute their values across the heap.
- //
- // This implementation is designed to be a drop-in replacement to gollrb.LLRB
- // trees, (http://github.com/petar/gollrb), an excellent and probably the most
- // widely used ordered tree implementation in the Go ecosystem currently.
- // Its functions, therefore, exactly mirror those of
- // llrb.LLRB where possible. Unlike gollrb, though, we currently don't
- // support storing multiple equivalent values.
- //
- // There are two implementations; those suffixed with 'G' are generics, usable
- // for any type, and require a passed-in "less" function to define their ordering.
- // Those without this prefix are specific to the 'Item' interface, and use
- // its 'Less' function for ordering.
- package btree
- import (
- "fmt"
- "io"
- "sort"
- "strings"
- "sync"
- )
- // Item represents a single object in the tree.
- type Item interface {
- // Less tests whether the current item is less than the given argument.
- //
- // This must provide a strict weak ordering.
- // If !a.Less(b) && !b.Less(a), we treat this to mean a == b (i.e. we can only
- // hold one of either a or b in the tree).
- Less(than Item) bool
- }
- const (
- DefaultFreeListSize = 32
- )
- // FreeListG represents a free list of btree nodes. By default each
- // BTree has its own FreeList, but multiple BTrees can share the same
- // FreeList, in particular when they're created with Clone.
- // Two Btrees using the same freelist are safe for concurrent write access.
- type FreeListG[T any] struct {
- mu sync.Mutex
- freelist []*node[T]
- }
- // NewFreeListG creates a new free list.
- // size is the maximum size of the returned free list.
- func NewFreeListG[T any](size int) *FreeListG[T] {
- return &FreeListG[T]{freelist: make([]*node[T], 0, size)}
- }
- func (f *FreeListG[T]) newNode() (n *node[T]) {
- f.mu.Lock()
- index := len(f.freelist) - 1
- if index < 0 {
- f.mu.Unlock()
- return new(node[T])
- }
- n = f.freelist[index]
- f.freelist[index] = nil
- f.freelist = f.freelist[:index]
- f.mu.Unlock()
- return
- }
- func (f *FreeListG[T]) freeNode(n *node[T]) (out bool) {
- f.mu.Lock()
- if len(f.freelist) < cap(f.freelist) {
- f.freelist = append(f.freelist, n)
- out = true
- }
- f.mu.Unlock()
- return
- }
- // ItemIteratorG allows callers of {A/De}scend* to iterate in-order over portions of
- // the tree. When this function returns false, iteration will stop and the
- // associated Ascend* function will immediately return.
- type ItemIteratorG[T any] func(item T) bool
- // Ordered represents the set of types for which the '<' operator work.
- type Ordered interface {
- ~int | ~int8 | ~int16 | ~int32 | ~int64 | ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~float32 | ~float64 | ~string
- }
- // Less[T] returns a default LessFunc that uses the '<' operator for types that support it.
- func Less[T Ordered]() LessFunc[T] {
- return func(a, b T) bool { return a < b }
- }
- // NewOrderedG creates a new B-Tree for ordered types.
- func NewOrderedG[T Ordered](degree int) *BTreeG[T] {
- return NewG[T](degree, Less[T]())
- }
- // NewG creates a new B-Tree with the given degree.
- //
- // NewG(2), for example, will create a 2-3-4 tree (each node contains 1-3 items
- // and 2-4 children).
- //
- // The passed-in LessFunc determines how objects of type T are ordered.
- func NewG[T any](degree int, less LessFunc[T]) *BTreeG[T] {
- return NewWithFreeListG(degree, less, NewFreeListG[T](DefaultFreeListSize))
- }
- // NewWithFreeListG creates a new B-Tree that uses the given node free list.
- func NewWithFreeListG[T any](degree int, less LessFunc[T], f *FreeListG[T]) *BTreeG[T] {
- if degree <= 1 {
- panic("bad degree")
- }
- return &BTreeG[T]{
- degree: degree,
- cow: ©OnWriteContext[T]{freelist: f, less: less},
- }
- }
- // items stores items in a node.
- type items[T any] []T
- // insertAt inserts a value into the given index, pushing all subsequent values
- // forward.
- func (s *items[T]) insertAt(index int, item T) {
- var zero T
- *s = append(*s, zero)
- if index < len(*s) {
- copy((*s)[index+1:], (*s)[index:])
- }
- (*s)[index] = item
- }
- // removeAt removes a value at a given index, pulling all subsequent values
- // back.
- func (s *items[T]) removeAt(index int) T {
- item := (*s)[index]
- copy((*s)[index:], (*s)[index+1:])
- var zero T
- (*s)[len(*s)-1] = zero
- *s = (*s)[:len(*s)-1]
- return item
- }
- // pop removes and returns the last element in the list.
- func (s *items[T]) pop() (out T) {
- index := len(*s) - 1
- out = (*s)[index]
- var zero T
- (*s)[index] = zero
- *s = (*s)[:index]
- return
- }
- // truncate truncates this instance at index so that it contains only the
- // first index items. index must be less than or equal to length.
- func (s *items[T]) truncate(index int) {
- var toClear items[T]
- *s, toClear = (*s)[:index], (*s)[index:]
- var zero T
- for i := 0; i < len(toClear); i++ {
- toClear[i] = zero
- }
- }
- // find returns the index where the given item should be inserted into this
- // list. 'found' is true if the item already exists in the list at the given
- // index.
- func (s items[T]) find(item T, less func(T, T) bool) (index int, found bool) {
- i := sort.Search(len(s), func(i int) bool {
- return less(item, s[i])
- })
- if i > 0 && !less(s[i-1], item) {
- return i - 1, true
- }
- return i, false
- }
- // node is an internal node in a tree.
- //
- // It must at all times maintain the invariant that either
- // * len(children) == 0, len(items) unconstrained
- // * len(children) == len(items) + 1
- type node[T any] struct {
- items items[T]
- children items[*node[T]]
- cow *copyOnWriteContext[T]
- }
- func (n *node[T]) mutableFor(cow *copyOnWriteContext[T]) *node[T] {
- if n.cow == cow {
- return n
- }
- out := cow.newNode()
- if cap(out.items) >= len(n.items) {
- out.items = out.items[:len(n.items)]
- } else {
- out.items = make(items[T], len(n.items), cap(n.items))
- }
- copy(out.items, n.items)
- // Copy children
- if cap(out.children) >= len(n.children) {
- out.children = out.children[:len(n.children)]
- } else {
- out.children = make(items[*node[T]], len(n.children), cap(n.children))
- }
- copy(out.children, n.children)
- return out
- }
- func (n *node[T]) mutableChild(i int) *node[T] {
- c := n.children[i].mutableFor(n.cow)
- n.children[i] = c
- return c
- }
- // split splits the given node at the given index. The current node shrinks,
- // and this function returns the item that existed at that index and a new node
- // containing all items/children after it.
- func (n *node[T]) split(i int) (T, *node[T]) {
- item := n.items[i]
- next := n.cow.newNode()
- next.items = append(next.items, n.items[i+1:]...)
- n.items.truncate(i)
- if len(n.children) > 0 {
- next.children = append(next.children, n.children[i+1:]...)
- n.children.truncate(i + 1)
- }
- return item, next
- }
- // maybeSplitChild checks if a child should be split, and if so splits it.
- // Returns whether or not a split occurred.
- func (n *node[T]) maybeSplitChild(i, maxItems int) bool {
- if len(n.children[i].items) < maxItems {
- return false
- }
- first := n.mutableChild(i)
- item, second := first.split(maxItems / 2)
- n.items.insertAt(i, item)
- n.children.insertAt(i+1, second)
- return true
- }
- // insert inserts an item into the subtree rooted at this node, making sure
- // no nodes in the subtree exceed maxItems items. Should an equivalent item be
- // be found/replaced by insert, it will be returned.
- func (n *node[T]) insert(item T, maxItems int) (_ T, _ bool) {
- i, found := n.items.find(item, n.cow.less)
- if found {
- out := n.items[i]
- n.items[i] = item
- return out, true
- }
- if len(n.children) == 0 {
- n.items.insertAt(i, item)
- return
- }
- if n.maybeSplitChild(i, maxItems) {
- inTree := n.items[i]
- switch {
- case n.cow.less(item, inTree):
- // no change, we want first split node
- case n.cow.less(inTree, item):
- i++ // we want second split node
- default:
- out := n.items[i]
- n.items[i] = item
- return out, true
- }
- }
- return n.mutableChild(i).insert(item, maxItems)
- }
- // get finds the given key in the subtree and returns it.
- func (n *node[T]) get(key T) (_ T, _ bool) {
- i, found := n.items.find(key, n.cow.less)
- if found {
- return n.items[i], true
- } else if len(n.children) > 0 {
- return n.children[i].get(key)
- }
- return
- }
- // min returns the first item in the subtree.
- func min[T any](n *node[T]) (_ T, found bool) {
- if n == nil {
- return
- }
- for len(n.children) > 0 {
- n = n.children[0]
- }
- if len(n.items) == 0 {
- return
- }
- return n.items[0], true
- }
- // max returns the last item in the subtree.
- func max[T any](n *node[T]) (_ T, found bool) {
- if n == nil {
- return
- }
- for len(n.children) > 0 {
- n = n.children[len(n.children)-1]
- }
- if len(n.items) == 0 {
- return
- }
- return n.items[len(n.items)-1], true
- }
- // toRemove details what item to remove in a node.remove call.
- type toRemove int
- const (
- removeItem toRemove = iota // removes the given item
- removeMin // removes smallest item in the subtree
- removeMax // removes largest item in the subtree
- )
- // remove removes an item from the subtree rooted at this node.
- func (n *node[T]) remove(item T, minItems int, typ toRemove) (_ T, _ bool) {
- var i int
- var found bool
- switch typ {
- case removeMax:
- if len(n.children) == 0 {
- return n.items.pop(), true
- }
- i = len(n.items)
- case removeMin:
- if len(n.children) == 0 {
- return n.items.removeAt(0), true
- }
- i = 0
- case removeItem:
- i, found = n.items.find(item, n.cow.less)
- if len(n.children) == 0 {
- if found {
- return n.items.removeAt(i), true
- }
- return
- }
- default:
- panic("invalid type")
- }
- // If we get to here, we have children.
- if len(n.children[i].items) <= minItems {
- return n.growChildAndRemove(i, item, minItems, typ)
- }
- child := n.mutableChild(i)
- // Either we had enough items to begin with, or we've done some
- // merging/stealing, because we've got enough now and we're ready to return
- // stuff.
- if found {
- // The item exists at index 'i', and the child we've selected can give us a
- // predecessor, since if we've gotten here it's got > minItems items in it.
- out := n.items[i]
- // We use our special-case 'remove' call with typ=maxItem to pull the
- // predecessor of item i (the rightmost leaf of our immediate left child)
- // and set it into where we pulled the item from.
- var zero T
- n.items[i], _ = child.remove(zero, minItems, removeMax)
- return out, true
- }
- // Final recursive call. Once we're here, we know that the item isn't in this
- // node and that the child is big enough to remove from.
- return child.remove(item, minItems, typ)
- }
- // growChildAndRemove grows child 'i' to make sure it's possible to remove an
- // item from it while keeping it at minItems, then calls remove to actually
- // remove it.
- //
- // Most documentation says we have to do two sets of special casing:
- // 1) item is in this node
- // 2) item is in child
- // In both cases, we need to handle the two subcases:
- // A) node has enough values that it can spare one
- // B) node doesn't have enough values
- // For the latter, we have to check:
- // a) left sibling has node to spare
- // b) right sibling has node to spare
- // c) we must merge
- // To simplify our code here, we handle cases #1 and #2 the same:
- // If a node doesn't have enough items, we make sure it does (using a,b,c).
- // We then simply redo our remove call, and the second time (regardless of
- // whether we're in case 1 or 2), we'll have enough items and can guarantee
- // that we hit case A.
- func (n *node[T]) growChildAndRemove(i int, item T, minItems int, typ toRemove) (T, bool) {
- if i > 0 && len(n.children[i-1].items) > minItems {
- // Steal from left child
- child := n.mutableChild(i)
- stealFrom := n.mutableChild(i - 1)
- stolenItem := stealFrom.items.pop()
- child.items.insertAt(0, n.items[i-1])
- n.items[i-1] = stolenItem
- if len(stealFrom.children) > 0 {
- child.children.insertAt(0, stealFrom.children.pop())
- }
- } else if i < len(n.items) && len(n.children[i+1].items) > minItems {
- // steal from right child
- child := n.mutableChild(i)
- stealFrom := n.mutableChild(i + 1)
- stolenItem := stealFrom.items.removeAt(0)
- child.items = append(child.items, n.items[i])
- n.items[i] = stolenItem
- if len(stealFrom.children) > 0 {
- child.children = append(child.children, stealFrom.children.removeAt(0))
- }
- } else {
- if i >= len(n.items) {
- i--
- }
- child := n.mutableChild(i)
- // merge with right child
- mergeItem := n.items.removeAt(i)
- mergeChild := n.children.removeAt(i + 1)
- child.items = append(child.items, mergeItem)
- child.items = append(child.items, mergeChild.items...)
- child.children = append(child.children, mergeChild.children...)
- n.cow.freeNode(mergeChild)
- }
- return n.remove(item, minItems, typ)
- }
- type direction int
- const (
- descend = direction(-1)
- ascend = direction(+1)
- )
- type optionalItem[T any] struct {
- item T
- valid bool
- }
- func optional[T any](item T) optionalItem[T] {
- return optionalItem[T]{item: item, valid: true}
- }
- func empty[T any]() optionalItem[T] {
- return optionalItem[T]{}
- }
- // iterate provides a simple method for iterating over elements in the tree.
- //
- // When ascending, the 'start' should be less than 'stop' and when descending,
- // the 'start' should be greater than 'stop'. Setting 'includeStart' to true
- // will force the iterator to include the first item when it equals 'start',
- // thus creating a "greaterOrEqual" or "lessThanEqual" rather than just a
- // "greaterThan" or "lessThan" queries.
- func (n *node[T]) iterate(dir direction, start, stop optionalItem[T], includeStart bool, hit bool, iter ItemIteratorG[T]) (bool, bool) {
- var ok, found bool
- var index int
- switch dir {
- case ascend:
- if start.valid {
- index, _ = n.items.find(start.item, n.cow.less)
- }
- for i := index; i < len(n.items); i++ {
- if len(n.children) > 0 {
- if hit, ok = n.children[i].iterate(dir, start, stop, includeStart, hit, iter); !ok {
- return hit, false
- }
- }
- if !includeStart && !hit && start.valid && !n.cow.less(start.item, n.items[i]) {
- hit = true
- continue
- }
- hit = true
- if stop.valid && !n.cow.less(n.items[i], stop.item) {
- return hit, false
- }
- if !iter(n.items[i]) {
- return hit, false
- }
- }
- if len(n.children) > 0 {
- if hit, ok = n.children[len(n.children)-1].iterate(dir, start, stop, includeStart, hit, iter); !ok {
- return hit, false
- }
- }
- case descend:
- if start.valid {
- index, found = n.items.find(start.item, n.cow.less)
- if !found {
- index = index - 1
- }
- } else {
- index = len(n.items) - 1
- }
- for i := index; i >= 0; i-- {
- if start.valid && !n.cow.less(n.items[i], start.item) {
- if !includeStart || hit || n.cow.less(start.item, n.items[i]) {
- continue
- }
- }
- if len(n.children) > 0 {
- if hit, ok = n.children[i+1].iterate(dir, start, stop, includeStart, hit, iter); !ok {
- return hit, false
- }
- }
- if stop.valid && !n.cow.less(stop.item, n.items[i]) {
- return hit, false // continue
- }
- hit = true
- if !iter(n.items[i]) {
- return hit, false
- }
- }
- if len(n.children) > 0 {
- if hit, ok = n.children[0].iterate(dir, start, stop, includeStart, hit, iter); !ok {
- return hit, false
- }
- }
- }
- return hit, true
- }
- // print is used for testing/debugging purposes.
- func (n *node[T]) print(w io.Writer, level int) {
- fmt.Fprintf(w, "%sNODE:%v\n", strings.Repeat(" ", level), n.items)
- for _, c := range n.children {
- c.print(w, level+1)
- }
- }
- // BTreeG is a generic implementation of a B-Tree.
- //
- // BTreeG stores items of type T in an ordered structure, allowing easy insertion,
- // removal, and iteration.
- //
- // Write operations are not safe for concurrent mutation by multiple
- // goroutines, but Read operations are.
- type BTreeG[T any] struct {
- degree int
- length int
- root *node[T]
- cow *copyOnWriteContext[T]
- }
- // LessFunc[T] determines how to order a type 'T'. It should implement a strict
- // ordering, and should return true if within that ordering, 'a' < 'b'.
- type LessFunc[T any] func(a, b T) bool
- // copyOnWriteContext pointers determine node ownership... a tree with a write
- // context equivalent to a node's write context is allowed to modify that node.
- // A tree whose write context does not match a node's is not allowed to modify
- // it, and must create a new, writable copy (IE: it's a Clone).
- //
- // When doing any write operation, we maintain the invariant that the current
- // node's context is equal to the context of the tree that requested the write.
- // We do this by, before we descend into any node, creating a copy with the
- // correct context if the contexts don't match.
- //
- // Since the node we're currently visiting on any write has the requesting
- // tree's context, that node is modifiable in place. Children of that node may
- // not share context, but before we descend into them, we'll make a mutable
- // copy.
- type copyOnWriteContext[T any] struct {
- freelist *FreeListG[T]
- less LessFunc[T]
- }
- // Clone clones the btree, lazily. Clone should not be called concurrently,
- // but the original tree (t) and the new tree (t2) can be used concurrently
- // once the Clone call completes.
- //
- // The internal tree structure of b is marked read-only and shared between t and
- // t2. Writes to both t and t2 use copy-on-write logic, creating new nodes
- // whenever one of b's original nodes would have been modified. Read operations
- // should have no performance degredation. Write operations for both t and t2
- // will initially experience minor slow-downs caused by additional allocs and
- // copies due to the aforementioned copy-on-write logic, but should converge to
- // the original performance characteristics of the original tree.
- func (t *BTreeG[T]) Clone() (t2 *BTreeG[T]) {
- // Create two entirely new copy-on-write contexts.
- // This operation effectively creates three trees:
- // the original, shared nodes (old b.cow)
- // the new b.cow nodes
- // the new out.cow nodes
- cow1, cow2 := *t.cow, *t.cow
- out := *t
- t.cow = &cow1
- out.cow = &cow2
- return &out
- }
- // maxItems returns the max number of items to allow per node.
- func (t *BTreeG[T]) maxItems() int {
- return t.degree*2 - 1
- }
- // minItems returns the min number of items to allow per node (ignored for the
- // root node).
- func (t *BTreeG[T]) minItems() int {
- return t.degree - 1
- }
- func (c *copyOnWriteContext[T]) newNode() (n *node[T]) {
- n = c.freelist.newNode()
- n.cow = c
- return
- }
- type freeType int
- const (
- ftFreelistFull freeType = iota // node was freed (available for GC, not stored in freelist)
- ftStored // node was stored in the freelist for later use
- ftNotOwned // node was ignored by COW, since it's owned by another one
- )
- // freeNode frees a node within a given COW context, if it's owned by that
- // context. It returns what happened to the node (see freeType const
- // documentation).
- func (c *copyOnWriteContext[T]) freeNode(n *node[T]) freeType {
- if n.cow == c {
- // clear to allow GC
- n.items.truncate(0)
- n.children.truncate(0)
- n.cow = nil
- if c.freelist.freeNode(n) {
- return ftStored
- } else {
- return ftFreelistFull
- }
- } else {
- return ftNotOwned
- }
- }
- // ReplaceOrInsert adds the given item to the tree. If an item in the tree
- // already equals the given one, it is removed from the tree and returned,
- // and the second return value is true. Otherwise, (zeroValue, false)
- //
- // nil cannot be added to the tree (will panic).
- func (t *BTreeG[T]) ReplaceOrInsert(item T) (_ T, _ bool) {
- if t.root == nil {
- t.root = t.cow.newNode()
- t.root.items = append(t.root.items, item)
- t.length++
- return
- } else {
- t.root = t.root.mutableFor(t.cow)
- if len(t.root.items) >= t.maxItems() {
- item2, second := t.root.split(t.maxItems() / 2)
- oldroot := t.root
- t.root = t.cow.newNode()
- t.root.items = append(t.root.items, item2)
- t.root.children = append(t.root.children, oldroot, second)
- }
- }
- out, outb := t.root.insert(item, t.maxItems())
- if !outb {
- t.length++
- }
- return out, outb
- }
- // Delete removes an item equal to the passed in item from the tree, returning
- // it. If no such item exists, returns (zeroValue, false).
- func (t *BTreeG[T]) Delete(item T) (T, bool) {
- return t.deleteItem(item, removeItem)
- }
- // DeleteMin removes the smallest item in the tree and returns it.
- // If no such item exists, returns (zeroValue, false).
- func (t *BTreeG[T]) DeleteMin() (T, bool) {
- var zero T
- return t.deleteItem(zero, removeMin)
- }
- // DeleteMax removes the largest item in the tree and returns it.
- // If no such item exists, returns (zeroValue, false).
- func (t *BTreeG[T]) DeleteMax() (T, bool) {
- var zero T
- return t.deleteItem(zero, removeMax)
- }
- func (t *BTreeG[T]) deleteItem(item T, typ toRemove) (_ T, _ bool) {
- if t.root == nil || len(t.root.items) == 0 {
- return
- }
- t.root = t.root.mutableFor(t.cow)
- out, outb := t.root.remove(item, t.minItems(), typ)
- if len(t.root.items) == 0 && len(t.root.children) > 0 {
- oldroot := t.root
- t.root = t.root.children[0]
- t.cow.freeNode(oldroot)
- }
- if outb {
- t.length--
- }
- return out, outb
- }
- // AscendRange calls the iterator for every value in the tree within the range
- // [greaterOrEqual, lessThan), until iterator returns false.
- func (t *BTreeG[T]) AscendRange(greaterOrEqual, lessThan T, iterator ItemIteratorG[T]) {
- if t.root == nil {
- return
- }
- t.root.iterate(ascend, optional[T](greaterOrEqual), optional[T](lessThan), true, false, iterator)
- }
- // AscendLessThan calls the iterator for every value in the tree within the range
- // [first, pivot), until iterator returns false.
- func (t *BTreeG[T]) AscendLessThan(pivot T, iterator ItemIteratorG[T]) {
- if t.root == nil {
- return
- }
- t.root.iterate(ascend, empty[T](), optional(pivot), false, false, iterator)
- }
- // AscendGreaterOrEqual calls the iterator for every value in the tree within
- // the range [pivot, last], until iterator returns false.
- func (t *BTreeG[T]) AscendGreaterOrEqual(pivot T, iterator ItemIteratorG[T]) {
- if t.root == nil {
- return
- }
- t.root.iterate(ascend, optional[T](pivot), empty[T](), true, false, iterator)
- }
- // Ascend calls the iterator for every value in the tree within the range
- // [first, last], until iterator returns false.
- func (t *BTreeG[T]) Ascend(iterator ItemIteratorG[T]) {
- if t.root == nil {
- return
- }
- t.root.iterate(ascend, empty[T](), empty[T](), false, false, iterator)
- }
- // DescendRange calls the iterator for every value in the tree within the range
- // [lessOrEqual, greaterThan), until iterator returns false.
- func (t *BTreeG[T]) DescendRange(lessOrEqual, greaterThan T, iterator ItemIteratorG[T]) {
- if t.root == nil {
- return
- }
- t.root.iterate(descend, optional[T](lessOrEqual), optional[T](greaterThan), true, false, iterator)
- }
- // DescendLessOrEqual calls the iterator for every value in the tree within the range
- // [pivot, first], until iterator returns false.
- func (t *BTreeG[T]) DescendLessOrEqual(pivot T, iterator ItemIteratorG[T]) {
- if t.root == nil {
- return
- }
- t.root.iterate(descend, optional[T](pivot), empty[T](), true, false, iterator)
- }
- // DescendGreaterThan calls the iterator for every value in the tree within
- // the range [last, pivot), until iterator returns false.
- func (t *BTreeG[T]) DescendGreaterThan(pivot T, iterator ItemIteratorG[T]) {
- if t.root == nil {
- return
- }
- t.root.iterate(descend, empty[T](), optional[T](pivot), false, false, iterator)
- }
- // Descend calls the iterator for every value in the tree within the range
- // [last, first], until iterator returns false.
- func (t *BTreeG[T]) Descend(iterator ItemIteratorG[T]) {
- if t.root == nil {
- return
- }
- t.root.iterate(descend, empty[T](), empty[T](), false, false, iterator)
- }
- // Get looks for the key item in the tree, returning it. It returns
- // (zeroValue, false) if unable to find that item.
- func (t *BTreeG[T]) Get(key T) (_ T, _ bool) {
- if t.root == nil {
- return
- }
- return t.root.get(key)
- }
- // Min returns the smallest item in the tree, or (zeroValue, false) if the tree is empty.
- func (t *BTreeG[T]) Min() (_ T, _ bool) {
- return min(t.root)
- }
- // Max returns the largest item in the tree, or (zeroValue, false) if the tree is empty.
- func (t *BTreeG[T]) Max() (_ T, _ bool) {
- return max(t.root)
- }
- // Has returns true if the given key is in the tree.
- func (t *BTreeG[T]) Has(key T) bool {
- _, ok := t.Get(key)
- return ok
- }
- // Len returns the number of items currently in the tree.
- func (t *BTreeG[T]) Len() int {
- return t.length
- }
- // Clear removes all items from the btree. If addNodesToFreelist is true,
- // t's nodes are added to its freelist as part of this call, until the freelist
- // is full. Otherwise, the root node is simply dereferenced and the subtree
- // left to Go's normal GC processes.
- //
- // This can be much faster
- // than calling Delete on all elements, because that requires finding/removing
- // each element in the tree and updating the tree accordingly. It also is
- // somewhat faster than creating a new tree to replace the old one, because
- // nodes from the old tree are reclaimed into the freelist for use by the new
- // one, instead of being lost to the garbage collector.
- //
- // This call takes:
- // O(1): when addNodesToFreelist is false, this is a single operation.
- // O(1): when the freelist is already full, it breaks out immediately
- // O(freelist size): when the freelist is empty and the nodes are all owned
- // by this tree, nodes are added to the freelist until full.
- // O(tree size): when all nodes are owned by another tree, all nodes are
- // iterated over looking for nodes to add to the freelist, and due to
- // ownership, none are.
- func (t *BTreeG[T]) Clear(addNodesToFreelist bool) {
- if t.root != nil && addNodesToFreelist {
- t.root.reset(t.cow)
- }
- t.root, t.length = nil, 0
- }
- // reset returns a subtree to the freelist. It breaks out immediately if the
- // freelist is full, since the only benefit of iterating is to fill that
- // freelist up. Returns true if parent reset call should continue.
- func (n *node[T]) reset(c *copyOnWriteContext[T]) bool {
- for _, child := range n.children {
- if !child.reset(c) {
- return false
- }
- }
- return c.freeNode(n) != ftFreelistFull
- }
- // Int implements the Item interface for integers.
- type Int int
- // Less returns true if int(a) < int(b).
- func (a Int) Less(b Item) bool {
- return a < b.(Int)
- }
- // BTree is an implementation of a B-Tree.
- //
- // BTree stores Item instances in an ordered structure, allowing easy insertion,
- // removal, and iteration.
- //
- // Write operations are not safe for concurrent mutation by multiple
- // goroutines, but Read operations are.
- type BTree BTreeG[Item]
- var itemLess LessFunc[Item] = func(a, b Item) bool {
- return a.Less(b)
- }
- // New creates a new B-Tree with the given degree.
- //
- // New(2), for example, will create a 2-3-4 tree (each node contains 1-3 items
- // and 2-4 children).
- func New(degree int) *BTree {
- return (*BTree)(NewG[Item](degree, itemLess))
- }
- // FreeList represents a free list of btree nodes. By default each
- // BTree has its own FreeList, but multiple BTrees can share the same
- // FreeList.
- // Two Btrees using the same freelist are safe for concurrent write access.
- type FreeList FreeListG[Item]
- // NewFreeList creates a new free list.
- // size is the maximum size of the returned free list.
- func NewFreeList(size int) *FreeList {
- return (*FreeList)(NewFreeListG[Item](size))
- }
- // NewWithFreeList creates a new B-Tree that uses the given node free list.
- func NewWithFreeList(degree int, f *FreeList) *BTree {
- return (*BTree)(NewWithFreeListG[Item](degree, itemLess, (*FreeListG[Item])(f)))
- }
- // ItemIterator allows callers of Ascend* to iterate in-order over portions of
- // the tree. When this function returns false, iteration will stop and the
- // associated Ascend* function will immediately return.
- type ItemIterator ItemIteratorG[Item]
- // Clone clones the btree, lazily. Clone should not be called concurrently,
- // but the original tree (t) and the new tree (t2) can be used concurrently
- // once the Clone call completes.
- //
- // The internal tree structure of b is marked read-only and shared between t and
- // t2. Writes to both t and t2 use copy-on-write logic, creating new nodes
- // whenever one of b's original nodes would have been modified. Read operations
- // should have no performance degredation. Write operations for both t and t2
- // will initially experience minor slow-downs caused by additional allocs and
- // copies due to the aforementioned copy-on-write logic, but should converge to
- // the original performance characteristics of the original tree.
- func (t *BTree) Clone() (t2 *BTree) {
- return (*BTree)((*BTreeG[Item])(t).Clone())
- }
- // Delete removes an item equal to the passed in item from the tree, returning
- // it. If no such item exists, returns nil.
- func (t *BTree) Delete(item Item) Item {
- i, _ := (*BTreeG[Item])(t).Delete(item)
- return i
- }
- // DeleteMax removes the largest item in the tree and returns it.
- // If no such item exists, returns nil.
- func (t *BTree) DeleteMax() Item {
- i, _ := (*BTreeG[Item])(t).DeleteMax()
- return i
- }
- // DeleteMin removes the smallest item in the tree and returns it.
- // If no such item exists, returns nil.
- func (t *BTree) DeleteMin() Item {
- i, _ := (*BTreeG[Item])(t).DeleteMin()
- return i
- }
- // Get looks for the key item in the tree, returning it. It returns nil if
- // unable to find that item.
- func (t *BTree) Get(key Item) Item {
- i, _ := (*BTreeG[Item])(t).Get(key)
- return i
- }
- // Max returns the largest item in the tree, or nil if the tree is empty.
- func (t *BTree) Max() Item {
- i, _ := (*BTreeG[Item])(t).Max()
- return i
- }
- // Min returns the smallest item in the tree, or nil if the tree is empty.
- func (t *BTree) Min() Item {
- i, _ := (*BTreeG[Item])(t).Min()
- return i
- }
- // Has returns true if the given key is in the tree.
- func (t *BTree) Has(key Item) bool {
- return (*BTreeG[Item])(t).Has(key)
- }
- // ReplaceOrInsert adds the given item to the tree. If an item in the tree
- // already equals the given one, it is removed from the tree and returned.
- // Otherwise, nil is returned.
- //
- // nil cannot be added to the tree (will panic).
- func (t *BTree) ReplaceOrInsert(item Item) Item {
- i, _ := (*BTreeG[Item])(t).ReplaceOrInsert(item)
- return i
- }
- // AscendRange calls the iterator for every value in the tree within the range
- // [greaterOrEqual, lessThan), until iterator returns false.
- func (t *BTree) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) {
- (*BTreeG[Item])(t).AscendRange(greaterOrEqual, lessThan, (ItemIteratorG[Item])(iterator))
- }
- // AscendLessThan calls the iterator for every value in the tree within the range
- // [first, pivot), until iterator returns false.
- func (t *BTree) AscendLessThan(pivot Item, iterator ItemIterator) {
- (*BTreeG[Item])(t).AscendLessThan(pivot, (ItemIteratorG[Item])(iterator))
- }
- // AscendGreaterOrEqual calls the iterator for every value in the tree within
- // the range [pivot, last], until iterator returns false.
- func (t *BTree) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {
- (*BTreeG[Item])(t).AscendGreaterOrEqual(pivot, (ItemIteratorG[Item])(iterator))
- }
- // Ascend calls the iterator for every value in the tree within the range
- // [first, last], until iterator returns false.
- func (t *BTree) Ascend(iterator ItemIterator) {
- (*BTreeG[Item])(t).Ascend((ItemIteratorG[Item])(iterator))
- }
- // DescendRange calls the iterator for every value in the tree within the range
- // [lessOrEqual, greaterThan), until iterator returns false.
- func (t *BTree) DescendRange(lessOrEqual, greaterThan Item, iterator ItemIterator) {
- (*BTreeG[Item])(t).DescendRange(lessOrEqual, greaterThan, (ItemIteratorG[Item])(iterator))
- }
- // DescendLessOrEqual calls the iterator for every value in the tree within the range
- // [pivot, first], until iterator returns false.
- func (t *BTree) DescendLessOrEqual(pivot Item, iterator ItemIterator) {
- (*BTreeG[Item])(t).DescendLessOrEqual(pivot, (ItemIteratorG[Item])(iterator))
- }
- // DescendGreaterThan calls the iterator for every value in the tree within
- // the range [last, pivot), until iterator returns false.
- func (t *BTree) DescendGreaterThan(pivot Item, iterator ItemIterator) {
- (*BTreeG[Item])(t).DescendGreaterThan(pivot, (ItemIteratorG[Item])(iterator))
- }
- // Descend calls the iterator for every value in the tree within the range
- // [last, first], until iterator returns false.
- func (t *BTree) Descend(iterator ItemIterator) {
- (*BTreeG[Item])(t).Descend((ItemIteratorG[Item])(iterator))
- }
- // Len returns the number of items currently in the tree.
- func (t *BTree) Len() int {
- return (*BTreeG[Item])(t).Len()
- }
- // Clear removes all items from the btree. If addNodesToFreelist is true,
- // t's nodes are added to its freelist as part of this call, until the freelist
- // is full. Otherwise, the root node is simply dereferenced and the subtree
- // left to Go's normal GC processes.
- //
- // This can be much faster
- // than calling Delete on all elements, because that requires finding/removing
- // each element in the tree and updating the tree accordingly. It also is
- // somewhat faster than creating a new tree to replace the old one, because
- // nodes from the old tree are reclaimed into the freelist for use by the new
- // one, instead of being lost to the garbage collector.
- //
- // This call takes:
- // O(1): when addNodesToFreelist is false, this is a single operation.
- // O(1): when the freelist is already full, it breaks out immediately
- // O(freelist size): when the freelist is empty and the nodes are all owned
- // by this tree, nodes are added to the freelist until full.
- // O(tree size): when all nodes are owned by another tree, all nodes are
- // iterated over looking for nodes to add to the freelist, and due to
- // ownership, none are.
- func (t *BTree) Clear(addNodesToFreelist bool) {
- (*BTreeG[Item])(t).Clear(addNodesToFreelist)
- }
|