3
0
mirror of https://github.com/ergochat/ergo.git synced 2024-11-28 15:09:28 +01:00
ergo/vendor/github.com/tidwall/btree/map.go

1057 lines
24 KiB
Go
Raw Permalink Normal View History

2022-06-16 19:36:02 +02:00
// Copyright 2020 Joshua J Baker. All rights reserved.
2023-01-15 14:26:32 +01:00
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.
2022-06-16 19:36:02 +02:00
package btree
2023-01-15 14:26:32 +01:00
import "sync/atomic"
2022-06-16 19:36:02 +02:00
2023-01-15 14:26:32 +01:00
type ordered interface {
~int | ~int8 | ~int16 | ~int32 | ~int64 |
~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr |
~float32 | ~float64 | ~string
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
type mapPair[K ordered, V any] struct {
// The `value` field should be before the `key` field because doing so
// allows for the Go compiler to optimize away the `value` field when
// it's a `struct{}`, which is the case for `btree.Set`.
value V
key K
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
type Map[K ordered, V any] struct {
cow uint64
root *mapNode[K, V]
count int
empty mapPair[K, V]
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
type mapNode[K ordered, V any] struct {
cow uint64
count int
items []mapPair[K, V]
children *[]*mapNode[K, V]
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
// 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 *Map[K, V]) copy(n *mapNode[K, V]) *mapNode[K, V] {
n2 := new(mapNode[K, V])
n2.cow = tr.cow
n2.count = n.count
n2.items = make([]mapPair[K, V], len(n.items), cap(n.items))
copy(n2.items, n.items)
if !n.leaf() {
n2.children = new([]*mapNode[K, V])
*n2.children = make([]*mapNode[K, V], len(*n.children), maxItems+1)
copy(*n2.children, *n.children)
}
return n2
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
// cowLoad loads the provided node and, if needed, performs a copy-on-write.
func (tr *Map[K, V]) cowLoad(cn **mapNode[K, V]) *mapNode[K, V] {
if (*cn).cow != tr.cow {
*cn = tr.copy(*cn)
}
return *cn
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) Copy() *Map[K, V] {
tr2 := new(Map[K, V])
*tr2 = *tr
tr2.cow = atomic.AddUint64(&gcow, 1)
tr.cow = atomic.AddUint64(&gcow, 1)
return tr2
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) newNode(leaf bool) *mapNode[K, V] {
n := new(mapNode[K, V])
n.cow = tr.cow
2022-06-16 19:36:02 +02:00
if !leaf {
2023-01-15 14:26:32 +01:00
n.children = new([]*mapNode[K, V])
2022-06-16 19:36:02 +02:00
}
return n
}
// leaf returns true if the node is a leaf.
2023-01-15 14:26:32 +01:00
func (n *mapNode[K, V]) leaf() bool {
2022-06-16 19:36:02 +02:00
return n.children == nil
}
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) bsearch(n *mapNode[K, V], key K) (index int, found bool) {
low, high := 0, len(n.items)
for low < high {
h := int(uint(low+high) >> 1)
if key >= n.items[h].key {
low = h + 1
2022-06-16 19:36:02 +02:00
} else {
2023-01-15 14:26:32 +01:00
high = h
2022-06-16 19:36:02 +02:00
}
}
2023-01-15 14:26:32 +01:00
if low > 0 && n.items[low-1].key >= key {
return low - 1, true
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
return low, false
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
// Set or replace a value for a key
func (tr *Map[K, V]) Set(key K, value V) (V, bool) {
item := mapPair[K, V]{key: key, value: value}
2022-06-16 19:36:02 +02:00
if tr.root == nil {
tr.root = tr.newNode(true)
2023-01-15 14:26:32 +01:00
tr.root.items = append([]mapPair[K, V]{}, item)
2022-06-16 19:36:02 +02:00
tr.root.count = 1
tr.count = 1
2023-01-15 14:26:32 +01:00
return tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
prev, replaced, split := tr.nodeSet(&tr.root, item)
2022-06-16 19:36:02 +02:00
if split {
2023-01-15 14:26:32 +01:00
left := tr.root
2022-06-16 19:36:02 +02:00
right, median := tr.nodeSplit(left)
tr.root = tr.newNode(false)
2023-01-15 14:26:32 +01:00
*tr.root.children = make([]*mapNode[K, V], 0, maxItems+1)
*tr.root.children = append([]*mapNode[K, V]{}, left, right)
tr.root.items = append([]mapPair[K, V]{}, median)
2022-06-16 19:36:02 +02:00
tr.root.updateCount()
2023-01-15 14:26:32 +01:00
return tr.Set(item.key, item.value)
2022-06-16 19:36:02 +02:00
}
if replaced {
return prev, true
}
tr.count++
2023-01-15 14:26:32 +01:00
return tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) nodeSplit(n *mapNode[K, V],
) (right *mapNode[K, V], median mapPair[K, V]) {
2022-06-16 19:36:02 +02:00
i := maxItems / 2
median = n.items[i]
2023-01-15 14:26:32 +01:00
const sliceItems = true
2022-06-16 19:36:02 +02:00
// right node
right = tr.newNode(n.leaf())
2023-01-15 14:26:32 +01:00
if sliceItems {
right.items = n.items[i+1:]
if !n.leaf() {
*right.children = (*n.children)[i+1:]
}
} else {
right.items = make([]mapPair[K, V], len(n.items[i+1:]), maxItems/2)
copy(right.items, n.items[i+1:])
if !n.leaf() {
*right.children = make([]*mapNode[K, V],
len((*n.children)[i+1:]), maxItems+1)
copy(*right.children, (*n.children)[i+1:])
}
2022-06-16 19:36:02 +02:00
}
right.updateCount()
2023-01-15 14:26:32 +01:00
// left node
if sliceItems {
n.items[i] = tr.empty
n.items = n.items[:i:i]
if !n.leaf() {
*n.children = (*n.children)[: i+1 : i+1]
}
} else {
for j := i; j < len(n.items); j++ {
n.items[j] = tr.empty
}
if !n.leaf() {
for j := i + 1; j < len((*n.children)); j++ {
(*n.children)[j] = nil
}
}
n.items = n.items[:i]
if !n.leaf() {
*n.children = (*n.children)[:i+1]
}
}
n.updateCount()
2022-06-16 19:36:02 +02:00
return right, median
}
2023-01-15 14:26:32 +01:00
func (n *mapNode[K, V]) updateCount() {
2022-06-16 19:36:02 +02:00
n.count = len(n.items)
if !n.leaf() {
for i := 0; i < len(*n.children); i++ {
n.count += (*n.children)[i].count
}
}
}
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) nodeSet(pn **mapNode[K, V], item mapPair[K, V],
) (prev V, replaced bool, split bool) {
n := tr.cowLoad(pn)
i, found := tr.bsearch(n, item.key)
2022-06-16 19:36:02 +02:00
if found {
2023-01-15 14:26:32 +01:00
prev = n.items[i].value
n.items[i].value = item.value
2022-06-16 19:36:02 +02:00
return prev, true, false
}
if n.leaf() {
if len(n.items) == maxItems {
2023-01-15 14:26:32 +01:00
return tr.empty.value, false, true
2022-06-16 19:36:02 +02:00
}
n.items = append(n.items, tr.empty)
copy(n.items[i+1:], n.items[i:])
n.items[i] = item
n.count++
2023-01-15 14:26:32 +01:00
return tr.empty.value, false, false
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
prev, replaced, split = tr.nodeSet(&(*n.children)[i], item)
2022-06-16 19:36:02 +02:00
if split {
if len(n.items) == maxItems {
2023-01-15 14:26:32 +01:00
return tr.empty.value, false, true
2022-06-16 19:36:02 +02:00
}
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
2023-01-15 14:26:32 +01:00
return tr.nodeSet(&n, item)
2022-06-16 19:36:02 +02:00
}
if !replaced {
n.count++
}
return prev, replaced, false
}
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) Scan(iter func(key K, value V) bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil {
return
}
tr.root.scan(iter)
}
2023-01-15 14:26:32 +01:00
func (n *mapNode[K, V]) scan(iter func(key K, value V) bool) bool {
2022-06-16 19:36:02 +02:00
if n.leaf() {
for i := 0; i < len(n.items); i++ {
2023-01-15 14:26:32 +01:00
if !iter(n.items[i].key, n.items[i].value) {
2022-06-16 19:36:02 +02:00
return false
}
}
return true
}
for i := 0; i < len(n.items); i++ {
if !(*n.children)[i].scan(iter) {
return false
}
2023-01-15 14:26:32 +01:00
if !iter(n.items[i].key, n.items[i].value) {
2022-06-16 19:36:02 +02:00
return false
}
}
return (*n.children)[len(*n.children)-1].scan(iter)
}
// Get a value for key
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) Get(key K) (V, bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil {
2023-01-15 14:26:32 +01:00
return tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
n := tr.root
for {
2023-01-15 14:26:32 +01:00
i, found := tr.bsearch(n, key)
2022-06-16 19:36:02 +02:00
if found {
2023-01-15 14:26:32 +01:00
return n.items[i].value, true
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
if n.leaf() {
return tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
n = (*n.children)[i]
}
}
// Len returns the number of items in the tree
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) Len() int {
2022-06-16 19:36:02 +02:00
return tr.count
}
2023-01-15 14:26:32 +01:00
// Delete a value for a key and returns the deleted value.
// Returns false if there was no value by that key found.
func (tr *Map[K, V]) Delete(key K) (V, bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil {
2023-01-15 14:26:32 +01:00
return tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
prev, deleted := tr.delete(&tr.root, false, key)
2022-06-16 19:36:02 +02:00
if !deleted {
2023-01-15 14:26:32 +01:00
return tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
if len(tr.root.items) == 0 && !tr.root.leaf() {
tr.root = (*tr.root.children)[0]
}
tr.count--
if tr.count == 0 {
tr.root = nil
}
2023-01-15 14:26:32 +01:00
return prev.value, true
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) delete(pn **mapNode[K, V], max bool, key K,
) (mapPair[K, V], bool) {
n := tr.cowLoad(pn)
2022-06-16 19:36:02 +02:00
var i int
var found bool
if max {
i, found = len(n.items)-1, true
} else {
2023-01-15 14:26:32 +01:00
i, found = tr.bsearch(n, key)
2022-06-16 19:36:02 +02:00
}
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
}
2023-01-15 14:26:32 +01:00
var prev mapPair[K, V]
2022-06-16 19:36:02 +02:00
var deleted bool
if found {
if max {
i++
2023-01-15 14:26:32 +01:00
prev, deleted = tr.delete(&(*n.children)[i], true, tr.empty.key)
2022-06-16 19:36:02 +02:00
} else {
prev = n.items[i]
2023-01-15 14:26:32 +01:00
maxItem, _ := tr.delete(&(*n.children)[i], true, tr.empty.key)
2022-06-16 19:36:02 +02:00
deleted = true
n.items[i] = maxItem
}
} else {
2023-01-15 14:26:32 +01:00
prev, deleted = tr.delete(&(*n.children)[i], max, key)
2022-06-16 19:36:02 +02:00
}
if !deleted {
return tr.empty, false
}
n.count--
if len((*n.children)[i].items) < minItems {
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.
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) nodeRebalance(n *mapNode[K, V], i int) {
2022-06-16 19:36:02 +02:00
if i == len(n.items) {
i--
}
// ensure copy-on-write
left := tr.cowLoad(&(*n.children)[i])
right := tr.cowLoad(&(*n.children)[i+1])
if len(left.items)+len(right.items) < maxItems {
// 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
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) Ascend(pivot K, iter func(key K, value V) bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil {
return
}
2023-01-15 14:26:32 +01:00
tr.ascend(tr.root, pivot, iter)
2022-06-16 19:36:02 +02:00
}
// The return value of this function determines whether we should keep iterating
// upon this functions return.
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) ascend(n *mapNode[K, V], pivot K,
iter func(key K, value V) bool,
2022-06-16 19:36:02 +02:00
) bool {
2023-01-15 14:26:32 +01:00
i, found := tr.bsearch(n, pivot)
2022-06-16 19:36:02 +02:00
if !found {
if !n.leaf() {
2023-01-15 14:26:32 +01:00
if !tr.ascend((*n.children)[i], pivot, iter) {
2022-06-16 19:36:02 +02:00
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++ {
2023-01-15 14:26:32 +01:00
if !iter(n.items[i].key, n.items[i].value) {
2022-06-16 19:36:02 +02:00
return false
}
if !n.leaf() {
if !(*n.children)[i+1].scan(iter) {
return false
}
}
}
return true
}
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) Reverse(iter func(key K, value V) bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil {
return
}
tr.root.reverse(iter)
}
2023-01-15 14:26:32 +01:00
func (n *mapNode[K, V]) reverse(iter func(key K, value V) bool) bool {
2022-06-16 19:36:02 +02:00
if n.leaf() {
for i := len(n.items) - 1; i >= 0; i-- {
2023-01-15 14:26:32 +01:00
if !iter(n.items[i].key, n.items[i].value) {
2022-06-16 19:36:02 +02:00
return false
}
}
return true
}
if !(*n.children)[len(*n.children)-1].reverse(iter) {
return false
}
for i := len(n.items) - 1; i >= 0; i-- {
2023-01-15 14:26:32 +01:00
if !iter(n.items[i].key, n.items[i].value) {
2022-06-16 19:36:02 +02:00
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
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) Descend(pivot K, iter func(key K, value V) bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil {
return
}
2023-01-15 14:26:32 +01:00
tr.descend(tr.root, pivot, iter)
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) descend(n *mapNode[K, V], pivot K,
iter func(key K, value V) bool,
2022-06-16 19:36:02 +02:00
) bool {
2023-01-15 14:26:32 +01:00
i, found := tr.bsearch(n, pivot)
2022-06-16 19:36:02 +02:00
if !found {
if !n.leaf() {
2023-01-15 14:26:32 +01:00
if !tr.descend((*n.children)[i], pivot, iter) {
2022-06-16 19:36:02 +02:00
return false
}
}
i--
}
for ; i >= 0; i-- {
2023-01-15 14:26:32 +01:00
if !iter(n.items[i].key, n.items[i].value) {
2022-06-16 19:36:02 +02:00
return false
}
if !n.leaf() {
if !(*n.children)[i].reverse(iter) {
return false
}
}
}
return true
}
// Load is for bulk loading pre-sorted items
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) Load(key K, value V) (V, bool) {
item := mapPair[K, V]{key: key, value: value}
2022-06-16 19:36:02 +02:00
if tr.root == nil {
2023-01-15 14:26:32 +01:00
return tr.Set(item.key, item.value)
2022-06-16 19:36:02 +02:00
}
n := tr.cowLoad(&tr.root)
for {
n.count++ // optimistically update counts
if n.leaf() {
if len(n.items) < maxItems {
2023-01-15 14:26:32 +01:00
if n.items[len(n.items)-1].key < item.key {
2022-06-16 19:36:02 +02:00
n.items = append(n.items, item)
tr.count++
2023-01-15 14:26:32 +01:00
return tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
}
break
}
n = tr.cowLoad(&(*n.children)[len(*n.children)-1])
}
// revert the counts
n = tr.root
for {
n.count--
if n.leaf() {
break
}
n = (*n.children)[len(*n.children)-1]
}
2023-01-15 14:26:32 +01:00
return tr.Set(item.key, item.value)
2022-06-16 19:36:02 +02:00
}
// Min returns the minimum item in tree.
2023-01-15 14:26:32 +01:00
// Returns nil if the treex has no items.
func (tr *Map[K, V]) Min() (K, V, bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil {
2023-01-15 14:26:32 +01:00
return tr.empty.key, tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
n := tr.root
for {
if n.leaf() {
2023-01-15 14:26:32 +01:00
item := n.items[0]
return item.key, item.value, true
2022-06-16 19:36:02 +02:00
}
n = (*n.children)[0]
}
}
// Max returns the maximum item in tree.
// Returns nil if the tree has no items.
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) Max() (K, V, bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil {
2023-01-15 14:26:32 +01:00
return tr.empty.key, tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
n := tr.root
for {
if n.leaf() {
2023-01-15 14:26:32 +01:00
item := n.items[len(n.items)-1]
return item.key, item.value, true
2022-06-16 19:36:02 +02:00
}
n = (*n.children)[len(*n.children)-1]
}
}
// PopMin removes the minimum item in tree and returns it.
// Returns nil if the tree has no items.
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) PopMin() (K, V, bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil {
2023-01-15 14:26:32 +01:00
return tr.empty.key, tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
n := tr.cowLoad(&tr.root)
2023-01-15 14:26:32 +01:00
var item mapPair[K, V]
2022-06-16 19:36:02 +02:00
for {
n.count-- // optimistically update counts
if n.leaf() {
item = n.items[0]
if len(n.items) == minItems {
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
}
2023-01-15 14:26:32 +01:00
return item.key, item.value, true
2022-06-16 19:36:02 +02:00
}
n = tr.cowLoad(&(*n.children)[0])
}
// revert the counts
n = tr.root
for {
n.count++
if n.leaf() {
break
}
n = (*n.children)[0]
}
2023-01-15 14:26:32 +01:00
value, deleted := tr.Delete(item.key)
if deleted {
return item.key, value, true
}
return tr.empty.key, tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
// PopMax removes the maximum item in tree and returns it.
2022-06-16 19:36:02 +02:00
// Returns nil if the tree has no items.
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) PopMax() (K, V, bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil {
2023-01-15 14:26:32 +01:00
return tr.empty.key, tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
n := tr.cowLoad(&tr.root)
2023-01-15 14:26:32 +01:00
var item mapPair[K, V]
2022-06-16 19:36:02 +02:00
for {
n.count-- // optimistically update counts
if n.leaf() {
item = n.items[len(n.items)-1]
if len(n.items) == minItems {
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
}
2023-01-15 14:26:32 +01:00
return item.key, item.value, true
2022-06-16 19:36:02 +02:00
}
n = tr.cowLoad(&(*n.children)[len(*n.children)-1])
}
// revert the counts
n = tr.root
for {
n.count++
if n.leaf() {
break
}
n = (*n.children)[len(*n.children)-1]
}
2023-01-15 14:26:32 +01:00
value, deleted := tr.Delete(item.key)
if deleted {
return item.key, value, true
}
return tr.empty.key, tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
// GetAt returns the value at index.
// Return nil if the tree is empty or the index is out of bounds.
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) GetAt(index int) (K, V, bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil || index < 0 || index >= tr.count {
2023-01-15 14:26:32 +01:00
return tr.empty.key, tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
n := tr.root
for {
if n.leaf() {
2023-01-15 14:26:32 +01:00
return n.items[index].key, n.items[index].value, true
2022-06-16 19:36:02 +02:00
}
i := 0
for ; i < len(n.items); i++ {
if index < (*n.children)[i].count {
break
} else if index == (*n.children)[i].count {
2023-01-15 14:26:32 +01:00
return n.items[i].key, n.items[i].value, true
2022-06-16 19:36:02 +02:00
}
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.
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) DeleteAt(index int) (K, V, bool) {
2022-06-16 19:36:02 +02:00
if tr.root == nil || index < 0 || index >= tr.count {
2023-01-15 14:26:32 +01:00
return tr.empty.key, tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
var pathbuf [8]uint8 // track the path
path := pathbuf[:0]
2023-01-15 14:26:32 +01:00
var item mapPair[K, V]
2022-06-16 19:36:02 +02:00
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 len(n.items) == minItems {
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
}
2023-01-15 14:26:32 +01:00
return item.key, item.value, true
2022-06-16 19:36:02 +02:00
}
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.cowLoad(&(*n.children)[i])
}
// revert the counts
n = tr.root
for i := 0; i < len(path); i++ {
n.count++
if !n.leaf() {
n = (*n.children)[uint8(path[i])]
}
}
2023-01-15 14:26:32 +01:00
value, deleted := tr.Delete(item.key)
if deleted {
return item.key, value, true
}
return tr.empty.key, tr.empty.value, false
2022-06-16 19:36:02 +02:00
}
// Height returns the height of the tree.
// Returns zero if tree has no items.
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) Height() int {
2022-06-16 19:36:02 +02:00
var height int
if tr.root != nil {
n := tr.root
for {
height++
if n.leaf() {
break
}
n = (*n.children)[0]
}
}
return height
}
2023-01-15 14:26:32 +01:00
// MapIter represents an iterator for btree.Map
type MapIter[K ordered, V any] struct {
tr *Map[K, V]
2022-06-16 19:36:02 +02:00
seeked bool
atstart bool
atend bool
2023-01-15 14:26:32 +01:00
stack []mapIterStackItem[K, V]
item mapPair[K, V]
2022-06-16 19:36:02 +02:00
}
2023-01-15 14:26:32 +01:00
type mapIterStackItem[K ordered, V any] struct {
n *mapNode[K, V]
2022-06-16 19:36:02 +02:00
i int
}
// Iter returns a read-only iterator.
2023-01-15 14:26:32 +01:00
func (tr *Map[K, V]) Iter() MapIter[K, V] {
var iter MapIter[K, V]
2022-06-16 19:36:02 +02:00
iter.tr = tr
return iter
}
// Seek to item greater-or-equal-to key.
// Returns false if there was no item found.
2023-01-15 14:26:32 +01:00
func (iter *MapIter[K, V]) Seek(key K) bool {
2022-06-16 19:36:02 +02:00
if iter.tr == nil {
return false
}
iter.seeked = true
iter.stack = iter.stack[:0]
if iter.tr.root == nil {
return false
}
n := iter.tr.root
for {
2023-01-15 14:26:32 +01:00
i, found := iter.tr.bsearch(n, key)
iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, i})
2022-06-16 19:36:02 +02:00
if found {
2023-01-15 14:26:32 +01:00
iter.item = n.items[i]
2022-06-16 19:36:02 +02:00
return true
}
if n.leaf() {
2023-01-15 14:26:32 +01:00
iter.stack[len(iter.stack)-1].i--
return iter.Next()
2022-06-16 19:36:02 +02:00
}
n = (*n.children)[i]
}
}
// First moves iterator to first item in tree.
// Returns false if the tree is empty.
2023-01-15 14:26:32 +01:00
func (iter *MapIter[K, V]) First() bool {
2022-06-16 19:36:02 +02:00
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.root
for {
2023-01-15 14:26:32 +01:00
iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, 0})
2022-06-16 19:36:02 +02:00
if n.leaf() {
break
}
n = (*n.children)[0]
}
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.
2023-01-15 14:26:32 +01:00
func (iter *MapIter[K, V]) Last() bool {
2022-06-16 19:36:02 +02:00
if iter.tr == nil {
return false
}
iter.seeked = true
iter.stack = iter.stack[:0]
if iter.tr.root == nil {
return false
}
n := iter.tr.root
for {
2023-01-15 14:26:32 +01:00
iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, len(n.items)})
2022-06-16 19:36:02 +02:00
if n.leaf() {
iter.stack[len(iter.stack)-1].i--
break
}
n = (*n.children)[len(n.items)]
}
s := &iter.stack[len(iter.stack)-1]
iter.item = s.n.items[s.i]
return true
}
// 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.
2023-01-15 14:26:32 +01:00
func (iter *MapIter[K, V]) Next() bool {
2022-06-16 19:36:02 +02:00
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 := (*s.n.children)[s.i]
for {
2023-01-15 14:26:32 +01:00
iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, 0})
2022-06-16 19:36:02 +02:00
if n.leaf() {
break
}
n = (*n.children)[0]
}
}
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.
2023-01-15 14:26:32 +01:00
func (iter *MapIter[K, V]) Prev() bool {
2022-06-16 19:36:02 +02:00
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 := (*s.n.children)[s.i]
for {
2023-01-15 14:26:32 +01:00
iter.stack = append(iter.stack,
mapIterStackItem[K, V]{n, len(n.items)})
2022-06-16 19:36:02 +02:00
if n.leaf() {
iter.stack[len(iter.stack)-1].i--
break
}
n = (*n.children)[len(n.items)]
}
}
s = &iter.stack[len(iter.stack)-1]
iter.item = s.n.items[s.i]
return true
}
2023-01-15 14:26:32 +01:00
// Key returns the current iterator item key.
func (iter *MapIter[K, V]) Key() K {
return iter.item.key
}
// Value returns the current iterator item value.
func (iter *MapIter[K, V]) Value() V {
return iter.item.value
}
// Values returns all the values in order.
func (tr *Map[K, V]) Values() []V {
values := make([]V, 0, tr.Len())
if tr.root != nil {
values = tr.root.values(values)
}
return values
}
func (n *mapNode[K, V]) values(values []V) []V {
if n.leaf() {
for i := 0; i < len(n.items); i++ {
values = append(values, n.items[i].value)
}
return values
}
for i := 0; i < len(n.items); i++ {
values = (*n.children)[i].values(values)
values = append(values, n.items[i].value)
}
return (*n.children)[len(*n.children)-1].values(values)
}
// Keys returns all the keys in order.
func (tr *Map[K, V]) Keys() []K {
keys := make([]K, 0, tr.Len())
if tr.root != nil {
keys = tr.root.keys(keys)
}
return keys
}
func (n *mapNode[K, V]) keys(keys []K) []K {
if n.leaf() {
for i := 0; i < len(n.items); i++ {
keys = append(keys, n.items[i].key)
}
return keys
}
for i := 0; i < len(n.items); i++ {
keys = (*n.children)[i].keys(keys)
keys = append(keys, n.items[i].key)
}
return (*n.children)[len(*n.children)-1].keys(keys)
}
// KeyValues returns all the keys and values in order.
func (tr *Map[K, V]) KeyValues() ([]K, []V) {
keys := make([]K, 0, tr.Len())
values := make([]V, 0, tr.Len())
if tr.root != nil {
keys, values = tr.root.keyValues(keys, values)
}
return keys, values
}
func (n *mapNode[K, V]) keyValues(keys []K, values []V) ([]K, []V) {
if n.leaf() {
for i := 0; i < len(n.items); i++ {
keys = append(keys, n.items[i].key)
values = append(values, n.items[i].value)
}
return keys, values
}
for i := 0; i < len(n.items); i++ {
keys, values = (*n.children)[i].keyValues(keys, values)
keys = append(keys, n.items[i].key)
values = append(values, n.items[i].value)
}
return (*n.children)[len(*n.children)-1].keyValues(keys, values)
2022-06-16 19:36:02 +02:00
}