| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391 |
- // 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()}
- }
|