Updated Btree implementation

buntdb.go

@@ -1,7 +1,7 @@
 // Package buntdb implements a low-level in-memory key/value store in pure Go.
 // It persists to disk, is ACID compliant, and uses locking for multiple
-// readers and a single writer. Bunt is ideal for projects that need
+// readers and a single writer. Bunt is ideal for projects that need a
-// a dependable database, and favor speed over data size.
+// dependable database, and favor speed over data size.
 package buntdb
 
 import (
@@ -122,10 +122,11 @@ type Config struct {
 	// has been expired.
 	OnExpired func(keys []string)
 
-	// OnExpiredSync will be called inside the same transaction that is performing
+	// OnExpiredSync will be called inside the same transaction that is
-	// the deletion of expired items. If OnExpired is present then this callback
+	// performing the deletion of expired items. If OnExpired is present then
-	// will not be called. If this callback is present, then the deletion of the
+	// this callback will not be called. If this callback is present, then the
-	// timeed-out item is the explicit responsibility of this callback.
+	// deletion of the timeed-out item is the explicit responsibility of this
+	// callback.
 	OnExpiredSync func(key, value string, tx *Tx) error
 }
 
@@ -142,8 +143,8 @@ const btreeDegrees = 64
 func Open(path string) (*DB, error) {
 	db := &DB{}
 	// initialize trees and indexes
-	db.keys = btree.New(btreeDegrees, nil)
+	db.keys = btree.New(lessCtx(nil))
-	db.exps = btree.New(btreeDegrees, &exctx{db})
+	db.exps = btree.New(lessCtx(&exctx{db}))
 	db.idxs = make(map[string]*index)
 	// initialize default configuration
 	db.config = Config{
@@ -204,7 +205,7 @@ func (db *DB) Save(wr io.Writer) error {
 	// use a buffered writer and flush every 4MB
 	var buf []byte
 	// iterated through every item in the database and write to the buffer
-	db.keys.Ascend(func(item btree.Item) bool {
+	btreeAscend(db.keys, func(item interface{}) bool {
 		dbi := item.(*dbItem)
 		buf = dbi.writeSetTo(buf)
 		if len(buf) > 1024*1024*4 {
@@ -285,7 +286,7 @@ func (idx *index) clearCopy() *index {
 	}
 	// initialize with empty trees
 	if nidx.less != nil {
-		nidx.btr = btree.New(btreeDegrees, nidx)
+		nidx.btr = btree.New(lessCtx(nidx))
 	}
 	if nidx.rect != nil {
 		nidx.rtr = rtree.New(nidx)
@@ -297,20 +298,20 @@ func (idx *index) clearCopy() *index {
 func (idx *index) rebuild() {
 	// initialize trees
 	if idx.less != nil {
-		idx.btr = btree.New(btreeDegrees, idx)
+		idx.btr = btree.New(lessCtx(idx))
 	}
 	if idx.rect != nil {
 		idx.rtr = rtree.New(idx)
 	}
 	// iterate through all keys and fill the index
-	idx.db.keys.Ascend(func(item btree.Item) bool {
+	btreeAscend(idx.db.keys, func(item interface{}) bool {
 		dbi := item.(*dbItem)
 		if !idx.match(dbi.key) {
 			// does not match the pattern, continue
 			return true
 		}
 		if idx.less != nil {
-			idx.btr.ReplaceOrInsert(dbi)
+			idx.btr.Set(dbi)
 		}
 		if idx.rect != nil {
 			idx.rtr.Insert(dbi)
@@ -456,7 +457,7 @@ func (db *DB) SetConfig(config Config) error {
 // will be replaced with the new one, and return the previous item.
 func (db *DB) insertIntoDatabase(item *dbItem) *dbItem {
 	var pdbi *dbItem
-	prev := db.keys.ReplaceOrInsert(item)
+	prev := db.keys.Set(item)
 	if prev != nil {
 		// A previous item was removed from the keys tree. Let's
 		// fully delete this item from all indexes.
@@ -479,7 +480,7 @@ func (db *DB) insertIntoDatabase(item *dbItem) *dbItem {
 	if item.opts != nil && item.opts.ex {
 		// The new item has eviction options. Add it to the
 		// expires tree
-		db.exps.ReplaceOrInsert(item)
+		db.exps.Set(item)
 	}
 	for _, idx := range db.idxs {
 		if !idx.match(item.key) {
@@ -487,7 +488,7 @@ func (db *DB) insertIntoDatabase(item *dbItem) *dbItem {
 		}
 		if idx.btr != nil {
 			// Add new item to btree index.
-			idx.btr.ReplaceOrInsert(item)
+			idx.btr.Set(item)
 		}
 		if idx.rtr != nil {
 			// Add new item to rtree index.
@@ -557,9 +558,9 @@ func (db *DB) backgroundManager() {
 				}
 			}
 			// produce a list of expired items that need removing
-			db.exps.AscendLessThan(&dbItem{
+			btreeAscendLessThan(db.exps, &dbItem{
 				opts: &dbItemOpts{ex: true, exat: time.Now()},
-			}, func(item btree.Item) bool {
+			}, func(item interface{}) bool {
 				expired = append(expired, item.(*dbItem))
 				return true
 			})
@@ -674,8 +675,8 @@ func (db *DB) Shrink() error {
 			}
 			done = true
 			var n int
-			db.keys.AscendGreaterOrEqual(&dbItem{key: pivot},
+			btreeAscendGreaterOrEqual(db.keys, &dbItem{key: pivot},
-				func(item btree.Item) bool {
+				func(item interface{}) bool {
 					dbi := item.(*dbItem)
 					// 1000 items or 64MB buffer
 					if n > 1000 || len(buf) > 64*1024*1024 {
@@ -889,8 +890,8 @@ func (db *DB) readLoad(rd io.Reader, modTime time.Time) error {
 			db.deleteFromDatabase(&dbItem{key: parts[1]})
 		} else if (parts[0][0] == 'f' || parts[0][0] == 'F') &&
 			strings.ToLower(parts[0]) == "flushdb" {
-			db.keys = btree.New(btreeDegrees, nil)
+			db.keys = btree.New(lessCtx(nil))
-			db.exps = btree.New(btreeDegrees, &exctx{db})
+			db.exps = btree.New(lessCtx(&exctx{db}))
 			db.idxs = make(map[string]*index)
 		} else {
 			return ErrInvalid
@@ -1025,8 +1026,8 @@ func (tx *Tx) DeleteAll() error {
 	}
 
 	// now reset the live database trees
-	tx.db.keys = btree.New(btreeDegrees, nil)
+	tx.db.keys = btree.New(lessCtx(nil))
-	tx.db.exps = btree.New(btreeDegrees, &exctx{tx.db})
+	tx.db.exps = btree.New(lessCtx(&exctx{tx.db}))
 	tx.db.idxs = make(map[string]*index)
 
 	// finally re-create the indexes
@@ -1264,8 +1265,7 @@ func (dbi *dbItem) expiresAt() time.Time {
 // to note that the ctx parameter is used to help with determine which
 // formula to use on an item. Each b-tree should use a different ctx when
 // sharing the same item.
-func (dbi *dbItem) Less(item btree.Item, ctx interface{}) bool {
+func (dbi *dbItem) Less(dbi2 *dbItem, ctx interface{}) bool {
-	dbi2 := item.(*dbItem)
 	switch ctx := ctx.(type) {
 	case *exctx:
 		// The expires b-tree formula
@@ -1295,6 +1295,12 @@ func (dbi *dbItem) Less(item btree.Item, ctx interface{}) bool {
 	return dbi.key < dbi2.key
 }
 
+func lessCtx(ctx interface{}) func(a, b interface{}) bool {
+	return func(a, b interface{}) bool {
+		return a.(*dbItem).Less(b.(*dbItem), ctx)
+	}
+}
+
 // Rect converts a string to a rectangle.
 // An invalid rectangle will cause a panic.
 func (dbi *dbItem) Rect(ctx interface{}) (min, max []float64) {
@@ -1498,7 +1504,7 @@ func (tx *Tx) scan(desc, gt, lt bool, index, start, stop string,
 		return ErrTxClosed
 	}
 	// wrap a btree specific iterator around the user-defined iterator.
-	iter := func(item btree.Item) bool {
+	iter := func(item interface{}) bool {
 		dbi := item.(*dbItem)
 		return iterator(dbi.key, dbi.val)
 	}
@@ -1542,26 +1548,26 @@ func (tx *Tx) scan(desc, gt, lt bool, index, start, stop string,
 	if desc {
 		if gt {
 			if lt {
-				tr.DescendRange(itemA, itemB, iter)
+				btreeDescendRange(tr, itemA, itemB, iter)
 			} else {
-				tr.DescendGreaterThan(itemA, iter)
+				btreeDescendGreaterThan(tr, itemA, iter)
 			}
 		} else if lt {
-			tr.DescendLessOrEqual(itemA, iter)
+			btreeDescendLessOrEqual(tr, itemA, iter)
 		} else {
-			tr.Descend(iter)
+			btreeDescend(tr, iter)
 		}
 	} else {
 		if gt {
 			if lt {
-				tr.AscendRange(itemA, itemB, iter)
+				btreeAscendRange(tr, itemA, itemB, iter)
 			} else {
-				tr.AscendGreaterOrEqual(itemA, iter)
+				btreeAscendGreaterOrEqual(tr, itemA, iter)
 			}
 		} else if lt {
-			tr.AscendLessThan(itemA, iter)
+			btreeAscendLessThan(tr, itemA, iter)
 		} else {
-			tr.Ascend(iter)
+			btreeAscend(tr, iter)
 		}
 	}
 	return nil
@@ -2014,7 +2020,8 @@ func (tx *Tx) createIndex(name string, pattern string,
 	if tx.wc.rbkeys == nil {
 		// store the index in the rollback map.
 		if _, ok := tx.wc.rollbackIndexes[name]; !ok {
-			// we use nil to indicate that the index should be removed upon rollback.
+			// we use nil to indicate that the index should be removed upon
+			// rollback.
 			tx.wc.rollbackIndexes[name] = nil
 		}
 	}
@@ -2044,8 +2051,8 @@ func (tx *Tx) DropIndex(name string) error {
 	if tx.wc.rbkeys == nil {
 		// store the index in the rollback map.
 		if _, ok := tx.wc.rollbackIndexes[name]; !ok {
-			// we use a non-nil copy of the index without the data to indicate that the
+			// we use a non-nil copy of the index without the data to indicate
-			// index should be rebuilt upon rollback.
+			// that the index should be rebuilt upon rollback.
 			tx.wc.rollbackIndexes[name] = idx.clearCopy()
 		}
 	}
@@ -2181,3 +2188,67 @@ func IndexJSONCaseSensitive(path string) func(a, b string) bool {
 func Desc(less func(a, b string) bool) func(a, b string) bool {
 	return func(a, b string) bool { return less(b, a) }
 }
+
+//// Wrappers around btree Ascend/Descend
+
+func bLT(tr *btree.BTree, a, b interface{}) bool { return tr.Less(a, b) }
+func bGT(tr *btree.BTree, a, b interface{}) bool { return tr.Less(b, a) }
+
+// func bLTE(tr *btree.BTree, a, b interface{}) bool { return !tr.Less(b, a) }
+// func bGTE(tr *btree.BTree, a, b interface{}) bool { return !tr.Less(a, b) }
+
+// Ascend
+
+func btreeAscend(tr *btree.BTree, iter func(item interface{}) bool) {
+	tr.Ascend(nil, iter)
+}
+
+func btreeAscendLessThan(tr *btree.BTree, pivot interface{},
+	iter func(item interface{}) bool,
+) {
+	tr.Ascend(nil, func(item interface{}) bool {
+		return bLT(tr, item, pivot) && iter(item)
+	})
+}
+
+func btreeAscendGreaterOrEqual(tr *btree.BTree, pivot interface{},
+	iter func(item interface{}) bool,
+) {
+	tr.Ascend(pivot, iter)
+}
+
+func btreeAscendRange(tr *btree.BTree, greaterOrEqual, lessThan interface{},
+	iter func(item interface{}) bool,
+) {
+	tr.Ascend(greaterOrEqual, func(item interface{}) bool {
+		return bLT(tr, item, lessThan) && iter(item)
+	})
+}
+
+// Descend
+
+func btreeDescend(tr *btree.BTree, iter func(item interface{}) bool) {
+	tr.Descend(nil, iter)
+}
+
+func btreeDescendGreaterThan(tr *btree.BTree, pivot interface{},
+	iter func(item interface{}) bool,
+) {
+	tr.Descend(nil, func(item interface{}) bool {
+		return bGT(tr, item, pivot) && iter(item)
+	})
+}
+
+func btreeDescendRange(tr *btree.BTree, lessOrEqual, greaterThan interface{},
+	iter func(item interface{}) bool,
+) {
+	tr.Descend(lessOrEqual, func(item interface{}) bool {
+		return bGT(tr, item, greaterThan) && iter(item)
+	})
+}
+
+func btreeDescendLessOrEqual(tr *btree.BTree, pivot interface{},
+	iter func(item interface{}) bool,
+) {
+	tr.Descend(pivot, iter)
+}

go.mod

@@ -3,7 +3,7 @@ module github.com/tidwall/buntdb
 go 1.13
 
 require (
-	github.com/tidwall/btree v0.0.0-20191029221954-400434d76274
+	github.com/tidwall/btree v0.2.1
 	github.com/tidwall/gjson v1.3.4
 	github.com/tidwall/grect v0.0.0-20161006141115-ba9a043346eb
 	github.com/tidwall/match v1.0.1

go.sum

@@ -1,5 +1,11 @@
 github.com/tidwall/btree v0.0.0-20191029221954-400434d76274 h1:G6Z6HvJuPjG6XfNGi/feOATzeJrfgTNJY+rGrHbA04E=
 github.com/tidwall/btree v0.0.0-20191029221954-400434d76274/go.mod h1:huei1BkDWJ3/sLXmO+bsCNELL+Bp2Kks9OLyQFkzvA8=
+github.com/tidwall/btree v0.1.0 h1:08K82J5GUBN+zdEE8/xWaBgU6+yYWW995LhHFEl8dzI=
+github.com/tidwall/btree v0.1.0/go.mod h1:huei1BkDWJ3/sLXmO+bsCNELL+Bp2Kks9OLyQFkzvA8=
+github.com/tidwall/btree v0.2.0 h1:+x3HF7KQ4uarFfyQpAViE+R3ip/zIjASmAC54ocDNng=
+github.com/tidwall/btree v0.2.0/go.mod h1:huei1BkDWJ3/sLXmO+bsCNELL+Bp2Kks9OLyQFkzvA8=
+github.com/tidwall/btree v0.2.1 h1:jToKrG8o1uRsYRr07sTpLyiP4deacqJetUk/+WI7H94=
+github.com/tidwall/btree v0.2.1/go.mod h1:huei1BkDWJ3/sLXmO+bsCNELL+Bp2Kks9OLyQFkzvA8=
 github.com/tidwall/gjson v1.3.4 h1:On5waDnyKKk3SWE4EthbjjirAWXp43xx5cKCUZY1eZw=
 github.com/tidwall/gjson v1.3.4/go.mod h1:P256ACg0Mn+j1RXIDXoss50DeIABTYK1PULOJHhxOls=
 github.com/tidwall/grect v0.0.0-20161006141115-ba9a043346eb h1:5NSYaAdrnblKByzd7XByQEJVT8+9v0W/tIY0Oo4OwrE=