165 lines
3.2 KiB
Go
165 lines
3.2 KiB
Go
package mdb
|
|
|
|
import (
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"github.com/google/btree"
|
|
)
|
|
|
|
type BTreeIndex[T any] struct {
|
|
c *Collection[T]
|
|
name string
|
|
modLock sync.Mutex
|
|
bt atomic.Value // *btree.BTreeG[*T]
|
|
getID func(*T) uint64
|
|
less func(*T, *T) bool
|
|
include func(*T) bool
|
|
}
|
|
|
|
func NewBTreeIndex[T any](
|
|
c *Collection[T],
|
|
name string,
|
|
less func(*T, *T) bool,
|
|
include func(*T) bool,
|
|
) *BTreeIndex[T] {
|
|
|
|
t := &BTreeIndex[T]{
|
|
c: c,
|
|
name: name,
|
|
getID: c.getID,
|
|
less: less,
|
|
include: include,
|
|
}
|
|
|
|
btree := btree.NewG(32, less)
|
|
t.bt.Store(btree)
|
|
c.indices = append(c.indices, t)
|
|
|
|
return t
|
|
}
|
|
|
|
func (t *BTreeIndex[T]) load(m map[uint64]*T) error {
|
|
btree := btree.NewG(64, t.less)
|
|
t.bt.Store(btree)
|
|
for _, item := range m {
|
|
if t.include == nil || t.include(item) {
|
|
if x, _ := btree.ReplaceOrInsert(item); x != nil {
|
|
return ErrDuplicate
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
|
|
func (t *BTreeIndex[T]) Ascend() *BTreeIterator[T] {
|
|
iter := newBTreeIterator[T]()
|
|
go func() {
|
|
t.btree().Ascend(iter.each)
|
|
iter.done()
|
|
}()
|
|
return iter
|
|
}
|
|
|
|
func (t *BTreeIndex[T]) AscendAfter(pivot T) *BTreeIterator[T] {
|
|
iter := newBTreeIterator[T]()
|
|
go func() {
|
|
t.btree().AscendGreaterOrEqual(&pivot, iter.each)
|
|
iter.done()
|
|
}()
|
|
return iter
|
|
}
|
|
|
|
func (t *BTreeIndex[T]) Descend() *BTreeIterator[T] {
|
|
iter := newBTreeIterator[T]()
|
|
go func() {
|
|
t.btree().Descend(iter.each)
|
|
iter.done()
|
|
}()
|
|
return iter
|
|
}
|
|
|
|
func (t *BTreeIndex[T]) DescendAfter(pivot T) *BTreeIterator[T] {
|
|
iter := newBTreeIterator[T]()
|
|
go func() {
|
|
t.btree().DescendLessOrEqual(&pivot, iter.each)
|
|
iter.done()
|
|
}()
|
|
return iter
|
|
}
|
|
|
|
func (t *BTreeIndex[T]) Get(item T) (T, bool) {
|
|
ptr, ok := t.btree().Get(&item)
|
|
if !ok {
|
|
return item, false
|
|
}
|
|
return *ptr, true
|
|
}
|
|
|
|
func (t *BTreeIndex[T]) Min() (item T, ok bool) {
|
|
if ptr, ok := t.btree().Min(); ok {
|
|
return *ptr, ok
|
|
}
|
|
return item, false
|
|
}
|
|
|
|
func (t *BTreeIndex[T]) Max() (item T, ok bool) {
|
|
if ptr, ok := t.btree().Max(); ok {
|
|
return *ptr, ok
|
|
}
|
|
return item, false
|
|
}
|
|
|
|
func (t *BTreeIndex[T]) Len() int {
|
|
return t.btree().Len()
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
|
|
func (t *BTreeIndex[T]) insert(item *T) {
|
|
if t.include == nil || t.include(item) {
|
|
t.modify(func(bt *btree.BTreeG[*T]) {
|
|
if _, ok := bt.ReplaceOrInsert(item); ok {
|
|
panic("Insert replaced an item.")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func (t *BTreeIndex[T]) update(old, new *T) {
|
|
t.modify(func(bt *btree.BTreeG[*T]) {
|
|
if t.include == nil || t.include(old) {
|
|
bt.Delete(old)
|
|
}
|
|
if t.include == nil || t.include(new) {
|
|
if _, ok := bt.ReplaceOrInsert(new); ok {
|
|
panic("Update insert replaced an item.")
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
func (t *BTreeIndex[T]) delete(item *T) {
|
|
if t.include == nil || t.include(item) {
|
|
t.modify(func(bt *btree.BTreeG[*T]) {
|
|
bt.Delete(item)
|
|
})
|
|
}
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
|
|
func (t *BTreeIndex[T]) btree() *btree.BTreeG[*T] {
|
|
return t.bt.Load().(*btree.BTreeG[*T])
|
|
}
|
|
|
|
func (t *BTreeIndex[T]) modify(mod func(clone *btree.BTreeG[*T])) {
|
|
t.modLock.Lock()
|
|
defer t.modLock.Unlock()
|
|
clone := t.btree().Clone()
|
|
mod(clone)
|
|
t.bt.Store(clone)
|
|
}
|