123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890 |
- package btree
- import (
- "fmt"
- "io"
- "sort"
- "strings"
- "sync"
- )
- type Item interface {
-
-
-
-
-
- Less(than Item) bool
- }
- const (
- DefaultFreeListSize = 32
- )
- var (
- nilItems = make(items, 16)
- nilChildren = make(children, 16)
- )
- type FreeList struct {
- mu sync.Mutex
- freelist []*node
- }
- 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) (out bool) {
- f.mu.Lock()
- if len(f.freelist) < cap(f.freelist) {
- f.freelist = append(f.freelist, n)
- out = true
- }
- f.mu.Unlock()
- return
- }
- type ItemIterator func(i Item) bool
- func New(degree int) *BTree {
- return NewWithFreeList(degree, NewFreeList(DefaultFreeListSize))
- }
- func NewWithFreeList(degree int, f *FreeList) *BTree {
- if degree <= 1 {
- panic("bad degree")
- }
- return &BTree{
- degree: degree,
- cow: ©OnWriteContext{freelist: f},
- }
- }
- type items []Item
- 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
- }
- 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
- }
- func (s *items) pop() (out Item) {
- index := len(*s) - 1
- out = (*s)[index]
- (*s)[index] = nil
- *s = (*s)[:index]
- return
- }
- func (s *items) truncate(index int) {
- var toClear items
- *s, toClear = (*s)[:index], (*s)[index:]
- for len(toClear) > 0 {
- toClear = toClear[copy(toClear, nilItems):]
- }
- }
- func (s items) find(item Item) (index int, found bool) {
- i := sort.Search(len(s), func(i int) bool {
- return item.Less(s[i])
- })
- if i > 0 && !s[i-1].Less(item) {
- return i - 1, true
- }
- return i, false
- }
- type children []*node
- 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
- }
- 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
- }
- func (s *children) pop() (out *node) {
- index := len(*s) - 1
- out = (*s)[index]
- (*s)[index] = nil
- *s = (*s)[:index]
- return
- }
- func (s *children) truncate(index int) {
- var toClear children
- *s, toClear = (*s)[:index], (*s)[index:]
- for len(toClear) > 0 {
- toClear = toClear[copy(toClear, nilChildren):]
- }
- }
- 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)
-
- 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
- }
- 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
- }
- 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
- }
- func (n *node) insert(item Item, maxItems int) Item {
- i, found := n.items.find(item)
- 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):
-
- case inTree.Less(item):
- i++
- default:
- out := n.items[i]
- n.items[i] = item
- return out
- }
- }
- return n.mutableChild(i).insert(item, maxItems)
- }
- func (n *node) get(key Item) Item {
- i, found := n.items.find(key)
- if found {
- return n.items[i]
- } else if len(n.children) > 0 {
- return n.children[i].get(key)
- }
- return nil
- }
- 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]
- }
- 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]
- }
- type toRemove int
- const (
- removeItem toRemove = iota
- removeMin
- removeMax
- )
- func (n *node) remove(item Item, minItems int, typ toRemove) 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)
- if len(n.children) == 0 {
- if found {
- return n.items.removeAt(i)
- }
- return nil
- }
- default:
- panic("invalid type")
- }
-
- if len(n.children[i].items) <= minItems {
- return n.growChildAndRemove(i, item, minItems, typ)
- }
- child := n.mutableChild(i)
-
-
-
- if found {
-
-
- out := n.items[i]
-
-
-
- n.items[i] = child.remove(nil, minItems, removeMax)
- return out
- }
-
-
- return child.remove(item, minItems, typ)
- }
- func (n *node) growChildAndRemove(i int, item Item, minItems int, typ toRemove) Item {
- if i > 0 && len(n.children[i-1].items) > minItems {
-
- 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 {
-
- 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)
-
- 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)
- }
- type direction int
- const (
- descend = direction(-1)
- ascend = direction(+1)
- )
- func (n *node) iterate(dir direction, start, stop Item, includeStart bool, hit bool, iter ItemIterator) (bool, bool) {
- var ok, found bool
- var index int
- switch dir {
- case ascend:
- if start != nil {
- index, _ = n.items.find(start)
- }
- for i := index; i < len(n.items); i++ {
- if len(n.children) > 0 {
- if hit, ok = n.children[i].iterate(dir, start, stop, includeStart, hit, iter); !ok {
- return hit, false
- }
- }
- if !includeStart && !hit && start != nil && !start.Less(n.items[i]) {
- hit = true
- continue
- }
- hit = true
- if stop != nil && !n.items[i].Less(stop) {
- 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); !ok {
- return hit, false
- }
- }
- case descend:
- if start != nil {
- index, found = n.items.find(start)
- if !found {
- index = index - 1
- }
- } else {
- index = len(n.items) - 1
- }
- for i := index; i >= 0; i-- {
- if start != nil && !n.items[i].Less(start) {
- if !includeStart || hit || start.Less(n.items[i]) {
- continue
- }
- }
- if len(n.children) > 0 {
- if hit, ok = n.children[i+1].iterate(dir, start, stop, includeStart, hit, iter); !ok {
- return hit, false
- }
- }
- if stop != nil && !stop.Less(n.items[i]) {
- return hit, false
- }
- 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); !ok {
- return hit, false
- }
- }
- }
- return hit, true
- }
- 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)
- }
- }
- type BTree struct {
- degree int
- length int
- root *node
- cow *copyOnWriteContext
- }
- type copyOnWriteContext struct {
- freelist *FreeList
- }
- func (t *BTree) Clone() (t2 *BTree) {
-
-
-
-
-
- cow1, cow2 := *t.cow, *t.cow
- out := *t
- t.cow = &cow1
- out.cow = &cow2
- return &out
- }
- func (t *BTree) maxItems() int {
- return t.degree*2 - 1
- }
- func (t *BTree) minItems() int {
- return t.degree - 1
- }
- func (c *copyOnWriteContext) newNode() (n *node) {
- n = c.freelist.newNode()
- n.cow = c
- return
- }
- type freeType int
- const (
- ftFreelistFull freeType = iota
- ftStored
- ftNotOwned
- )
- func (c *copyOnWriteContext) freeNode(n *node) freeType {
- if n.cow == c {
-
- n.items.truncate(0)
- n.children.truncate(0)
- n.cow = nil
- if c.freelist.freeNode(n) {
- return ftStored
- } else {
- return ftFreelistFull
- }
- } else {
- return ftNotOwned
- }
- }
- 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())
- if out == nil {
- t.length++
- }
- return out
- }
- func (t *BTree) Delete(item Item) Item {
- return t.deleteItem(item, removeItem)
- }
- func (t *BTree) DeleteMin() Item {
- return t.deleteItem(nil, removeMin)
- }
- func (t *BTree) DeleteMax() Item {
- return t.deleteItem(nil, removeMax)
- }
- func (t *BTree) deleteItem(item Item, typ toRemove) 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)
- 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
- }
- func (t *BTree) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) {
- if t.root == nil {
- return
- }
- t.root.iterate(ascend, greaterOrEqual, lessThan, true, false, iterator)
- }
- func (t *BTree) AscendLessThan(pivot Item, iterator ItemIterator) {
- if t.root == nil {
- return
- }
- t.root.iterate(ascend, nil, pivot, false, false, iterator)
- }
- func (t *BTree) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {
- if t.root == nil {
- return
- }
- t.root.iterate(ascend, pivot, nil, true, false, iterator)
- }
- func (t *BTree) Ascend(iterator ItemIterator) {
- if t.root == nil {
- return
- }
- t.root.iterate(ascend, nil, nil, false, false, iterator)
- }
- func (t *BTree) DescendRange(lessOrEqual, greaterThan Item, iterator ItemIterator) {
- if t.root == nil {
- return
- }
- t.root.iterate(descend, lessOrEqual, greaterThan, true, false, iterator)
- }
- func (t *BTree) DescendLessOrEqual(pivot Item, iterator ItemIterator) {
- if t.root == nil {
- return
- }
- t.root.iterate(descend, pivot, nil, true, false, iterator)
- }
- func (t *BTree) DescendGreaterThan(pivot Item, iterator ItemIterator) {
- if t.root == nil {
- return
- }
- t.root.iterate(descend, nil, pivot, false, false, iterator)
- }
- func (t *BTree) Descend(iterator ItemIterator) {
- if t.root == nil {
- return
- }
- t.root.iterate(descend, nil, nil, false, false, iterator)
- }
- func (t *BTree) Get(key Item) Item {
- if t.root == nil {
- return nil
- }
- return t.root.get(key)
- }
- func (t *BTree) Min() Item {
- return min(t.root)
- }
- func (t *BTree) Max() Item {
- return max(t.root)
- }
- func (t *BTree) Has(key Item) bool {
- return t.Get(key) != nil
- }
- func (t *BTree) Len() int {
- return t.length
- }
- func (t *BTree) Clear(addNodesToFreelist bool) {
- if t.root != nil && addNodesToFreelist {
- t.root.reset(t.cow)
- }
- t.root, t.length = nil, 0
- }
- func (n *node) reset(c *copyOnWriteContext) bool {
- for _, child := range n.children {
- if !child.reset(c) {
- return false
- }
- }
- return c.freeNode(n) != ftFreelistFull
- }
- type Int int
- func (a Int) Less(b Item) bool {
- return a < b.(Int)
- }
|