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 { db *Database name string collectionID uint64 indexID uint64 include func(*T) bool copy func(*T) *T } func (i *Index[T]) ensureSnapshot(tx *Snapshot) *Snapshot { if tx == nil { tx = i.db.Snapshot() } return tx } func (i *Index[T]) Get(tx *Snapshot, in *T) *T { tx = i.ensureSnapshot(tx) if tPtr, ok := i.get(tx, in); ok { return i.copy(tPtr) } return nil } 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 { tx = i.ensureSnapshot(tx) return i.btree(tx).Has(in) } func (i *Index[T]) Min(tx *Snapshot) *T { tx = i.ensureSnapshot(tx) if tPtr, ok := i.btree(tx).Min(); ok { return i.copy(tPtr) } return nil } func (i *Index[T]) Max(tx *Snapshot) *T { tx = i.ensureSnapshot(tx) if tPtr, ok := i.btree(tx).Max(); ok { return i.copy(tPtr) } return nil } func (i *Index[T]) Ascend(tx *Snapshot, each func(*T) bool) { tx = i.ensureSnapshot(tx) 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) { tx = i.ensureSnapshot(tx) 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) { tx = i.ensureSnapshot(tx) 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) { tx = i.ensureSnapshot(tx) 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 { tx = i.ensureSnapshot(tx) if args == nil { args = &ListArgs[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))) }