mirror of
https://github.com/ergochat/ergo.git
synced 2024-12-30 14:42:50 +01:00
979 lines
28 KiB
Go
979 lines
28 KiB
Go
// Copyright 2014 Google Inc.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
// Package btree implements in-memory B-Trees of arbitrary degree.
|
|
//
|
|
// btree implements an in-memory B-Tree for use as an ordered data structure.
|
|
// It is not meant for persistent storage solutions.
|
|
//
|
|
// It has a flatter structure than an equivalent red-black or other binary tree,
|
|
// which in some cases yields better memory usage and/or performance.
|
|
// See some discussion on the matter here:
|
|
// http://google-opensource.blogspot.com/2013/01/c-containers-that-save-memory-and-time.html
|
|
// Note, though, that this project is in no way related to the C++ B-Tree
|
|
// implementation written about there.
|
|
//
|
|
// Within this tree, each node contains a slice of items and a (possibly nil)
|
|
// slice of children. For basic numeric values or raw structs, this can cause
|
|
// efficiency differences when compared to equivalent C++ template code that
|
|
// stores values in arrays within the node:
|
|
// * Due to the overhead of storing values as interfaces (each
|
|
// value needs to be stored as the value itself, then 2 words for the
|
|
// interface pointing to that value and its type), resulting in higher
|
|
// memory use.
|
|
// * Since interfaces can point to values anywhere in memory, values are
|
|
// most likely not stored in contiguous blocks, resulting in a higher
|
|
// number of cache misses.
|
|
// These issues don't tend to matter, though, when working with strings or other
|
|
// heap-allocated structures, since C++-equivalent structures also must store
|
|
// pointers and also distribute their values across the heap.
|
|
//
|
|
// This implementation is designed to be a drop-in replacement to gollrb.LLRB
|
|
// trees, (http://github.com/petar/gollrb), an excellent and probably the most
|
|
// widely used ordered tree implementation in the Go ecosystem currently.
|
|
// Its functions, therefore, exactly mirror those of
|
|
// llrb.LLRB where possible. Unlike gollrb, though, we currently don't
|
|
// support storing multiple equivalent values.
|
|
package btree
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"strings"
|
|
"sync"
|
|
)
|
|
|
|
// Item represents a single object in the tree.
|
|
type Item interface {
|
|
// Less tests whether the current item is less than the given argument.
|
|
//
|
|
// This must provide a strict weak ordering.
|
|
// If !a.Less(b) && !b.Less(a), we treat this to mean a == b (i.e. we can only
|
|
// hold one of either a or b in the tree).
|
|
//
|
|
// There is a user-defined ctx argument that is equal to the ctx value which
|
|
// is set at time of the btree contruction.
|
|
Less(than Item, ctx interface{}) bool
|
|
}
|
|
|
|
const (
|
|
DefaultFreeListSize = 32
|
|
)
|
|
|
|
var (
|
|
nilItems = make(items, 16)
|
|
nilChildren = make(children, 16)
|
|
)
|
|
|
|
// FreeList represents a free list of btree nodes. By default each
|
|
// BTree has its own FreeList, but multiple BTrees can share the same
|
|
// FreeList.
|
|
// Two Btrees using the same freelist are safe for concurrent write access.
|
|
type FreeList struct {
|
|
mu sync.Mutex
|
|
freelist []*node
|
|
}
|
|
|
|
// NewFreeList creates a new free list.
|
|
// size is the maximum size of the returned free list.
|
|
func NewFreeList(size int) *FreeList {
|
|
return &FreeList{freelist: make([]*node, 0, size)}
|
|
}
|
|
|
|
func (f *FreeList) newNode() (n *node) {
|
|
f.mu.Lock()
|
|
index := len(f.freelist) - 1
|
|
if index < 0 {
|
|
f.mu.Unlock()
|
|
return new(node)
|
|
}
|
|
n = f.freelist[index]
|
|
f.freelist[index] = nil
|
|
f.freelist = f.freelist[:index]
|
|
f.mu.Unlock()
|
|
return
|
|
}
|
|
|
|
func (f *FreeList) freeNode(n *node) {
|
|
f.mu.Lock()
|
|
if len(f.freelist) < cap(f.freelist) {
|
|
f.freelist = append(f.freelist, n)
|
|
}
|
|
f.mu.Unlock()
|
|
}
|
|
|
|
// ItemIterator allows callers of Ascend* to iterate in-order over portions of
|
|
// the tree. When this function returns false, iteration will stop and the
|
|
// associated Ascend* function will immediately return.
|
|
type ItemIterator func(i Item) bool
|
|
|
|
// New creates a new B-Tree with the given degree.
|
|
//
|
|
// New(2), for example, will create a 2-3-4 tree (each node contains 1-3 items
|
|
// and 2-4 children).
|
|
func New(degree int, ctx interface{}) *BTree {
|
|
return NewWithFreeList(degree, NewFreeList(DefaultFreeListSize), ctx)
|
|
}
|
|
|
|
// NewWithFreeList creates a new B-Tree that uses the given node free list.
|
|
func NewWithFreeList(degree int, f *FreeList, ctx interface{}) *BTree {
|
|
if degree <= 1 {
|
|
panic("bad degree")
|
|
}
|
|
return &BTree{
|
|
degree: degree,
|
|
cow: ©OnWriteContext{freelist: f},
|
|
ctx: ctx,
|
|
}
|
|
}
|
|
|
|
// items stores items in a node.
|
|
type items []Item
|
|
|
|
// insertAt inserts a value into the given index, pushing all subsequent values
|
|
// forward.
|
|
func (s *items) insertAt(index int, item Item) {
|
|
*s = append(*s, nil)
|
|
if index < len(*s) {
|
|
copy((*s)[index+1:], (*s)[index:])
|
|
}
|
|
(*s)[index] = item
|
|
}
|
|
|
|
// removeAt removes a value at a given index, pulling all subsequent values
|
|
// back.
|
|
func (s *items) removeAt(index int) Item {
|
|
item := (*s)[index]
|
|
copy((*s)[index:], (*s)[index+1:])
|
|
(*s)[len(*s)-1] = nil
|
|
*s = (*s)[:len(*s)-1]
|
|
return item
|
|
}
|
|
|
|
// pop removes and returns the last element in the list.
|
|
func (s *items) pop() (out Item) {
|
|
index := len(*s) - 1
|
|
out = (*s)[index]
|
|
(*s)[index] = nil
|
|
*s = (*s)[:index]
|
|
return
|
|
}
|
|
|
|
// truncate truncates this instance at index so that it contains only the
|
|
// first index items. index must be less than or equal to length.
|
|
func (s *items) truncate(index int) {
|
|
var toClear items
|
|
*s, toClear = (*s)[:index], (*s)[index:]
|
|
for len(toClear) > 0 {
|
|
toClear = toClear[copy(toClear, nilItems):]
|
|
}
|
|
}
|
|
|
|
// find returns the index where the given item should be inserted into this
|
|
// list. 'found' is true if the item already exists in the list at the given
|
|
// index.
|
|
func (s items) find(item Item, ctx interface{}) (index int, found bool) {
|
|
i, j := 0, len(s)
|
|
for i < j {
|
|
h := i + (j-i)/2
|
|
if !item.Less(s[h], ctx) {
|
|
i = h + 1
|
|
} else {
|
|
j = h
|
|
}
|
|
}
|
|
if i > 0 && !s[i-1].Less(item, ctx) {
|
|
return i - 1, true
|
|
}
|
|
return i, false
|
|
}
|
|
|
|
// children stores child nodes in a node.
|
|
type children []*node
|
|
|
|
// insertAt inserts a value into the given index, pushing all subsequent values
|
|
// forward.
|
|
func (s *children) insertAt(index int, n *node) {
|
|
*s = append(*s, nil)
|
|
if index < len(*s) {
|
|
copy((*s)[index+1:], (*s)[index:])
|
|
}
|
|
(*s)[index] = n
|
|
}
|
|
|
|
// removeAt removes a value at a given index, pulling all subsequent values
|
|
// back.
|
|
func (s *children) removeAt(index int) *node {
|
|
n := (*s)[index]
|
|
copy((*s)[index:], (*s)[index+1:])
|
|
(*s)[len(*s)-1] = nil
|
|
*s = (*s)[:len(*s)-1]
|
|
return n
|
|
}
|
|
|
|
// pop removes and returns the last element in the list.
|
|
func (s *children) pop() (out *node) {
|
|
index := len(*s) - 1
|
|
out = (*s)[index]
|
|
(*s)[index] = nil
|
|
*s = (*s)[:index]
|
|
return
|
|
}
|
|
|
|
// truncate truncates this instance at index so that it contains only the
|
|
// first index children. index must be less than or equal to length.
|
|
func (s *children) truncate(index int) {
|
|
var toClear children
|
|
*s, toClear = (*s)[:index], (*s)[index:]
|
|
for len(toClear) > 0 {
|
|
toClear = toClear[copy(toClear, nilChildren):]
|
|
}
|
|
}
|
|
|
|
// node is an internal node in a tree.
|
|
//
|
|
// It must at all times maintain the invariant that either
|
|
// * len(children) == 0, len(items) unconstrained
|
|
// * len(children) == len(items) + 1
|
|
type node struct {
|
|
items items
|
|
children children
|
|
cow *copyOnWriteContext
|
|
}
|
|
|
|
func (n *node) mutableFor(cow *copyOnWriteContext) *node {
|
|
if n.cow == cow {
|
|
return n
|
|
}
|
|
out := cow.newNode()
|
|
if cap(out.items) >= len(n.items) {
|
|
out.items = out.items[:len(n.items)]
|
|
} else {
|
|
out.items = make(items, len(n.items), cap(n.items))
|
|
}
|
|
copy(out.items, n.items)
|
|
// Copy children
|
|
if cap(out.children) >= len(n.children) {
|
|
out.children = out.children[:len(n.children)]
|
|
} else {
|
|
out.children = make(children, len(n.children), cap(n.children))
|
|
}
|
|
copy(out.children, n.children)
|
|
return out
|
|
}
|
|
|
|
func (n *node) mutableChild(i int) *node {
|
|
c := n.children[i].mutableFor(n.cow)
|
|
n.children[i] = c
|
|
return c
|
|
}
|
|
|
|
// split splits the given node at the given index. The current node shrinks,
|
|
// and this function returns the item that existed at that index and a new node
|
|
// containing all items/children after it.
|
|
func (n *node) split(i int) (Item, *node) {
|
|
item := n.items[i]
|
|
next := n.cow.newNode()
|
|
next.items = append(next.items, n.items[i+1:]...)
|
|
n.items.truncate(i)
|
|
if len(n.children) > 0 {
|
|
next.children = append(next.children, n.children[i+1:]...)
|
|
n.children.truncate(i + 1)
|
|
}
|
|
return item, next
|
|
}
|
|
|
|
// maybeSplitChild checks if a child should be split, and if so splits it.
|
|
// Returns whether or not a split occurred.
|
|
func (n *node) maybeSplitChild(i, maxItems int) bool {
|
|
if len(n.children[i].items) < maxItems {
|
|
return false
|
|
}
|
|
first := n.mutableChild(i)
|
|
item, second := first.split(maxItems / 2)
|
|
n.items.insertAt(i, item)
|
|
n.children.insertAt(i+1, second)
|
|
return true
|
|
}
|
|
|
|
// insert inserts an item into the subtree rooted at this node, making sure
|
|
// no nodes in the subtree exceed maxItems items. Should an equivalent item be
|
|
// be found/replaced by insert, it will be returned.
|
|
func (n *node) insert(item Item, maxItems int, ctx interface{}) Item {
|
|
i, found := n.items.find(item, ctx)
|
|
if found {
|
|
out := n.items[i]
|
|
n.items[i] = item
|
|
return out
|
|
}
|
|
if len(n.children) == 0 {
|
|
n.items.insertAt(i, item)
|
|
return nil
|
|
}
|
|
if n.maybeSplitChild(i, maxItems) {
|
|
inTree := n.items[i]
|
|
switch {
|
|
case item.Less(inTree, ctx):
|
|
// no change, we want first split node
|
|
case inTree.Less(item, ctx):
|
|
i++ // we want second split node
|
|
default:
|
|
out := n.items[i]
|
|
n.items[i] = item
|
|
return out
|
|
}
|
|
}
|
|
return n.mutableChild(i).insert(item, maxItems, ctx)
|
|
}
|
|
|
|
// get finds the given key in the subtree and returns it.
|
|
func (n *node) get(key Item, ctx interface{}) Item {
|
|
i, found := n.items.find(key, ctx)
|
|
if found {
|
|
return n.items[i]
|
|
} else if len(n.children) > 0 {
|
|
return n.children[i].get(key, ctx)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// min returns the first item in the subtree.
|
|
func min(n *node) Item {
|
|
if n == nil {
|
|
return nil
|
|
}
|
|
for len(n.children) > 0 {
|
|
n = n.children[0]
|
|
}
|
|
if len(n.items) == 0 {
|
|
return nil
|
|
}
|
|
return n.items[0]
|
|
}
|
|
|
|
// max returns the last item in the subtree.
|
|
func max(n *node) Item {
|
|
if n == nil {
|
|
return nil
|
|
}
|
|
for len(n.children) > 0 {
|
|
n = n.children[len(n.children)-1]
|
|
}
|
|
if len(n.items) == 0 {
|
|
return nil
|
|
}
|
|
return n.items[len(n.items)-1]
|
|
}
|
|
|
|
// toRemove details what item to remove in a node.remove call.
|
|
type toRemove int
|
|
|
|
const (
|
|
removeItem toRemove = iota // removes the given item
|
|
removeMin // removes smallest item in the subtree
|
|
removeMax // removes largest item in the subtree
|
|
)
|
|
|
|
// remove removes an item from the subtree rooted at this node.
|
|
func (n *node) remove(item Item, minItems int, typ toRemove, ctx interface{}) Item {
|
|
var i int
|
|
var found bool
|
|
switch typ {
|
|
case removeMax:
|
|
if len(n.children) == 0 {
|
|
return n.items.pop()
|
|
}
|
|
i = len(n.items)
|
|
case removeMin:
|
|
if len(n.children) == 0 {
|
|
return n.items.removeAt(0)
|
|
}
|
|
i = 0
|
|
case removeItem:
|
|
i, found = n.items.find(item, ctx)
|
|
if len(n.children) == 0 {
|
|
if found {
|
|
return n.items.removeAt(i)
|
|
}
|
|
return nil
|
|
}
|
|
default:
|
|
panic("invalid type")
|
|
}
|
|
// If we get to here, we have children.
|
|
if len(n.children[i].items) <= minItems {
|
|
return n.growChildAndRemove(i, item, minItems, typ, ctx)
|
|
}
|
|
child := n.mutableChild(i)
|
|
// Either we had enough items to begin with, or we've done some
|
|
// merging/stealing, because we've got enough now and we're ready to return
|
|
// stuff.
|
|
if found {
|
|
// The item exists at index 'i', and the child we've selected can give us a
|
|
// predecessor, since if we've gotten here it's got > minItems items in it.
|
|
out := n.items[i]
|
|
// We use our special-case 'remove' call with typ=maxItem to pull the
|
|
// predecessor of item i (the rightmost leaf of our immediate left child)
|
|
// and set it into where we pulled the item from.
|
|
n.items[i] = child.remove(nil, minItems, removeMax, ctx)
|
|
return out
|
|
}
|
|
// Final recursive call. Once we're here, we know that the item isn't in this
|
|
// node and that the child is big enough to remove from.
|
|
return child.remove(item, minItems, typ, ctx)
|
|
}
|
|
|
|
// growChildAndRemove grows child 'i' to make sure it's possible to remove an
|
|
// item from it while keeping it at minItems, then calls remove to actually
|
|
// remove it.
|
|
//
|
|
// Most documentation says we have to do two sets of special casing:
|
|
// 1) item is in this node
|
|
// 2) item is in child
|
|
// In both cases, we need to handle the two subcases:
|
|
// A) node has enough values that it can spare one
|
|
// B) node doesn't have enough values
|
|
// For the latter, we have to check:
|
|
// a) left sibling has node to spare
|
|
// b) right sibling has node to spare
|
|
// c) we must merge
|
|
// To simplify our code here, we handle cases #1 and #2 the same:
|
|
// If a node doesn't have enough items, we make sure it does (using a,b,c).
|
|
// We then simply redo our remove call, and the second time (regardless of
|
|
// whether we're in case 1 or 2), we'll have enough items and can guarantee
|
|
// that we hit case A.
|
|
func (n *node) growChildAndRemove(i int, item Item, minItems int, typ toRemove, ctx interface{}) Item {
|
|
if i > 0 && len(n.children[i-1].items) > minItems {
|
|
// Steal from left child
|
|
child := n.mutableChild(i)
|
|
stealFrom := n.mutableChild(i - 1)
|
|
stolenItem := stealFrom.items.pop()
|
|
child.items.insertAt(0, n.items[i-1])
|
|
n.items[i-1] = stolenItem
|
|
if len(stealFrom.children) > 0 {
|
|
child.children.insertAt(0, stealFrom.children.pop())
|
|
}
|
|
} else if i < len(n.items) && len(n.children[i+1].items) > minItems {
|
|
// steal from right child
|
|
child := n.mutableChild(i)
|
|
stealFrom := n.mutableChild(i + 1)
|
|
stolenItem := stealFrom.items.removeAt(0)
|
|
child.items = append(child.items, n.items[i])
|
|
n.items[i] = stolenItem
|
|
if len(stealFrom.children) > 0 {
|
|
child.children = append(child.children, stealFrom.children.removeAt(0))
|
|
}
|
|
} else {
|
|
if i >= len(n.items) {
|
|
i--
|
|
}
|
|
child := n.mutableChild(i)
|
|
// merge with right child
|
|
mergeItem := n.items.removeAt(i)
|
|
mergeChild := n.children.removeAt(i + 1)
|
|
child.items = append(child.items, mergeItem)
|
|
child.items = append(child.items, mergeChild.items...)
|
|
child.children = append(child.children, mergeChild.children...)
|
|
n.cow.freeNode(mergeChild)
|
|
}
|
|
return n.remove(item, minItems, typ, ctx)
|
|
}
|
|
|
|
type direction int
|
|
|
|
const (
|
|
descend = direction(-1)
|
|
ascend = direction(+1)
|
|
)
|
|
|
|
// iterate provides a simple method for iterating over elements in the tree.
|
|
//
|
|
// When ascending, the 'start' should be less than 'stop' and when descending,
|
|
// the 'start' should be greater than 'stop'. Setting 'includeStart' to true
|
|
// will force the iterator to include the first item when it equals 'start',
|
|
// thus creating a "greaterOrEqual" or "lessThanEqual" rather than just a
|
|
// "greaterThan" or "lessThan" queries.
|
|
func (n *node) iterate(dir direction, start, stop Item, includeStart bool, hit bool, iter ItemIterator, ctx interface{}) (bool, bool) {
|
|
var ok bool
|
|
switch dir {
|
|
case ascend:
|
|
for i := 0; i < len(n.items); i++ {
|
|
if start != nil && n.items[i].Less(start, ctx) {
|
|
continue
|
|
}
|
|
if len(n.children) > 0 {
|
|
if hit, ok = n.children[i].iterate(dir, start, stop, includeStart, hit, iter, ctx); !ok {
|
|
return hit, false
|
|
}
|
|
}
|
|
if !includeStart && !hit && start != nil && !start.Less(n.items[i], ctx) {
|
|
hit = true
|
|
continue
|
|
}
|
|
hit = true
|
|
if stop != nil && !n.items[i].Less(stop, ctx) {
|
|
return hit, false
|
|
}
|
|
if !iter(n.items[i]) {
|
|
return hit, false
|
|
}
|
|
}
|
|
if len(n.children) > 0 {
|
|
if hit, ok = n.children[len(n.children)-1].iterate(dir, start, stop, includeStart, hit, iter, ctx); !ok {
|
|
return hit, false
|
|
}
|
|
}
|
|
case descend:
|
|
for i := len(n.items) - 1; i >= 0; i-- {
|
|
if start != nil && !n.items[i].Less(start, ctx) {
|
|
if !includeStart || hit || start.Less(n.items[i], ctx) {
|
|
continue
|
|
}
|
|
}
|
|
if len(n.children) > 0 {
|
|
if hit, ok = n.children[i+1].iterate(dir, start, stop, includeStart, hit, iter, ctx); !ok {
|
|
return hit, false
|
|
}
|
|
}
|
|
if stop != nil && !stop.Less(n.items[i], ctx) {
|
|
return hit, false // continue
|
|
}
|
|
hit = true
|
|
if !iter(n.items[i]) {
|
|
return hit, false
|
|
}
|
|
}
|
|
if len(n.children) > 0 {
|
|
if hit, ok = n.children[0].iterate(dir, start, stop, includeStart, hit, iter, ctx); !ok {
|
|
return hit, false
|
|
}
|
|
}
|
|
}
|
|
return hit, true
|
|
}
|
|
|
|
// Used for testing/debugging purposes.
|
|
func (n *node) print(w io.Writer, level int) {
|
|
fmt.Fprintf(w, "%sNODE:%v\n", strings.Repeat(" ", level), n.items)
|
|
for _, c := range n.children {
|
|
c.print(w, level+1)
|
|
}
|
|
}
|
|
|
|
// BTree is an implementation of a B-Tree.
|
|
//
|
|
// BTree stores Item instances in an ordered structure, allowing easy insertion,
|
|
// removal, and iteration.
|
|
//
|
|
// Write operations are not safe for concurrent mutation by multiple
|
|
// goroutines, but Read operations are.
|
|
type BTree struct {
|
|
degree int
|
|
length int
|
|
root *node
|
|
ctx interface{}
|
|
cow *copyOnWriteContext
|
|
}
|
|
|
|
// copyOnWriteContext pointers determine node ownership... a tree with a write
|
|
// context equivalent to a node's write context is allowed to modify that node.
|
|
// A tree whose write context does not match a node's is not allowed to modify
|
|
// it, and must create a new, writable copy (IE: it's a Clone).
|
|
//
|
|
// When doing any write operation, we maintain the invariant that the current
|
|
// node's context is equal to the context of the tree that requested the write.
|
|
// We do this by, before we descend into any node, creating a copy with the
|
|
// correct context if the contexts don't match.
|
|
//
|
|
// Since the node we're currently visiting on any write has the requesting
|
|
// tree's context, that node is modifiable in place. Children of that node may
|
|
// not share context, but before we descend into them, we'll make a mutable
|
|
// copy.
|
|
type copyOnWriteContext struct {
|
|
freelist *FreeList
|
|
}
|
|
|
|
// Clone clones the btree, lazily. Clone should not be called concurrently,
|
|
// but the original tree (t) and the new tree (t2) can be used concurrently
|
|
// once the Clone call completes.
|
|
//
|
|
// The internal tree structure of b is marked read-only and shared between t and
|
|
// t2. Writes to both t and t2 use copy-on-write logic, creating new nodes
|
|
// whenever one of b's original nodes would have been modified. Read operations
|
|
// should have no performance degredation. Write operations for both t and t2
|
|
// will initially experience minor slow-downs caused by additional allocs and
|
|
// copies due to the aforementioned copy-on-write logic, but should converge to
|
|
// the original performance characteristics of the original tree.
|
|
func (t *BTree) Clone() (t2 *BTree) {
|
|
// Create two entirely new copy-on-write contexts.
|
|
// This operation effectively creates three trees:
|
|
// the original, shared nodes (old b.cow)
|
|
// the new b.cow nodes
|
|
// the new out.cow nodes
|
|
cow1, cow2 := *t.cow, *t.cow
|
|
out := *t
|
|
t.cow = &cow1
|
|
out.cow = &cow2
|
|
return &out
|
|
}
|
|
|
|
// maxItems returns the max number of items to allow per node.
|
|
func (t *BTree) maxItems() int {
|
|
return t.degree*2 - 1
|
|
}
|
|
|
|
// minItems returns the min number of items to allow per node (ignored for the
|
|
// root node).
|
|
func (t *BTree) minItems() int {
|
|
return t.degree - 1
|
|
}
|
|
|
|
func (c *copyOnWriteContext) newNode() (n *node) {
|
|
n = c.freelist.newNode()
|
|
n.cow = c
|
|
return
|
|
}
|
|
|
|
func (c *copyOnWriteContext) freeNode(n *node) {
|
|
if n.cow == c {
|
|
// clear to allow GC
|
|
n.items.truncate(0)
|
|
n.children.truncate(0)
|
|
n.cow = nil
|
|
c.freelist.freeNode(n)
|
|
}
|
|
}
|
|
|
|
// ReplaceOrInsert adds the given item to the tree. If an item in the tree
|
|
// already equals the given one, it is removed from the tree and returned.
|
|
// Otherwise, nil is returned.
|
|
//
|
|
// nil cannot be added to the tree (will panic).
|
|
func (t *BTree) ReplaceOrInsert(item Item) Item {
|
|
if item == nil {
|
|
panic("nil item being added to BTree")
|
|
}
|
|
if t.root == nil {
|
|
t.root = t.cow.newNode()
|
|
t.root.items = append(t.root.items, item)
|
|
t.length++
|
|
return nil
|
|
} else {
|
|
t.root = t.root.mutableFor(t.cow)
|
|
if len(t.root.items) >= t.maxItems() {
|
|
item2, second := t.root.split(t.maxItems() / 2)
|
|
oldroot := t.root
|
|
t.root = t.cow.newNode()
|
|
t.root.items = append(t.root.items, item2)
|
|
t.root.children = append(t.root.children, oldroot, second)
|
|
}
|
|
}
|
|
out := t.root.insert(item, t.maxItems(), t.ctx)
|
|
if out == nil {
|
|
t.length++
|
|
}
|
|
return out
|
|
}
|
|
|
|
// Delete removes an item equal to the passed in item from the tree, returning
|
|
// it. If no such item exists, returns nil.
|
|
func (t *BTree) Delete(item Item) Item {
|
|
return t.deleteItem(item, removeItem, t.ctx)
|
|
}
|
|
|
|
// DeleteMin removes the smallest item in the tree and returns it.
|
|
// If no such item exists, returns nil.
|
|
func (t *BTree) DeleteMin() Item {
|
|
return t.deleteItem(nil, removeMin, t.ctx)
|
|
}
|
|
|
|
// DeleteMax removes the largest item in the tree and returns it.
|
|
// If no such item exists, returns nil.
|
|
func (t *BTree) DeleteMax() Item {
|
|
return t.deleteItem(nil, removeMax, t.ctx)
|
|
}
|
|
|
|
func (t *BTree) deleteItem(item Item, typ toRemove, ctx interface{}) Item {
|
|
if t.root == nil || len(t.root.items) == 0 {
|
|
return nil
|
|
}
|
|
t.root = t.root.mutableFor(t.cow)
|
|
out := t.root.remove(item, t.minItems(), typ, ctx)
|
|
if len(t.root.items) == 0 && len(t.root.children) > 0 {
|
|
oldroot := t.root
|
|
t.root = t.root.children[0]
|
|
t.cow.freeNode(oldroot)
|
|
}
|
|
if out != nil {
|
|
t.length--
|
|
}
|
|
return out
|
|
}
|
|
|
|
// AscendRange calls the iterator for every value in the tree within the range
|
|
// [greaterOrEqual, lessThan), until iterator returns false.
|
|
func (t *BTree) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) {
|
|
if t.root == nil {
|
|
return
|
|
}
|
|
t.root.iterate(ascend, greaterOrEqual, lessThan, true, false, iterator, t.ctx)
|
|
}
|
|
|
|
// AscendLessThan calls the iterator for every value in the tree within the range
|
|
// [first, pivot), until iterator returns false.
|
|
func (t *BTree) AscendLessThan(pivot Item, iterator ItemIterator) {
|
|
if t.root == nil {
|
|
return
|
|
}
|
|
t.root.iterate(ascend, nil, pivot, false, false, iterator, t.ctx)
|
|
}
|
|
|
|
// AscendGreaterOrEqual calls the iterator for every value in the tree within
|
|
// the range [pivot, last], until iterator returns false.
|
|
func (t *BTree) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {
|
|
if t.root == nil {
|
|
return
|
|
}
|
|
t.root.iterate(ascend, pivot, nil, true, false, iterator, t.ctx)
|
|
}
|
|
|
|
// Ascend calls the iterator for every value in the tree within the range
|
|
// [first, last], until iterator returns false.
|
|
func (t *BTree) Ascend(iterator ItemIterator) {
|
|
if t.root == nil {
|
|
return
|
|
}
|
|
t.root.iterate(ascend, nil, nil, false, false, iterator, t.ctx)
|
|
}
|
|
|
|
// DescendRange calls the iterator for every value in the tree within the range
|
|
// [lessOrEqual, greaterThan), until iterator returns false.
|
|
func (t *BTree) DescendRange(lessOrEqual, greaterThan Item, iterator ItemIterator) {
|
|
if t.root == nil {
|
|
return
|
|
}
|
|
t.root.iterate(descend, lessOrEqual, greaterThan, true, false, iterator, t.ctx)
|
|
}
|
|
|
|
// DescendLessOrEqual calls the iterator for every value in the tree within the range
|
|
// [pivot, first], until iterator returns false.
|
|
func (t *BTree) DescendLessOrEqual(pivot Item, iterator ItemIterator) {
|
|
if t.root == nil {
|
|
return
|
|
}
|
|
t.root.iterate(descend, pivot, nil, true, false, iterator, t.ctx)
|
|
}
|
|
|
|
// DescendGreaterThan calls the iterator for every value in the tree within
|
|
// the range (pivot, last], until iterator returns false.
|
|
func (t *BTree) DescendGreaterThan(pivot Item, iterator ItemIterator) {
|
|
if t.root == nil {
|
|
return
|
|
}
|
|
t.root.iterate(descend, nil, pivot, false, false, iterator, t.ctx)
|
|
}
|
|
|
|
// Descend calls the iterator for every value in the tree within the range
|
|
// [last, first], until iterator returns false.
|
|
func (t *BTree) Descend(iterator ItemIterator) {
|
|
if t.root == nil {
|
|
return
|
|
}
|
|
t.root.iterate(descend, nil, nil, false, false, iterator, t.ctx)
|
|
}
|
|
|
|
// Get looks for the key item in the tree, returning it. It returns nil if
|
|
// unable to find that item.
|
|
func (t *BTree) Get(key Item) Item {
|
|
if t.root == nil {
|
|
return nil
|
|
}
|
|
return t.root.get(key, t.ctx)
|
|
}
|
|
|
|
// Min returns the smallest item in the tree, or nil if the tree is empty.
|
|
func (t *BTree) Min() Item {
|
|
return min(t.root)
|
|
}
|
|
|
|
// Max returns the largest item in the tree, or nil if the tree is empty.
|
|
func (t *BTree) Max() Item {
|
|
return max(t.root)
|
|
}
|
|
|
|
// Has returns true if the given key is in the tree.
|
|
func (t *BTree) Has(key Item) bool {
|
|
return t.Get(key) != nil
|
|
}
|
|
|
|
// Len returns the number of items currently in the tree.
|
|
func (t *BTree) Len() int {
|
|
return t.length
|
|
}
|
|
|
|
// Context returns the context of the tree.
|
|
func (t *BTree) Context() interface{} {
|
|
return t.ctx
|
|
}
|
|
|
|
// SetContext will replace the context of the tree.
|
|
func (t *BTree) SetContext(ctx interface{}) {
|
|
t.ctx = ctx
|
|
}
|
|
|
|
// Int implements the Item interface for integers.
|
|
type Int int
|
|
|
|
// Less returns true if int(a) < int(b).
|
|
func (a Int) Less(b Item, ctx interface{}) bool {
|
|
return a < b.(Int)
|
|
}
|
|
|
|
type stackItem struct {
|
|
n *node // current node
|
|
i int // index of the next child/item.
|
|
}
|
|
|
|
// Cursor represents an iterator that can traverse over all items in the tree
|
|
// in sorted order.
|
|
//
|
|
// Changing data while traversing a cursor may result in unexpected items to
|
|
// be returned. You must reposition your cursor after mutating data.
|
|
type Cursor struct {
|
|
t *BTree
|
|
stack []stackItem
|
|
}
|
|
|
|
// Cursor returns a new cursor used to traverse over items in the tree.
|
|
func (t *BTree) Cursor() *Cursor {
|
|
return &Cursor{t: t}
|
|
}
|
|
|
|
// First moves the cursor to the first item in the tree and returns that item.
|
|
func (c *Cursor) First() Item {
|
|
c.stack = c.stack[:0]
|
|
n := c.t.root
|
|
if n == nil {
|
|
return nil
|
|
}
|
|
c.stack = append(c.stack, stackItem{n: n})
|
|
for len(n.children) > 0 {
|
|
n = n.children[0]
|
|
c.stack = append(c.stack, stackItem{n: n})
|
|
}
|
|
if len(n.items) == 0 {
|
|
return nil
|
|
}
|
|
return n.items[0]
|
|
}
|
|
|
|
// Next moves the cursor to the next item and returns that item.
|
|
func (c *Cursor) Next() Item {
|
|
if len(c.stack) == 0 {
|
|
return nil
|
|
}
|
|
si := len(c.stack) - 1
|
|
c.stack[si].i++
|
|
n := c.stack[si].n
|
|
i := c.stack[si].i
|
|
if i == len(n.children)+len(n.items) {
|
|
c.stack = c.stack[:len(c.stack)-1]
|
|
return c.Next()
|
|
}
|
|
if len(n.children) == 0 {
|
|
if i >= len(n.items) {
|
|
c.stack = c.stack[:len(c.stack)-1]
|
|
return c.Next()
|
|
}
|
|
return n.items[i]
|
|
} else if i%2 == 1 {
|
|
return n.items[i/2]
|
|
}
|
|
c.stack = append(c.stack, stackItem{n: n.children[i/2], i: -1})
|
|
return c.Next()
|
|
|
|
}
|
|
|
|
// Last moves the cursor to the last item in the tree and returns that item.
|
|
func (c *Cursor) Last() Item {
|
|
c.stack = c.stack[:0]
|
|
n := c.t.root
|
|
if n == nil {
|
|
return nil
|
|
}
|
|
c.stack = append(c.stack, stackItem{n: n, i: len(n.children) + len(n.items) - 1})
|
|
for len(n.children) > 0 {
|
|
n = n.children[len(n.children)-1]
|
|
c.stack = append(c.stack, stackItem{n: n, i: len(n.children) + len(n.items) - 1})
|
|
}
|
|
if len(n.items) == 0 {
|
|
return nil
|
|
}
|
|
return n.items[len(n.items)-1]
|
|
}
|
|
|
|
// Prev moves the cursor to the previous item and returns that item.
|
|
func (c *Cursor) Prev() Item {
|
|
if len(c.stack) == 0 {
|
|
return nil
|
|
}
|
|
si := len(c.stack) - 1
|
|
c.stack[si].i--
|
|
n := c.stack[si].n
|
|
i := c.stack[si].i
|
|
if i == -1 {
|
|
c.stack = c.stack[:len(c.stack)-1]
|
|
return c.Prev()
|
|
}
|
|
if len(n.children) == 0 {
|
|
return n.items[i]
|
|
} else if i%2 == 1 {
|
|
return n.items[i/2]
|
|
}
|
|
child := n.children[i/2]
|
|
c.stack = append(c.stack, stackItem{n: child,
|
|
i: len(child.children) + len(child.items)})
|
|
return c.Prev()
|
|
}
|
|
|
|
// Seek moves the cursor to provided item and returns that item.
|
|
// If the item does not exist then the next item is returned.
|
|
func (c *Cursor) Seek(pivot Item) Item {
|
|
c.stack = c.stack[:0]
|
|
n := c.t.root
|
|
for n != nil {
|
|
i, found := n.items.find(pivot, c.t.ctx)
|
|
c.stack = append(c.stack, stackItem{n: n})
|
|
if found {
|
|
if len(n.children) == 0 {
|
|
c.stack[len(c.stack)-1].i = i
|
|
} else {
|
|
c.stack[len(c.stack)-1].i = i*2 + 1
|
|
}
|
|
return n.items[i]
|
|
}
|
|
if len(n.children) == 0 {
|
|
if i == len(n.items) {
|
|
c.stack[len(c.stack)-1].i = i + 1
|
|
return c.Next()
|
|
}
|
|
c.stack[len(c.stack)-1].i = i
|
|
return n.items[i]
|
|
}
|
|
c.stack[len(c.stack)-1].i = i * 2
|
|
n = n.children[i]
|
|
}
|
|
return nil
|
|
}
|