jldb/mdb/index.go

package mdb

import (
	"unsafe"

	"github.com/google/btree"
)

func NewIndex[T any](
	c *Collection[T],
	name string,
	compare func(lhs, rhs *T) int,
) Index[T] {
	return c.addIndex(indexConfig[T]{
		Name:    name,
		Unique:  false,
		Compare: compare,
		Include: nil,
	})
}

func NewPartialIndex[T any](
	c *Collection[T],
	name string,
	compare func(lhs, rhs *T) int,
	include func(*T) bool,
) Index[T] {
	return c.addIndex(indexConfig[T]{
		Name:    name,
		Unique:  false,
		Compare: compare,
		Include: include,
	})
}

func NewUniqueIndex[T any](
	c *Collection[T],
	name string,
	compare func(lhs, rhs *T) int,
) Index[T] {
	return c.addIndex(indexConfig[T]{
		Name:    name,
		Unique:  true,
		Compare: compare,
		Include: nil,
	})
}

func NewUniquePartialIndex[T any](
	c *Collection[T],
	name string,
	compare func(lhs, rhs *T) int,
	include func(*T) bool,
) Index[T] {
	return c.addIndex(indexConfig[T]{
		Name:    name,
		Unique:  true,
		Compare: compare,
		Include: include,
	})
}

// ----------------------------------------------------------------------------

type Index[T any] struct {
	name         string
	collectionID uint64
	indexID      uint64
	include      func(*T) bool
	copy         func(*T) *T
}

func (i Index[T]) Get(tx *Snapshot, in *T) (item *T, ok bool) {
	tPtr, ok := i.get(tx, in)
	if !ok {
		return item, false
	}
	return i.copy(tPtr), true
}

func (i Index[T]) get(tx *Snapshot, in *T) (*T, bool) {
	return i.btree(tx).Get(in)
}

func (i Index[T]) Has(tx *Snapshot, in *T) bool {
	return i.btree(tx).Has(in)
}

func (i Index[T]) Min(tx *Snapshot) (item *T, ok bool) {
	tPtr, ok := i.btree(tx).Min()
	if !ok {
		return item, false
	}
	return i.copy(tPtr), true
}

func (i Index[T]) Max(tx *Snapshot) (item *T, ok bool) {
	tPtr, ok := i.btree(tx).Max()
	if !ok {
		return item, false
	}
	return i.copy(tPtr), true
}

func (i Index[T]) Ascend(tx *Snapshot, each func(*T) bool) {
	i.btreeForIter(tx).Ascend(func(t *T) bool {
		return each(i.copy(t))
	})
}

func (i Index[T]) AscendAfter(tx *Snapshot, after *T, each func(*T) bool) {
	i.btreeForIter(tx).AscendGreaterOrEqual(after, func(t *T) bool {
		return each(i.copy(t))
	})
}

func (i Index[T]) Descend(tx *Snapshot, each func(*T) bool) {
	i.btreeForIter(tx).Descend(func(t *T) bool {
		return each(i.copy(t))
	})
}

func (i Index[T]) DescendAfter(tx *Snapshot, after *T, each func(*T) bool) {
	i.btreeForIter(tx).DescendLessOrEqual(after, func(t *T) bool {
		return each(i.copy(t))
	})
}

type ListArgs[T any] struct {
	Desc  bool          // True for descending order, otherwise ascending.
	After *T            // If after is given, iterate after (and including) the value.
	While func(*T) bool // Continue iterating until While is false.
	Limit int           // Maximum number of items to return. 0 => All.
}

func (i Index[T]) List(tx *Snapshot, args ListArgs[T], out []*T) []*T {
	if args.Limit < 0 {
		return nil
	}

	if args.While == nil {
		args.While = func(*T) bool { return true }
	}

	size := args.Limit
	if size == 0 {
		size = 32 // Why not?
	}

	items := out[:0]

	each := func(item *T) bool {
		if !args.While(item) {
			return false
		}
		items = append(items, item)
		return args.Limit == 0 || len(items) < args.Limit
	}

	if args.Desc {
		if args.After != nil {
			i.DescendAfter(tx, args.After, each)
		} else {
			i.Descend(tx, each)
		}
	} else {
		if args.After != nil {
			i.AscendAfter(tx, args.After, each)
		} else {
			i.Ascend(tx, each)
		}
	}

	return items
}

// ----------------------------------------------------------------------------

func (i Index[T]) insertConflict(tx *Snapshot, item *T) bool {
	return i.btree(tx).Has(item)
}

func (i Index[T]) updateConflict(tx *Snapshot, item *T) bool {
	current, ok := i.btree(tx).Get(item)
	return ok && i.getID(current) != i.getID(item)
}

// This should only be called after insertConflict. Additionally, the caller
// should ensure that the index has been properly cloned for write before
// writing.
func (i Index[T]) insert(tx *Snapshot, item *T) {
	if i.include != nil && !i.include(item) {
		return
	}

	i.btree(tx).ReplaceOrInsert(item)
}

func (i Index[T]) update(tx *Snapshot, old, new *T) {
	bt := i.btree(tx)
	bt.Delete(old)

	// The insert call will also check the include function if available.
	i.insert(tx, new)
}

func (i Index[T]) delete(tx *Snapshot, item *T) {
	i.btree(tx).Delete(item)
}

// ----------------------------------------------------------------------------

func (i Index[T]) getState(tx *Snapshot) indexState[T] {
	return tx.collections[i.collectionID].(*collectionState[T]).Indices[i.indexID]
}

// Get the current btree for get/has/update/delete, etc.
func (i Index[T]) btree(tx *Snapshot) *btree.BTreeG[*T] {
	return i.getState(tx).BTree
}

func (i Index[T]) btreeForIter(tx *Snapshot) *btree.BTreeG[*T] {
	cState := tx.collections[i.collectionID].(*collectionState[T])
	bt := cState.Indices[i.indexID].BTree

	// If snapshot and index are writable, return a clone.
	if tx.writable() && cState.Version == tx.version {
		bt = bt.Clone()
	}

	return bt
}

func (i Index[T]) getID(t *T) uint64 {
	return *((*uint64)(unsafe.Pointer(t)))
}