Initial commit

mdb/index.go

@@ -0,0 +1,236 @@
+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)))
+}