// Copyright 2020 Joshua J Baker. All rights reserved. // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. package btree import "sync" type BTreeG[T any] struct { isoid uint64 mu *sync.RWMutex root *node[T] count int locks bool copyItems bool isoCopyItems bool less func(a, b T) bool empty T max int min int } type node[T any] struct { isoid uint64 count int items []T children *[]*node[T] } // PathHint is a utility type used with the *Hint() functions. Hints provide // faster operations for clustered keys. type PathHint struct { used [8]bool path [8]uint8 } // Options for passing to New when creating a new BTree. type Options struct { // Degree is used to define how many items and children each internal node // can contain before it must branch. For example, a degree of 2 will // create a 2-3-4 tree, where each node may contains 1-3 items and // 2-4 children. See https://en.wikipedia.org/wiki/2–3–4_tree. // Default is 32 Degree int // NoLocks will disable locking. Otherwide a sync.RWMutex is used to // ensure all operations are safe across multiple goroutines. NoLocks bool } // New returns a new BTree func NewBTreeG[T any](less func(a, b T) bool) *BTreeG[T] { return NewBTreeGOptions(less, Options{}) } func NewBTreeGOptions[T any](less func(a, b T) bool, opts Options) *BTreeG[T] { tr := new(BTreeG[T]) tr.isoid = newIsoID() tr.mu = new(sync.RWMutex) tr.locks = !opts.NoLocks tr.less = less tr.init(opts.Degree) return tr } func (tr *BTreeG[T]) init(degree int) { if tr.min != 0 { return } tr.min, tr.max = degreeToMinMax(degree) _, tr.copyItems = ((interface{})(tr.empty)).(copier[T]) if !tr.copyItems { _, tr.isoCopyItems = ((interface{})(tr.empty)).(isoCopier[T]) } } // Less is a convenience function that performs a comparison of two items // using the same "less" function provided to New. func (tr *BTreeG[T]) Less(a, b T) bool { return tr.less(a, b) } func (tr *BTreeG[T]) newNode(leaf bool) *node[T] { n := &node[T]{isoid: tr.isoid} if !leaf { n.children = new([]*node[T]) } return n } // leaf returns true if the node is a leaf. func (n *node[T]) leaf() bool { return n.children == nil } func (tr *BTreeG[T]) bsearch(n *node[T], key T) (index int, found bool) { low, high := 0, len(n.items) for low < high { h := (low + high) / 2 if !tr.less(key, n.items[h]) { low = h + 1 } else { high = h } } if low > 0 && !tr.less(n.items[low-1], key) { return low - 1, true } return low, false } func (tr *BTreeG[T]) find(n *node[T], key T, hint *PathHint, depth int, ) (index int, found bool) { if hint == nil { return tr.bsearch(n, key) } return tr.hintsearch(n, key, hint, depth) } func (tr *BTreeG[T]) hintsearch(n *node[T], key T, hint *PathHint, depth int, ) (index int, found bool) { // Best case finds the exact match, updates the hint and returns. // Worst case, updates the low and high bounds to binary search between. low := 0 high := len(n.items) - 1 if depth < 8 && hint.used[depth] { index = int(hint.path[depth]) if index >= len(n.items) { // tail item if tr.Less(n.items[len(n.items)-1], key) { index = len(n.items) goto path_match } index = len(n.items) - 1 } if tr.Less(key, n.items[index]) { if index == 0 || tr.Less(n.items[index-1], key) { goto path_match } high = index - 1 } else if tr.Less(n.items[index], key) { low = index + 1 } else { found = true goto path_match } } // Do a binary search between low and high // keep on going until low > high, where the guarantee on low is that // key >= items[low - 1] for low <= high { mid := low + ((high+1)-low)/2 // if key >= n.items[mid], low = mid + 1 // which implies that key >= everything below low if !tr.Less(key, n.items[mid]) { low = mid + 1 } else { high = mid - 1 } } // if low > 0, n.items[low - 1] >= key, // we have from before that key >= n.items[low - 1] // therefore key = n.items[low - 1], // and we have found the entry for key. // Otherwise we must keep searching for the key in index `low`. if low > 0 && !tr.Less(n.items[low-1], key) { index = low - 1 found = true } else { index = low found = false } path_match: if depth < 8 { hint.used[depth] = true var pathIndex uint8 if n.leaf() && found { pathIndex = uint8(index + 1) } else { pathIndex = uint8(index) } if pathIndex != hint.path[depth] { hint.path[depth] = pathIndex for i := depth + 1; i < 8; i++ { hint.used[i] = false } } } return index, found } // SetHint sets or replace a value for a key using a path hint func (tr *BTreeG[T]) SetHint(item T, hint *PathHint) (prev T, replaced bool) { if tr.locks { tr.mu.Lock() prev, replaced = tr.setHint(item, hint) tr.mu.Unlock() } else { prev, replaced = tr.setHint(item, hint) } return prev, replaced } func (tr *BTreeG[T]) setHint(item T, hint *PathHint) (prev T, replaced bool) { if tr.root == nil { tr.init(0) tr.root = tr.newNode(true) tr.root.items = append([]T{}, item) tr.root.count = 1 tr.count = 1 return tr.empty, false } prev, replaced, split := tr.nodeSet(&tr.root, item, hint, 0) if split { left := tr.isoLoad(&tr.root, true) right, median := tr.nodeSplit(left) tr.root = tr.newNode(false) *tr.root.children = make([]*node[T], 0, tr.max+1) *tr.root.children = append([]*node[T]{}, left, right) tr.root.items = append([]T{}, median) tr.root.updateCount() return tr.setHint(item, hint) } if replaced { return prev, true } tr.count++ return tr.empty, false } // Set or replace a value for a key func (tr *BTreeG[T]) Set(item T) (T, bool) { return tr.SetHint(item, nil) } func (tr *BTreeG[T]) nodeSplit(n *node[T]) (right *node[T], median T) { i := tr.max / 2 median = n.items[i] // right node right = tr.newNode(n.leaf()) right.items = n.items[i+1:] if !n.leaf() { *right.children = (*n.children)[i+1:] } right.updateCount() // left node n.items[i] = tr.empty n.items = n.items[:i:i] if !n.leaf() { *n.children = (*n.children)[: i+1 : i+1] } n.updateCount() return right, median } func (n *node[T]) updateCount() { n.count = len(n.items) if !n.leaf() { for i := 0; i < len(*n.children); i++ { n.count += (*n.children)[i].count } } } // Copy the node for safe isolation. func (tr *BTreeG[T]) copy(n *node[T]) *node[T] { n2 := new(node[T]) n2.isoid = tr.isoid n2.count = n.count n2.items = make([]T, len(n.items), cap(n.items)) copy(n2.items, n.items) if tr.copyItems { for i := 0; i < len(n2.items); i++ { n2.items[i] = ((interface{})(n2.items[i])).(copier[T]).Copy() } } else if tr.isoCopyItems { for i := 0; i < len(n2.items); i++ { n2.items[i] = ((interface{})(n2.items[i])).(isoCopier[T]).IsoCopy() } } if !n.leaf() { n2.children = new([]*node[T]) *n2.children = make([]*node[T], len(*n.children), tr.max+1) copy(*n2.children, *n.children) } return n2 } // isoLoad loads the provided node and, if needed, performs a copy-on-write. func (tr *BTreeG[T]) isoLoad(cn **node[T], mut bool) *node[T] { if mut && (*cn).isoid != tr.isoid { *cn = tr.copy(*cn) } return *cn } func (tr *BTreeG[T]) nodeSet(cn **node[T], item T, hint *PathHint, depth int, ) (prev T, replaced bool, split bool) { if (*cn).isoid != tr.isoid { *cn = tr.copy(*cn) } n := *cn var i int var found bool if hint == nil { i, found = tr.bsearch(n, item) } else { i, found = tr.hintsearch(n, item, hint, depth) } if found { prev = n.items[i] n.items[i] = item return prev, true, false } if n.leaf() { if len(n.items) == tr.max { return tr.empty, false, true } n.items = append(n.items, tr.empty) copy(n.items[i+1:], n.items[i:]) n.items[i] = item n.count++ return tr.empty, false, false } prev, replaced, split = tr.nodeSet(&(*n.children)[i], item, hint, depth+1) if split { if len(n.items) == tr.max { return tr.empty, false, true } right, median := tr.nodeSplit((*n.children)[i]) *n.children = append(*n.children, nil) copy((*n.children)[i+1:], (*n.children)[i:]) (*n.children)[i+1] = right n.items = append(n.items, tr.empty) copy(n.items[i+1:], n.items[i:]) n.items[i] = median return tr.nodeSet(&n, item, hint, depth) } if !replaced { n.count++ } return prev, replaced, false } func (tr *BTreeG[T]) Scan(iter func(item T) bool) { tr.scan(iter, false) } func (tr *BTreeG[T]) ScanMut(iter func(item T) bool) { tr.scan(iter, true) } func (tr *BTreeG[T]) scan(iter func(item T) bool, mut bool) { if tr.lock(mut) { defer tr.unlock(mut) } if tr.root == nil { return } tr.nodeScan(&tr.root, iter, mut) } func (tr *BTreeG[T]) nodeScan(cn **node[T], iter func(item T) bool, mut bool, ) bool { n := tr.isoLoad(cn, mut) if n.leaf() { for i := 0; i < len(n.items); i++ { if !iter(n.items[i]) { return false } } return true } for i := 0; i < len(n.items); i++ { if !tr.nodeScan(&(*n.children)[i], iter, mut) { return false } if !iter(n.items[i]) { return false } } return tr.nodeScan(&(*n.children)[len(*n.children)-1], iter, mut) } // Get a value for key func (tr *BTreeG[T]) Get(key T) (T, bool) { return tr.getHint(key, nil, false) } func (tr *BTreeG[T]) GetMut(key T) (T, bool) { return tr.getHint(key, nil, true) } // GetHint gets a value for key using a path hint func (tr *BTreeG[T]) GetHint(key T, hint *PathHint) (value T, ok bool) { return tr.getHint(key, hint, false) } func (tr *BTreeG[T]) GetHintMut(key T, hint *PathHint) (value T, ok bool) { return tr.getHint(key, hint, true) } // GetHint gets a value for key using a path hint func (tr *BTreeG[T]) getHint(key T, hint *PathHint, mut bool) (T, bool) { if tr.lock(mut) { defer tr.unlock(mut) } if tr.root == nil { return tr.empty, false } n := tr.isoLoad(&tr.root, mut) depth := 0 for { i, found := tr.find(n, key, hint, depth) if found { return n.items[i], true } if n.children == nil { return tr.empty, false } n = tr.isoLoad(&(*n.children)[i], mut) depth++ } } // Len returns the number of items in the tree func (tr *BTreeG[T]) Len() int { return tr.count } // Delete a value for a key and returns the deleted value. // Returns false if there was no value by that key found. func (tr *BTreeG[T]) Delete(key T) (T, bool) { return tr.DeleteHint(key, nil) } // DeleteHint deletes a value for a key using a path hint and returns the // deleted value. // Returns false if there was no value by that key found. func (tr *BTreeG[T]) DeleteHint(key T, hint *PathHint) (T, bool) { if tr.lock(true) { defer tr.unlock(true) } return tr.deleteHint(key, hint) } func (tr *BTreeG[T]) deleteHint(key T, hint *PathHint) (T, bool) { if tr.root == nil { return tr.empty, false } prev, deleted := tr.delete(&tr.root, false, key, hint, 0) if !deleted { return tr.empty, false } if len(tr.root.items) == 0 && !tr.root.leaf() { tr.root = (*tr.root.children)[0] } tr.count-- if tr.count == 0 { tr.root = nil } return prev, true } func (tr *BTreeG[T]) delete(cn **node[T], max bool, key T, hint *PathHint, depth int, ) (T, bool) { n := tr.isoLoad(cn, true) var i int var found bool if max { i, found = len(n.items)-1, true } else { i, found = tr.find(n, key, hint, depth) } if n.leaf() { if found { // found the items at the leaf, remove it and return. prev := n.items[i] copy(n.items[i:], n.items[i+1:]) n.items[len(n.items)-1] = tr.empty n.items = n.items[:len(n.items)-1] n.count-- return prev, true } return tr.empty, false } var prev T var deleted bool if found { if max { i++ prev, deleted = tr.delete(&(*n.children)[i], true, tr.empty, nil, 0) } else { prev = n.items[i] maxItem, _ := tr.delete(&(*n.children)[i], true, tr.empty, nil, 0) deleted = true n.items[i] = maxItem } } else { prev, deleted = tr.delete(&(*n.children)[i], max, key, hint, depth+1) } if !deleted { return tr.empty, false } n.count-- if len((*n.children)[i].items) < tr.min { tr.nodeRebalance(n, i) } return prev, true } // nodeRebalance rebalances the child nodes following a delete operation. // Provide the index of the child node with the number of items that fell // below minItems. func (tr *BTreeG[T]) nodeRebalance(n *node[T], i int) { if i == len(n.items) { i-- } // ensure copy-on-write left := tr.isoLoad(&(*n.children)[i], true) right := tr.isoLoad(&(*n.children)[i+1], true) if len(left.items)+len(right.items) < tr.max { // Merges the left and right children nodes together as a single node // that includes (left,item,right), and places the contents into the // existing left node. Delete the right node altogether and move the // following items and child nodes to the left by one slot. // merge (left,item,right) left.items = append(left.items, n.items[i]) left.items = append(left.items, right.items...) if !left.leaf() { *left.children = append(*left.children, *right.children...) } left.count += right.count + 1 // move the items over one slot copy(n.items[i:], n.items[i+1:]) n.items[len(n.items)-1] = tr.empty n.items = n.items[:len(n.items)-1] // move the children over one slot copy((*n.children)[i+1:], (*n.children)[i+2:]) (*n.children)[len(*n.children)-1] = nil (*n.children) = (*n.children)[:len(*n.children)-1] } else if len(left.items) > len(right.items) { // move left -> right over one slot // Move the item of the parent node at index into the right-node first // slot, and move the left-node last item into the previously moved // parent item slot. right.items = append(right.items, tr.empty) copy(right.items[1:], right.items) right.items[0] = n.items[i] right.count++ n.items[i] = left.items[len(left.items)-1] left.items[len(left.items)-1] = tr.empty left.items = left.items[:len(left.items)-1] left.count-- if !left.leaf() { // move the left-node last child into the right-node first slot *right.children = append(*right.children, nil) copy((*right.children)[1:], *right.children) (*right.children)[0] = (*left.children)[len(*left.children)-1] (*left.children)[len(*left.children)-1] = nil (*left.children) = (*left.children)[:len(*left.children)-1] left.count -= (*right.children)[0].count right.count += (*right.children)[0].count } } else { // move left <- right over one slot // Same as above but the other direction left.items = append(left.items, n.items[i]) left.count++ n.items[i] = right.items[0] copy(right.items, right.items[1:]) right.items[len(right.items)-1] = tr.empty right.items = right.items[:len(right.items)-1] right.count-- if !left.leaf() { *left.children = append(*left.children, (*right.children)[0]) copy(*right.children, (*right.children)[1:]) (*right.children)[len(*right.children)-1] = nil *right.children = (*right.children)[:len(*right.children)-1] left.count += (*left.children)[len(*left.children)-1].count right.count -= (*left.children)[len(*left.children)-1].count } } } // Ascend the tree within the range [pivot, last] // Pass nil for pivot to scan all item in ascending order // Return false to stop iterating func (tr *BTreeG[T]) Ascend(pivot T, iter func(item T) bool) { tr.ascend(pivot, iter, false) } func (tr *BTreeG[T]) AscendMut(pivot T, iter func(item T) bool) { tr.ascend(pivot, iter, true) } func (tr *BTreeG[T]) ascend(pivot T, iter func(item T) bool, mut bool) { if tr.lock(mut) { defer tr.unlock(mut) } if tr.root == nil { return } tr.nodeAscend(&tr.root, pivot, nil, 0, iter, mut) } // The return value of this function determines whether we should keep iterating // upon this functions return. func (tr *BTreeG[T]) nodeAscend(cn **node[T], pivot T, hint *PathHint, depth int, iter func(item T) bool, mut bool, ) bool { n := tr.isoLoad(cn, mut) i, found := tr.find(n, pivot, hint, depth) if !found { if !n.leaf() { if !tr.nodeAscend(&(*n.children)[i], pivot, hint, depth+1, iter, mut) { return false } } } // We are either in the case that // - node is found, we should iterate through it starting at `i`, // the index it was located at. // - node is not found, and TODO: fill in. for ; i < len(n.items); i++ { if !iter(n.items[i]) { return false } if !n.leaf() { if !tr.nodeScan(&(*n.children)[i+1], iter, mut) { return false } } } return true } func (tr *BTreeG[T]) Reverse(iter func(item T) bool) { tr.reverse(iter, false) } func (tr *BTreeG[T]) ReverseMut(iter func(item T) bool) { tr.reverse(iter, true) } func (tr *BTreeG[T]) reverse(iter func(item T) bool, mut bool) { if tr.lock(mut) { defer tr.unlock(mut) } if tr.root == nil { return } tr.nodeReverse(&tr.root, iter, mut) } func (tr *BTreeG[T]) nodeReverse(cn **node[T], iter func(item T) bool, mut bool, ) bool { n := tr.isoLoad(cn, mut) if n.leaf() { for i := len(n.items) - 1; i >= 0; i-- { if !iter(n.items[i]) { return false } } return true } if !tr.nodeReverse(&(*n.children)[len(*n.children)-1], iter, mut) { return false } for i := len(n.items) - 1; i >= 0; i-- { if !iter(n.items[i]) { return false } if !tr.nodeReverse(&(*n.children)[i], iter, mut) { return false } } return true } // Descend the tree within the range [pivot, first] // Pass nil for pivot to scan all item in descending order // Return false to stop iterating func (tr *BTreeG[T]) Descend(pivot T, iter func(item T) bool) { tr.descend(pivot, iter, false) } func (tr *BTreeG[T]) DescendMut(pivot T, iter func(item T) bool) { tr.descend(pivot, iter, true) } func (tr *BTreeG[T]) descend(pivot T, iter func(item T) bool, mut bool) { if tr.lock(mut) { defer tr.unlock(mut) } if tr.root == nil { return } tr.nodeDescend(&tr.root, pivot, nil, 0, iter, mut) } func (tr *BTreeG[T]) nodeDescend(cn **node[T], pivot T, hint *PathHint, depth int, iter func(item T) bool, mut bool, ) bool { n := tr.isoLoad(cn, mut) i, found := tr.find(n, pivot, hint, depth) if !found { if !n.leaf() { if !tr.nodeDescend(&(*n.children)[i], pivot, hint, depth+1, iter, mut) { return false } } i-- } for ; i >= 0; i-- { if !iter(n.items[i]) { return false } if !n.leaf() { if !tr.nodeReverse(&(*n.children)[i], iter, mut) { return false } } } return true } // Load is for bulk loading pre-sorted items func (tr *BTreeG[T]) Load(item T) (T, bool) { if tr.lock(true) { defer tr.unlock(true) } if tr.root == nil { return tr.setHint(item, nil) } n := tr.isoLoad(&tr.root, true) for { n.count++ // optimistically update counts if n.leaf() { if len(n.items) < tr.max { if tr.Less(n.items[len(n.items)-1], item) { n.items = append(n.items, item) tr.count++ return tr.empty, false } } break } n = tr.isoLoad(&(*n.children)[len(*n.children)-1], true) } // revert the counts n = tr.root for { n.count-- if n.leaf() { break } n = (*n.children)[len(*n.children)-1] } return tr.setHint(item, nil) } // Min returns the minimum item in tree. // Returns nil if the treex has no items. func (tr *BTreeG[T]) Min() (T, bool) { return tr.minMut(false) } func (tr *BTreeG[T]) MinMut() (T, bool) { return tr.minMut(true) } func (tr *BTreeG[T]) minMut(mut bool) (T, bool) { if tr.lock(mut) { defer tr.unlock(mut) } if tr.root == nil { return tr.empty, false } n := tr.isoLoad(&tr.root, mut) for { if n.leaf() { return n.items[0], true } n = tr.isoLoad(&(*n.children)[0], mut) } } // Max returns the maximum item in tree. // Returns nil if the tree has no items. func (tr *BTreeG[T]) Max() (T, bool) { return tr.maxMut(false) } func (tr *BTreeG[T]) MaxMut() (T, bool) { return tr.maxMut(true) } func (tr *BTreeG[T]) maxMut(mut bool) (T, bool) { if tr.lock(mut) { defer tr.unlock(mut) } if tr.root == nil { return tr.empty, false } n := tr.isoLoad(&tr.root, mut) for { if n.leaf() { return n.items[len(n.items)-1], true } n = tr.isoLoad(&(*n.children)[len(*n.children)-1], mut) } } // PopMin removes the minimum item in tree and returns it. // Returns nil if the tree has no items. func (tr *BTreeG[T]) PopMin() (T, bool) { if tr.lock(true) { defer tr.unlock(true) } if tr.root == nil { return tr.empty, false } n := tr.isoLoad(&tr.root, true) var item T for { n.count-- // optimistically update counts if n.leaf() { item = n.items[0] if len(n.items) == tr.min { break } copy(n.items[:], n.items[1:]) n.items[len(n.items)-1] = tr.empty n.items = n.items[:len(n.items)-1] tr.count-- if tr.count == 0 { tr.root = nil } return item, true } n = tr.isoLoad(&(*n.children)[0], true) } // revert the counts n = tr.root for { n.count++ if n.leaf() { break } n = (*n.children)[0] } return tr.deleteHint(item, nil) } // PopMax removes the maximum item in tree and returns it. // Returns nil if the tree has no items. func (tr *BTreeG[T]) PopMax() (T, bool) { if tr.lock(true) { defer tr.unlock(true) } if tr.root == nil { return tr.empty, false } n := tr.isoLoad(&tr.root, true) var item T for { n.count-- // optimistically update counts if n.leaf() { item = n.items[len(n.items)-1] if len(n.items) == tr.min { break } n.items[len(n.items)-1] = tr.empty n.items = n.items[:len(n.items)-1] tr.count-- if tr.count == 0 { tr.root = nil } return item, true } n = tr.isoLoad(&(*n.children)[len(*n.children)-1], true) } // revert the counts n = tr.root for { n.count++ if n.leaf() { break } n = (*n.children)[len(*n.children)-1] } return tr.deleteHint(item, nil) } // GetAt returns the value at index. // Return nil if the tree is empty or the index is out of bounds. func (tr *BTreeG[T]) GetAt(index int) (T, bool) { return tr.getAt(index, false) } func (tr *BTreeG[T]) GetAtMut(index int) (T, bool) { return tr.getAt(index, true) } func (tr *BTreeG[T]) getAt(index int, mut bool) (T, bool) { if tr.lock(mut) { defer tr.unlock(mut) } if tr.root == nil || index < 0 || index >= tr.count { return tr.empty, false } n := tr.isoLoad(&tr.root, mut) for { if n.leaf() { return n.items[index], true } i := 0 for ; i < len(n.items); i++ { if index < (*n.children)[i].count { break } else if index == (*n.children)[i].count { return n.items[i], true } index -= (*n.children)[i].count + 1 } n = tr.isoLoad(&(*n.children)[i], mut) } } // DeleteAt deletes the item at index. // Return nil if the tree is empty or the index is out of bounds. func (tr *BTreeG[T]) DeleteAt(index int) (T, bool) { if tr.lock(true) { defer tr.unlock(true) } if tr.root == nil || index < 0 || index >= tr.count { return tr.empty, false } var pathbuf [8]uint8 // track the path path := pathbuf[:0] var item T n := tr.isoLoad(&tr.root, true) outer: for { n.count-- // optimistically update counts if n.leaf() { // the index is the item position item = n.items[index] if len(n.items) == tr.min { path = append(path, uint8(index)) break outer } copy(n.items[index:], n.items[index+1:]) n.items[len(n.items)-1] = tr.empty n.items = n.items[:len(n.items)-1] tr.count-- if tr.count == 0 { tr.root = nil } return item, true } i := 0 for ; i < len(n.items); i++ { if index < (*n.children)[i].count { break } else if index == (*n.children)[i].count { item = n.items[i] path = append(path, uint8(i)) break outer } index -= (*n.children)[i].count + 1 } path = append(path, uint8(i)) n = tr.isoLoad(&(*n.children)[i], true) } // revert the counts var hint PathHint n = tr.root for i := 0; i < len(path); i++ { if i < len(hint.path) { hint.path[i] = uint8(path[i]) hint.used[i] = true } n.count++ if !n.leaf() { n = (*n.children)[uint8(path[i])] } } return tr.deleteHint(item, &hint) } // Height returns the height of the tree. // Returns zero if tree has no items. func (tr *BTreeG[T]) Height() int { if tr.lock(false) { defer tr.unlock(false) } var height int if tr.root != nil { n := tr.root for { height++ if n.leaf() { break } n = (*n.children)[0] } } return height } // Walk iterates over all items in tree, in order. // The items param will contain one or more items. func (tr *BTreeG[T]) Walk(iter func(item []T) bool) { tr.walk(iter, false) } func (tr *BTreeG[T]) WalkMut(iter func(item []T) bool) { tr.walk(iter, true) } func (tr *BTreeG[T]) walk(iter func(item []T) bool, mut bool) { if tr.lock(mut) { defer tr.unlock(mut) } if tr.root == nil { return } tr.nodeWalk(&tr.root, iter, mut) } func (tr *BTreeG[T]) nodeWalk(cn **node[T], iter func(item []T) bool, mut bool, ) bool { n := tr.isoLoad(cn, mut) if n.leaf() { if !iter(n.items) { return false } } else { for i := 0; i < len(n.items); i++ { if !tr.nodeWalk(&(*n.children)[i], iter, mut) { return false } if !iter(n.items[i : i+1]) { return false } } if !tr.nodeWalk(&(*n.children)[len(n.items)], iter, mut) { return false } } return true } // Copy the tree. This is a copy-on-write operation and is very fast because // it only performs a shadowed copy. func (tr *BTreeG[T]) Copy() *BTreeG[T] { return tr.IsoCopy() } func (tr *BTreeG[T]) IsoCopy() *BTreeG[T] { if tr.lock(true) { defer tr.unlock(true) } tr.isoid = newIsoID() tr2 := new(BTreeG[T]) *tr2 = *tr tr2.mu = new(sync.RWMutex) tr2.isoid = newIsoID() return tr2 } func (tr *BTreeG[T]) lock(write bool) bool { if tr.locks { if write { tr.mu.Lock() } else { tr.mu.RLock() } } return tr.locks } func (tr *BTreeG[T]) unlock(write bool) { if write { tr.mu.Unlock() } else { tr.mu.RUnlock() } } // Iter represents an iterator type IterG[T any] struct { tr *BTreeG[T] mut bool locked bool seeked bool atstart bool atend bool stack []iterStackItemG[T] item T } type iterStackItemG[T any] struct { n *node[T] i int } // Iter returns a read-only iterator. // The Release method must be called finished with iterator. func (tr *BTreeG[T]) Iter() IterG[T] { return tr.iter(false) } func (tr *BTreeG[T]) IterMut() IterG[T] { return tr.iter(true) } func (tr *BTreeG[T]) iter(mut bool) IterG[T] { var iter IterG[T] iter.tr = tr iter.mut = mut iter.locked = tr.lock(iter.mut) return iter } // Seek to item greater-or-equal-to key. // Returns false if there was no item found. func (iter *IterG[T]) Seek(key T) bool { if iter.tr == nil { return false } iter.seeked = true iter.stack = iter.stack[:0] if iter.tr.root == nil { return false } n := iter.tr.isoLoad(&iter.tr.root, iter.mut) for { i, found := iter.tr.find(n, key, nil, 0) iter.stack = append(iter.stack, iterStackItemG[T]{n, i}) if found { iter.item = n.items[i] return true } if n.leaf() { iter.stack[len(iter.stack)-1].i-- return iter.Next() } n = iter.tr.isoLoad(&(*n.children)[i], iter.mut) } } // First moves iterator to first item in tree. // Returns false if the tree is empty. func (iter *IterG[T]) First() bool { if iter.tr == nil { return false } iter.atend = false iter.atstart = false iter.seeked = true iter.stack = iter.stack[:0] if iter.tr.root == nil { return false } n := iter.tr.isoLoad(&iter.tr.root, iter.mut) for { iter.stack = append(iter.stack, iterStackItemG[T]{n, 0}) if n.leaf() { break } n = iter.tr.isoLoad(&(*n.children)[0], iter.mut) } s := &iter.stack[len(iter.stack)-1] iter.item = s.n.items[s.i] return true } // Last moves iterator to last item in tree. // Returns false if the tree is empty. func (iter *IterG[T]) Last() bool { if iter.tr == nil { return false } iter.seeked = true iter.stack = iter.stack[:0] if iter.tr.root == nil { return false } n := iter.tr.isoLoad(&iter.tr.root, iter.mut) for { iter.stack = append(iter.stack, iterStackItemG[T]{n, len(n.items)}) if n.leaf() { iter.stack[len(iter.stack)-1].i-- break } n = iter.tr.isoLoad(&(*n.children)[len(n.items)], iter.mut) } s := &iter.stack[len(iter.stack)-1] iter.item = s.n.items[s.i] return true } // Release the iterator. func (iter *IterG[T]) Release() { if iter.tr == nil { return } if iter.locked { iter.tr.unlock(iter.mut) iter.locked = false } iter.stack = nil iter.tr = nil } // Next moves iterator to the next item in iterator. // Returns false if the tree is empty or the iterator is at the end of // the tree. func (iter *IterG[T]) Next() bool { if iter.tr == nil { return false } if !iter.seeked { return iter.First() } if len(iter.stack) == 0 { if iter.atstart { return iter.First() && iter.Next() } return false } s := &iter.stack[len(iter.stack)-1] s.i++ if s.n.leaf() { if s.i == len(s.n.items) { for { iter.stack = iter.stack[:len(iter.stack)-1] if len(iter.stack) == 0 { iter.atend = true return false } s = &iter.stack[len(iter.stack)-1] if s.i < len(s.n.items) { break } } } } else { n := iter.tr.isoLoad(&(*s.n.children)[s.i], iter.mut) for { iter.stack = append(iter.stack, iterStackItemG[T]{n, 0}) if n.leaf() { break } n = iter.tr.isoLoad(&(*n.children)[0], iter.mut) } } s = &iter.stack[len(iter.stack)-1] iter.item = s.n.items[s.i] return true } // Prev moves iterator to the previous item in iterator. // Returns false if the tree is empty or the iterator is at the beginning of // the tree. func (iter *IterG[T]) Prev() bool { if iter.tr == nil { return false } if !iter.seeked { return false } if len(iter.stack) == 0 { if iter.atend { return iter.Last() && iter.Prev() } return false } s := &iter.stack[len(iter.stack)-1] if s.n.leaf() { s.i-- if s.i == -1 { for { iter.stack = iter.stack[:len(iter.stack)-1] if len(iter.stack) == 0 { iter.atstart = true return false } s = &iter.stack[len(iter.stack)-1] s.i-- if s.i > -1 { break } } } } else { n := iter.tr.isoLoad(&(*s.n.children)[s.i], iter.mut) for { iter.stack = append(iter.stack, iterStackItemG[T]{n, len(n.items)}) if n.leaf() { iter.stack[len(iter.stack)-1].i-- break } n = iter.tr.isoLoad(&(*n.children)[len(n.items)], iter.mut) } } s = &iter.stack[len(iter.stack)-1] iter.item = s.n.items[s.i] return true } // Item returns the current iterator item. func (iter *IterG[T]) Item() T { return iter.item } // Items returns all the items in order. func (tr *BTreeG[T]) Items() []T { return tr.items(false) } func (tr *BTreeG[T]) ItemsMut() []T { return tr.items(true) } func (tr *BTreeG[T]) items(mut bool) []T { if tr.lock(mut) { defer tr.unlock(mut) } items := make([]T, 0, tr.Len()) if tr.root != nil { items = tr.nodeItems(&tr.root, items, mut) } return items } func (tr *BTreeG[T]) nodeItems(cn **node[T], items []T, mut bool) []T { n := tr.isoLoad(cn, mut) if n.leaf() { return append(items, n.items...) } for i := 0; i < len(n.items); i++ { items = tr.nodeItems(&(*n.children)[i], items, mut) items = append(items, n.items[i]) } return tr.nodeItems(&(*n.children)[len(*n.children)-1], items, mut) } // Clear will delete all items. func (tr *BTreeG[T]) Clear() { if tr.lock(true) { defer tr.unlock(true) } tr.root = nil tr.count = 0 } // Generic BTree // // Deprecated: use BTreeG type Generic[T any] struct { *BTreeG[T] } // NewGeneric returns a generic BTree // // Deprecated: use NewBTreeG func NewGeneric[T any](less func(a, b T) bool) *Generic[T] { return &Generic[T]{NewBTreeGOptions(less, Options{})} } // NewGenericOptions returns a generic BTree // // Deprecated: use NewBTreeGOptions func NewGenericOptions[T any](less func(a, b T) bool, opts Options, ) *Generic[T] { return &Generic[T]{NewBTreeGOptions(less, opts)} } func (tr *Generic[T]) Copy() *Generic[T] { return &Generic[T]{tr.BTreeG.Copy()} }