123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693 |
- // 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"
-
- const maxItems = 255
- const minItems = maxItems * 40 / 100
-
- type cow struct {
- _ int // it cannot be an empty struct
- }
-
- type node struct {
- cow *cow
- leaf bool
- numItems int16
- items [maxItems]interface{}
- children *[maxItems + 1]*node
- }
-
- // BTree is an ordered set items
- type BTree struct {
- mu *sync.RWMutex
- cow *cow
- root *node
- length int
- less func(a, b interface{}) bool
- lnode *node
- }
-
- func (tr *BTree) newNode(leaf bool) *node {
- n := &node{leaf: leaf}
- if !leaf {
- n.children = new([maxItems + 1]*node)
- }
- n.cow = tr.cow
- return n
- }
-
- // PathHint is a utility type used with the *Hint() functions. Hints provide
- // faster operations for clustered keys.
- type PathHint struct {
- path [8]uint8
- }
-
- // New returns a new BTree
- func New(less func(a, b interface{}) bool) *BTree {
- if less == nil {
- panic("nil less")
- }
- tr := new(BTree)
- tr.mu = new(sync.RWMutex)
- tr.less = less
- return tr
- }
-
- // Less is a convenience function that performs a comparison of two items
- // using the same "less" function provided to New.
- func (tr *BTree) Less(a, b interface{}) bool {
- return tr.less(a, b)
- }
-
- func (n *node) find(key interface{}, less func(a, b interface{}) bool,
- hint *PathHint, depth int,
- ) (index int16, found bool) {
- low := int16(0)
- high := n.numItems - 1
- if hint != nil && depth < 8 {
- index = int16(hint.path[depth])
- if index > n.numItems-1 {
- index = n.numItems - 1
- }
- if less(key, n.items[index]) {
- high = index - 1
- } else if less(n.items[index], key) {
- low = index + 1
- } else {
- found = true
- goto done
- }
- }
- for low <= high {
- mid := low + ((high+1)-low)/2
- if !less(key, n.items[mid]) {
- low = mid + 1
- } else {
- high = mid - 1
- }
- }
- if low > 0 && !less(n.items[low-1], key) {
- index = low - 1
- found = true
- } else {
- index = low
- found = false
- }
- done:
- if hint != nil && depth < 8 {
- if n.leaf && found {
- hint.path[depth] = byte(index + 1)
- } else {
- hint.path[depth] = byte(index)
- }
- }
- return index, found
- }
-
- // SetHint sets or replace a value for a key using a path hint
- func (tr *BTree) SetHint(item interface{}, hint *PathHint) (prev interface{}) {
- if item == nil {
- panic("nil item")
- }
- tr.mu.Lock()
- prev = tr.setHint(item, hint)
- tr.mu.Unlock()
- return prev
- }
-
- func (tr *BTree) setHint(item interface{}, hint *PathHint) (prev interface{}) {
- if tr.root == nil {
- tr.root = tr.newNode(true)
- tr.root.items[0] = item
- tr.root.numItems = 1
- tr.length = 1
- return
- }
- prev = tr.nodeSet(&tr.root, item, tr.less, hint, 0)
- if prev != nil {
- return prev
- }
- tr.lnode = nil
- if tr.root.numItems == maxItems {
- n := tr.cowLoad(&tr.root)
- right, median := tr.nodeSplit(n)
- tr.root = tr.newNode(false)
- tr.root.children[0] = n
- tr.root.items[0] = median
- tr.root.children[1] = right
- tr.root.numItems = 1
- }
- tr.length++
- return prev
- }
-
- // Set or replace a value for a key
- func (tr *BTree) Set(item interface{}) (prev interface{}) {
- return tr.SetHint(item, nil)
- }
-
- func (tr *BTree) nodeSplit(n *node) (right *node, median interface{}) {
- right = tr.newNode(n.leaf)
- median = n.items[maxItems/2]
- copy(right.items[:maxItems/2], n.items[maxItems/2+1:])
- if !n.leaf {
- copy(right.children[:maxItems/2+1], n.children[maxItems/2+1:])
- }
- right.numItems = maxItems / 2
- if !n.leaf {
- for i := maxItems/2 + 1; i < maxItems+1; i++ {
- n.children[i] = nil
- }
- }
- for i := maxItems / 2; i < maxItems; i++ {
- n.items[i] = nil
- }
- n.numItems = maxItems / 2
- return right, median
- }
-
- //go:noinline
- func (tr *BTree) copy(n *node) *node {
- n2 := *n
- n2.cow = tr.cow
- copy(n2.items[:], n.items[:])
- if n.children != nil {
- n2.children = new([maxItems + 1]*node)
- copy(n2.children[:], n.children[:])
- }
- return &n2
- }
-
- // cowLoad loaded the provide node and, if needed, performs a copy-on-write.
- func (tr *BTree) cowLoad(cn **node) *node {
- if (*cn).cow != tr.cow {
- *cn = tr.copy(*cn)
- }
- return *cn
- }
-
- func (tr *BTree) nodeSet(cn **node, item interface{},
- less func(a, b interface{}) bool, hint *PathHint, depth int,
- ) (prev interface{}) {
- n := tr.cowLoad(cn)
- i, found := n.find(item, less, hint, depth)
- if found {
- prev = n.items[i]
- n.items[i] = item
- return prev
- }
- if n.leaf {
- copy(n.items[i+1:n.numItems+1], n.items[i:n.numItems])
- n.items[i] = item
- n.numItems++
- return nil
- }
- prev = tr.nodeSet(&n.children[i], item, less, hint, depth+1)
- if prev != nil {
- return prev
- }
- if n.children[i].numItems == maxItems {
- right, median := tr.nodeSplit(n.children[i])
- copy(n.children[i+1:], n.children[i:])
- copy(n.items[i+1:], n.items[i:])
- n.items[i] = median
- n.children[i+1] = right
- n.numItems++
- }
- return prev
- }
-
- func (n *node) scan(iter func(item interface{}) bool) bool {
- if n.leaf {
- for i := int16(0); i < n.numItems; i++ {
- if !iter(n.items[i]) {
- return false
- }
- }
- return true
- }
- for i := int16(0); i < n.numItems; i++ {
- if !n.children[i].scan(iter) {
- return false
- }
- if !iter(n.items[i]) {
- return false
- }
- }
- return n.children[n.numItems].scan(iter)
- }
-
- // Get a value for key
- func (tr *BTree) Get(key interface{}) interface{} {
- return tr.GetHint(key, nil)
- }
-
- // GetHint gets a value for key using a path hint
- func (tr *BTree) GetHint(key interface{}, hint *PathHint) interface{} {
- tr.mu.RLock()
- defer tr.mu.RUnlock()
- if tr.root == nil || key == nil {
- return nil
- }
- depth := 0
- n := tr.root
- for {
- i, found := n.find(key, tr.less, hint, depth)
- if found {
- return n.items[i]
- }
- if n.leaf {
- return nil
- }
- n = n.children[i]
- depth++
- }
- }
-
- // Len returns the number of items in the tree
- func (tr *BTree) Len() int {
- return tr.length
- }
-
- // Delete a value for a key
- func (tr *BTree) Delete(key interface{}) interface{} {
- return tr.DeleteHint(key, nil)
- }
-
- // DeleteHint deletes a value for a key using a path hint
- func (tr *BTree) DeleteHint(key interface{}, hint *PathHint) interface{} {
- tr.mu.Lock()
- prev := tr.deleteHint(key, hint)
- tr.mu.Unlock()
- return prev
- }
-
- func (tr *BTree) deleteHint(key interface{}, hint *PathHint) interface{} {
- if tr.root == nil || key == nil {
- return nil
- }
- prev := tr.delete(&tr.root, false, key, tr.less, hint, 0)
- if prev == nil {
- return nil
- }
- tr.lnode = nil
- if tr.root.numItems == 0 && !tr.root.leaf {
- tr.root = tr.root.children[0]
- }
- tr.length--
- if tr.length == 0 {
- tr.root = nil
- }
- return prev
- }
-
- func (tr *BTree) delete(cn **node, max bool, key interface{},
- less func(a, b interface{}) bool, hint *PathHint, depth int,
- ) interface{} {
- n := tr.cowLoad(cn)
- var i int16
- var found bool
- if max {
- i, found = n.numItems-1, true
- } else {
- i, found = n.find(key, less, hint, depth)
- }
- if n.leaf {
- if found {
- prev := n.items[i]
- // found the items at the leaf, remove it and return.
- copy(n.items[i:], n.items[i+1:n.numItems])
- n.items[n.numItems-1] = nil
- n.numItems--
- return prev
- }
- return nil
- }
-
- var prev interface{}
- if found {
- if max {
- i++
- prev = tr.delete(&n.children[i], true, "", less, nil, 0)
- } else {
- prev = n.items[i]
- maxItem := tr.delete(&n.children[i], true, "", less, nil, 0)
- n.items[i] = maxItem
- }
- } else {
- prev = tr.delete(&n.children[i], max, key, less, hint, depth+1)
- }
- if prev == nil {
- return nil
- }
- if n.children[i].numItems >= minItems {
- return prev
- }
-
- // merge / rebalance nodes
- if i == n.numItems {
- i--
- }
- n.children[i] = tr.cowLoad(&n.children[i])
- n.children[i+1] = tr.cowLoad(&n.children[i+1])
- if n.children[i].numItems+n.children[i+1].numItems+1 < maxItems {
- // merge left + item + right
- n.children[i].items[n.children[i].numItems] = n.items[i]
- copy(n.children[i].items[n.children[i].numItems+1:],
- n.children[i+1].items[:n.children[i+1].numItems])
- if !n.children[0].leaf {
- copy(n.children[i].children[n.children[i].numItems+1:],
- n.children[i+1].children[:n.children[i+1].numItems+1])
- }
- n.children[i].numItems += n.children[i+1].numItems + 1
- copy(n.items[i:], n.items[i+1:n.numItems])
- copy(n.children[i+1:], n.children[i+2:n.numItems+1])
- n.items[n.numItems] = nil
- n.children[n.numItems+1] = nil
- n.numItems--
- } else if n.children[i].numItems > n.children[i+1].numItems {
- // move left -> right
- copy(n.children[i+1].items[1:],
- n.children[i+1].items[:n.children[i+1].numItems])
- if !n.children[0].leaf {
- copy(n.children[i+1].children[1:],
- n.children[i+1].children[:n.children[i+1].numItems+1])
- }
- n.children[i+1].items[0] = n.items[i]
- if !n.children[0].leaf {
- n.children[i+1].children[0] =
- n.children[i].children[n.children[i].numItems]
- }
- n.children[i+1].numItems++
- n.items[i] = n.children[i].items[n.children[i].numItems-1]
- n.children[i].items[n.children[i].numItems-1] = nil
- if !n.children[0].leaf {
- n.children[i].children[n.children[i].numItems] = nil
- }
- n.children[i].numItems--
- } else {
- // move right -> left
- n.children[i].items[n.children[i].numItems] = n.items[i]
- if !n.children[0].leaf {
- n.children[i].children[n.children[i].numItems+1] =
- n.children[i+1].children[0]
- }
- n.children[i].numItems++
- n.items[i] = n.children[i+1].items[0]
- copy(n.children[i+1].items[:],
- n.children[i+1].items[1:n.children[i+1].numItems])
- if !n.children[0].leaf {
- copy(n.children[i+1].children[:],
- n.children[i+1].children[1:n.children[i+1].numItems+1])
- }
- n.children[i+1].numItems--
- }
- return prev
- }
-
- // 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 *BTree) Ascend(pivot interface{}, iter func(item interface{}) bool) {
- tr.mu.RLock()
- defer tr.mu.RUnlock()
- if tr.root == nil {
- return
- }
- if pivot == nil {
- tr.root.scan(iter)
- } else if tr.root != nil {
- tr.root.ascend(pivot, tr.less, nil, 0, iter)
- }
- }
-
- func (n *node) ascend(pivot interface{}, less func(a, b interface{}) bool,
- hint *PathHint, depth int, iter func(item interface{}) bool,
- ) bool {
- i, found := n.find(pivot, less, hint, depth)
- if !found {
- if !n.leaf {
- if !n.children[i].ascend(pivot, less, hint, depth+1, iter) {
- return false
- }
- }
- }
- for ; i < n.numItems; i++ {
- if !iter(n.items[i]) {
- return false
- }
- if !n.leaf {
- if !n.children[i+1].scan(iter) {
- return false
- }
- }
- }
- return true
- }
-
- func (n *node) reverse(iter func(item interface{}) bool) bool {
- if n.leaf {
- for i := n.numItems - 1; i >= 0; i-- {
- if !iter(n.items[i]) {
- return false
- }
- }
- return true
- }
- if !n.children[n.numItems].reverse(iter) {
- return false
- }
- for i := n.numItems - 1; i >= 0; i-- {
- if !iter(n.items[i]) {
- return false
- }
- if !n.children[i].reverse(iter) {
- 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 *BTree) Descend(pivot interface{}, iter func(item interface{}) bool) {
- tr.mu.RLock()
- defer tr.mu.RUnlock()
- if tr.root == nil {
- return
- }
- if pivot == nil {
- tr.root.reverse(iter)
- } else if tr.root != nil {
- tr.root.descend(pivot, tr.less, nil, 0, iter)
- }
- }
-
- func (n *node) descend(pivot interface{}, less func(a, b interface{}) bool,
- hint *PathHint, depth int, iter func(item interface{}) bool,
- ) bool {
- i, found := n.find(pivot, less, hint, depth)
- if !found {
- if !n.leaf {
- if !n.children[i].descend(pivot, less, hint, depth+1, iter) {
- return false
- }
- }
- i--
- }
- for ; i >= 0; i-- {
- if !iter(n.items[i]) {
- return false
- }
- if !n.leaf {
- if !n.children[i].reverse(iter) {
- return false
- }
- }
- }
- return true
- }
-
- // Load is for bulk loading pre-sorted items
- func (tr *BTree) Load(item interface{}) interface{} {
- if item == nil {
- panic("nil item")
- }
- tr.mu.Lock()
- defer tr.mu.Unlock()
-
- // Load does not need a cowGrid because the Copy operation sets the
- // lnode to nil.
-
- if tr.lnode != nil && tr.lnode.numItems < maxItems-2 {
- if tr.less(tr.lnode.items[tr.lnode.numItems-1], item) {
- tr.lnode.items[tr.lnode.numItems] = item
- tr.lnode.numItems++
- tr.length++
- return nil
- }
- }
- prev := tr.setHint(item, nil)
- if prev != nil {
- return prev
- }
- n := tr.root
- for {
- if n.leaf {
- tr.lnode = n
- break
- }
- n = n.children[n.numItems]
- }
- return nil
- }
-
- // Min returns the minimum item in tree.
- // Returns nil if the tree has no items.
- func (tr *BTree) Min() interface{} {
- tr.mu.RLock()
- defer tr.mu.RUnlock()
- if tr.root == nil {
- return nil
- }
- n := tr.root
- for {
- if n.leaf {
- return n.items[0]
- }
- n = n.children[0]
- }
- }
-
- // Max returns the maximum item in tree.
- // Returns nil if the tree has no items.
- func (tr *BTree) Max() interface{} {
- tr.mu.RLock()
- defer tr.mu.RUnlock()
- if tr.root == nil {
- return nil
- }
- n := tr.root
- for {
- if n.leaf {
- return n.items[n.numItems-1]
- }
- n = n.children[n.numItems]
- }
- }
-
- // PopMin removes the minimum item in tree and returns it.
- // Returns nil if the tree has no items.
- func (tr *BTree) PopMin() interface{} {
- tr.mu.Lock()
- defer tr.mu.Unlock()
- if tr.root == nil {
- return nil
- }
- tr.lnode = nil
- n := tr.cowLoad(&tr.root)
- for {
- if n.leaf {
- item := n.items[0]
- if n.numItems == minItems {
- return tr.deleteHint(item, nil)
- }
- copy(n.items[:], n.items[1:])
- n.items[n.numItems-1] = nil
- n.numItems--
- tr.length--
- if tr.length == 0 {
- tr.root = nil
- }
- return item
- }
- n = tr.cowLoad(&n.children[0])
- }
- }
-
- // PopMax removes the minimum item in tree and returns it.
- // Returns nil if the tree has no items.
- func (tr *BTree) PopMax() interface{} {
- tr.mu.Lock()
- defer tr.mu.Unlock()
- if tr.root == nil {
- return nil
- }
- tr.lnode = nil
- n := tr.cowLoad(&tr.root)
- for {
- if n.leaf {
- item := n.items[n.numItems-1]
- if n.numItems == minItems {
- return tr.deleteHint(item, nil)
- }
- n.items[n.numItems-1] = nil
- n.numItems--
- tr.length--
- if tr.length == 0 {
- tr.root = nil
- }
- return item
- }
- n = tr.cowLoad(&n.children[n.numItems])
- }
- }
-
- // Height returns the height of the tree.
- // Returns zero if tree has no items.
- func (tr *BTree) Height() int {
- tr.mu.RLock()
- defer tr.mu.RUnlock()
- var height int
- if tr.root != nil {
- n := tr.root
- for {
- height++
- if n.leaf {
- break
- }
- n = n.children[n.numItems]
- }
- }
- return height
- }
-
- // Walk iterates over all items in tree, in order.
- // The items param will contain one or more items.
- func (tr *BTree) Walk(iter func(item []interface{})) {
- tr.mu.RLock()
- defer tr.mu.RUnlock()
- if tr.root != nil {
- tr.root.walk(iter)
- }
- }
-
- func (n *node) walk(iter func(item []interface{})) {
- if n.leaf {
- iter(n.items[:n.numItems])
- } else {
- for i := int16(0); i < n.numItems; i++ {
- n.children[i].walk(iter)
- iter(n.items[i : i+1])
- }
- n.children[n.numItems].walk(iter)
- }
- }
-
- // Copy the tree. This operation is very fast because it only performs a
- // shadowed copy.
- func (tr *BTree) Copy() *BTree {
- tr.mu.Lock()
- tr.lnode = nil
- tr.cow = new(cow)
- tr2 := *tr
- tr2.mu = new(sync.RWMutex)
- tr2.cow = new(cow)
- tr.mu.Unlock()
- return &tr2
- }
|