Update btree

go.mod

@@ -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

go.sum

@@ -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=

vendor/github.com/tidwall/btree/btree.go

@@ -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
+	mu    *sync.RWMutex
-	cow    *cow
+	cow   *cow
-	root   *node
+	root  *node
-	length int
+	count int
-	less   func(a, b interface{}) bool
+	less  func(a, b interface{}) bool
-	lnode  *node
 }
 
 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 hint != nil && depth < 8 {
+		return index, found
-		if n.leaf && found {
+	}
-			hint.path[depth] = byte(index + 1)
+
-		} else {
+path_match:
-			hint.path[depth] = byte(index)
+	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--
+	tr.count--
-	if tr.length == 0 {
+	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.items[n.numItems-1] = nil
-		n.children[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()
-
+	if tr.root == nil {
-	// Load does not need a cowGrid because the Copy operation sets the
+		return tr.setHint(item, nil)
-	// 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)
+	n := tr.cowLoad(&tr.root)
-	if prev != nil {
-		return prev
-	}
-	n := 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--
+			tr.count--
-			if tr.length == 0 {
+			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--
+			tr.count--
-			if tr.length == 0 {
+			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)

vendor/github.com/tidwall/btree/go.mod

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

vendor/modules.txt

@@ -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