Update btree

This commit is contained in:
tidwall 2021-04-26 10:50:09 -07:00
parent c01cb2a5eb
commit 48b46f2752
5 changed files with 228 additions and 62 deletions

2
go.mod
View File

@ -14,7 +14,7 @@ require (
github.com/nats-io/nats.go v1.10.0
github.com/peterh/liner v1.2.1
github.com/streadway/amqp v1.0.0
github.com/tidwall/btree v0.4.2
github.com/tidwall/btree v0.5.0
github.com/tidwall/buntdb v1.2.3
github.com/tidwall/geoindex v1.4.3
github.com/tidwall/geojson v1.2.7

3
go.sum
View File

@ -182,8 +182,9 @@ github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81P
github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/tidwall/btree v0.4.2 h1:aLwwJlG+InuFzdAPuBf9YCAR1LvSQ9zhC5aorFPlIPs=
github.com/tidwall/btree v0.4.2/go.mod h1:huei1BkDWJ3/sLXmO+bsCNELL+Bp2Kks9OLyQFkzvA8=
github.com/tidwall/btree v0.5.0 h1:IBfCtOj4uOMQcodv3wzYVo0zPqSJObm71mE039/dlXY=
github.com/tidwall/btree v0.5.0/go.mod h1:TzIRzen6yHbibdSfK6t8QimqbUnoxUSrZfeW7Uob0q4=
github.com/tidwall/buntdb v1.2.3 h1:AoGVe4yrhKmnEPHrPrW5EUOATHOCIk4VtFvd8xn/ZtU=
github.com/tidwall/buntdb v1.2.3/go.mod h1:+i/gBwYOHWG19wLgwMXFLkl00twh9+VWkkaOhuNQ4PA=
github.com/tidwall/cities v0.1.0 h1:CVNkmMf7NEC9Bvokf5GoSsArHCKRMTgLuubRTHnH0mE=

View File

@ -6,7 +6,7 @@ package btree
import "sync"
const maxItems = 255
const maxItems = 255 // max items per node. max children is +1
const minItems = maxItems * 40 / 100
type cow struct {
@ -17,18 +17,18 @@ type node struct {
cow *cow
leaf bool
numItems int16
count int
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
mu *sync.RWMutex
cow *cow
root *node
count int
less func(a, b interface{}) bool
}
func (tr *BTree) newNode(leaf bool) *node {
@ -43,6 +43,7 @@ func (tr *BTree) newNode(leaf bool) *node {
// 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
}
@ -68,18 +69,32 @@ func (n *node) find(key interface{}, less func(a, b interface{}) bool,
) (index int16, found bool) {
low := int16(0)
high := n.numItems - 1
if hint != nil && depth < 8 {
if hint != nil && depth < 8 && hint.used[depth] {
index = int16(hint.path[depth])
if index > n.numItems-1 {
if index >= n.numItems {
// tail item
if less(n.items[n.numItems-1], key) {
if less(key, n.items[n.numItems-1]) {
index = n.numItems - 1
found = true
goto path_match
} else {
index = n.numItems
goto path_match
}
}
index = n.numItems - 1
}
if less(key, n.items[index]) {
if index == 0 || less(n.items[index-1], key) {
goto path_match
}
high = index - 1
} else if less(n.items[index], key) {
low = index + 1
} else {
found = true
goto done
goto path_match
}
}
for low <= high {
@ -97,13 +112,16 @@ func (n *node) find(key interface{}, less func(a, b interface{}) bool,
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)
}
if hint == nil || depth >= 8 {
return index, found
}
path_match:
hint.used[depth] = true
if n.leaf && found {
hint.path[depth] = byte(index + 1)
} else {
hint.path[depth] = byte(index)
}
return index, found
}
@ -124,14 +142,14 @@ func (tr *BTree) setHint(item interface{}, hint *PathHint) (prev interface{}) {
tr.root = tr.newNode(true)
tr.root.items[0] = item
tr.root.numItems = 1
tr.length = 1
tr.root.count = 1
tr.count = 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)
@ -140,8 +158,9 @@ func (tr *BTree) setHint(item interface{}, hint *PathHint) (prev interface{}) {
tr.root.items[0] = median
tr.root.children[1] = right
tr.root.numItems = 1
tr.root.count = n.count + 1 + right.count
}
tr.length++
tr.count++
return prev
}
@ -167,10 +186,25 @@ func (tr *BTree) nodeSplit(n *node) (right *node, median interface{}) {
n.items[i] = nil
}
n.numItems = maxItems / 2
// update counts
n.updateCount()
right.updateCount()
return right, median
}
//go:noinline
func (n *node) updateCount() {
n.count = int(n.numItems)
if !n.leaf {
for i := 0; i <= int(n.numItems); i++ {
n.count += n.children[i].count
}
}
}
// This operation should not be inlined because it's expensive and rarely
// called outside of heavy copy-on-write situations. Marking it "noinline"
// allows for the parent cowLoad to be inlined.
// go:noinline
func (tr *BTree) copy(n *node) *node {
n2 := *n
n2.cow = tr.cow
@ -182,7 +216,7 @@ func (tr *BTree) copy(n *node) *node {
return &n2
}
// cowLoad loaded the provide node and, if needed, performs a copy-on-write.
// cowLoad loads 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)
@ -204,6 +238,7 @@ func (tr *BTree) nodeSet(cn **node, item interface{},
copy(n.items[i+1:n.numItems+1], n.items[i:n.numItems])
n.items[i] = item
n.numItems++
n.count++
return nil
}
prev = tr.nodeSet(&n.children[i], item, less, hint, depth+1)
@ -218,7 +253,8 @@ func (tr *BTree) nodeSet(cn **node, item interface{},
n.children[i+1] = right
n.numItems++
}
return prev
n.count++
return nil
}
func (n *node) scan(iter func(item interface{}) bool) bool {
@ -270,7 +306,7 @@ func (tr *BTree) GetHint(key interface{}, hint *PathHint) interface{} {
// Len returns the number of items in the tree
func (tr *BTree) Len() int {
return tr.length
return tr.count
}
// Delete a value for a key
@ -294,12 +330,11 @@ func (tr *BTree) deleteHint(key interface{}, hint *PathHint) interface{} {
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.count--
if tr.count == 0 {
tr.root = nil
}
return prev
@ -323,6 +358,7 @@ func (tr *BTree) delete(cn **node, max bool, key interface{},
copy(n.items[i:], n.items[i+1:n.numItems])
n.items[n.numItems-1] = nil
n.numItems--
n.count--
return prev
}
return nil
@ -344,6 +380,7 @@ func (tr *BTree) delete(cn **node, max bool, key interface{},
if prev == nil {
return nil
}
n.count--
if n.children[i].numItems >= minItems {
return prev
}
@ -364,10 +401,11 @@ func (tr *BTree) delete(cn **node, max bool, key interface{},
n.children[i+1].children[:n.children[i+1].numItems+1])
}
n.children[i].numItems += n.children[i+1].numItems + 1
n.children[i].count += n.children[i+1].count + 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.items[n.numItems-1] = nil
n.children[n.numItems] = nil
n.numItems--
} else if n.children[i].numItems > n.children[i+1].numItems {
// move left -> right
@ -381,30 +419,42 @@ func (tr *BTree) delete(cn **node, max bool, key interface{},
if !n.children[0].leaf {
n.children[i+1].children[0] =
n.children[i].children[n.children[i].numItems]
n.children[i+1].count += n.children[i+1].children[0].count
}
n.children[i+1].numItems++
n.children[i+1].count++
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].count -= n.children[i+1].children[0].count
}
n.children[i].numItems--
n.children[i].count--
} else {
// move right -> left
// move left <- right
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].count +=
n.children[i].children[n.children[i].numItems+1].count
}
n.children[i].numItems++
n.children[i].count++
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])
n.children[i+1].items[n.children[i+1].numItems-1] = nil
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].children[n.children[i+1].numItems] = nil
n.children[i+1].count -=
n.children[i].children[n.children[i].numItems].count
}
n.children[i+1].numItems--
n.children[i+1].count--
}
return prev
}
@ -520,31 +570,35 @@ func (tr *BTree) Load(item interface{}) interface{} {
}
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
}
if tr.root == nil {
return tr.setHint(item, nil)
}
prev := tr.setHint(item, nil)
if prev != nil {
return prev
}
n := tr.root
n := tr.cowLoad(&tr.root)
for {
n.count++ // optimistically update counts
if n.leaf {
if n.numItems < maxItems-2 {
if tr.less(n.items[n.numItems-1], item) {
n.items[n.numItems] = item
n.numItems++
tr.count++
return nil
}
}
break
}
n = tr.cowLoad(&n.children[n.numItems])
}
// revert the counts
n = tr.root
for {
n.count--
if n.leaf {
tr.lnode = n
break
}
n = n.children[n.numItems]
}
return nil
return tr.setHint(item, nil)
}
// Min returns the minimum item in tree.
@ -589,25 +643,36 @@ func (tr *BTree) PopMin() interface{} {
if tr.root == nil {
return nil
}
tr.lnode = nil
n := tr.cowLoad(&tr.root)
var item interface{}
for {
n.count-- // optimistically update counts
if n.leaf {
item := n.items[0]
item = n.items[0]
if n.numItems == minItems {
return tr.deleteHint(item, nil)
break
}
copy(n.items[:], n.items[1:])
n.items[n.numItems-1] = nil
n.numItems--
tr.length--
if tr.length == 0 {
tr.count--
if tr.count == 0 {
tr.root = nil
}
return item
}
n = tr.cowLoad(&n.children[0])
}
// 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 minimum item in tree and returns it.
@ -618,24 +683,122 @@ func (tr *BTree) PopMax() interface{} {
if tr.root == nil {
return nil
}
tr.lnode = nil
n := tr.cowLoad(&tr.root)
var item interface{}
for {
n.count-- // optimistically update counts
if n.leaf {
item := n.items[n.numItems-1]
item = n.items[n.numItems-1]
if n.numItems == minItems {
return tr.deleteHint(item, nil)
break
}
n.items[n.numItems-1] = nil
n.numItems--
tr.length--
if tr.length == 0 {
tr.count--
if tr.count == 0 {
tr.root = nil
}
return item
}
n = tr.cowLoad(&n.children[n.numItems])
}
// revert the counts
n = tr.root
for {
n.count++
if n.leaf {
break
}
n = n.children[n.numItems]
}
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 *BTree) GetAt(index int) interface{} {
tr.mu.RLock()
defer tr.mu.RUnlock()
if tr.root == nil || index < 0 || index >= tr.count {
return nil
}
n := tr.root
for {
if n.leaf {
return n.items[index]
}
i := 0
for ; i < int(n.numItems); i++ {
if index < n.children[i].count {
break
} else if index == n.children[i].count {
return n.items[i]
}
index -= n.children[i].count + 1
}
n = n.children[i]
}
}
// DeleteAt deletes the item at index.
// Return nil if the tree is empty or the index is out of bounds.
func (tr *BTree) DeleteAt(index int) interface{} {
tr.mu.Lock()
defer tr.mu.Unlock()
if tr.root == nil || index < 0 || index >= tr.count {
return nil
}
var pathbuf [8]uint8 // track the path
path := pathbuf[:0]
var item interface{}
n := tr.cowLoad(&tr.root)
outer:
for {
n.count-- // optimistically update counts
if n.leaf {
// the index is the item position
item = n.items[index]
if n.numItems == minItems {
path = append(path, uint8(index))
break outer
}
copy(n.items[index:], n.items[index+1:n.numItems])
n.items[n.numItems-1] = nil
n.numItems--
tr.count--
if tr.count == 0 {
tr.root = nil
}
return item
}
i := 0
for ; i < int(n.numItems); 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.cowLoad(&n.children[i])
}
// revert the counts
var hint PathHint
n = tr.root
for i := 0; i < len(path); i++ {
if i < len(hint.path) {
hint.path[i] = 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.
@ -683,7 +846,6 @@ func (n *node) walk(iter func(item []interface{})) {
// shadowed copy.
func (tr *BTree) Copy() *BTree {
tr.mu.Lock()
tr.lnode = nil
tr.cow = new(cow)
tr2 := *tr
tr2.mu = new(sync.RWMutex)

3
vendor/github.com/tidwall/btree/go.mod generated vendored Normal file
View File

@ -0,0 +1,3 @@
module github.com/tidwall/btree
go 1.16

2
vendor/modules.txt vendored
View File

@ -145,7 +145,7 @@ github.com/rcrowley/go-metrics
# github.com/streadway/amqp v1.0.0
## explicit
github.com/streadway/amqp
# github.com/tidwall/btree v0.4.2
# github.com/tidwall/btree v0.5.0
## explicit
github.com/tidwall/btree
# github.com/tidwall/buntdb v1.2.3