Added AscendHint, DescendHint, and SeekHint to BTree and BTreeG

These new functions allow for using a path hint with iterators, making it potentially faster when seeking to the first item in the iteration. Benchmarks: https://github.com/tidwall/btree-benchmark About path hints: https://github.com/tidwall/btree/blob/master/PATH_HINT.md Usage: ```go tr.AscendHint(key, iter, &hint) // iterate items that are >= key, ascending tr.DescendHint(key, iter, &hint) // iterate items that are <= key, descending iter.SeekHint(key, &hint) // seek to item that is >= key ```
dzfranklin · Sep 6, 2023 · 8d3586d · 8d3586d
1 parent 6b26091
commit 8d3586d
Show file tree

Hide file tree

Showing 4 changed files with 231 additions and 9 deletions.
diff --git a/README.md b/README.md
@@ -219,6  219,9 @@ Load(item)              // load presorted items into tree
 SetHint(item, *hint)    // insert or replace an existing item
 GetHint(item, *hint)    // get an existing item
 DeleteHint(item, *hint) // delete an item
 AscendHint(key, iter, *hint)
 DescendHint(key, iter, *hint)
 SeekHint(key, iter, *hint)
 
 // Copy-on-write
 Copy()                  // copy the btree
@@ -339,6  342,9 @@ Load(item)              // load presorted items into tree
 SetHint(item, *hint)    // insert or replace an existing item
 GetHint(item, *hint)    // get an existing item
 DeleteHint(item, *hint) // delete an item
 AscendHint(key, iter, *hint)
 DescendHint(key, iter, *hint)
 SeekHint(key, iter, *hint)
 
 // Copy-on-write
 Copy()                  // copy the btree

diff --git a/btree.go b/btree.go
@@ -142,6  142,26 @@ func (tr *BTree) AscendMut(pivot any, iter func(item any) bool) {
 	}
 }
 
 func (tr *BTree) AscendHint(pivot any, iter func(item any) bool,
 	hint *PathHint,
 ) {
 	if pivot == nil {
 		tr.base.Scan(iter)
 	} else {
 		tr.base.AscendHint(pivot, iter, hint)
 	}
 }
 
 func (tr *BTree) AscendHintMut(pivot any, iter func(item any) bool,
 	hint *PathHint,
 ) {
 	if pivot == nil {
 		tr.base.ScanMut(iter)
 	} else {
 		tr.base.AscendHintMut(pivot, iter, hint)
 	}
 }
 
 // Descend the tree within the range [pivot, first]
 // Pass nil for pivot to scan all item in descending order
 // Return false to stop iterating
@@ -161,6  181,26 @@ func (tr *BTree) DescendMut(pivot any, iter func(item any) bool) {
 	}
 }
 
 func (tr *BTree) DescendHint(pivot any, iter func(item any) bool,
 	hint *PathHint,
 ) {
 	if pivot == nil {
 		tr.base.Reverse(iter)
 	} else {
 		tr.base.DescendHint(pivot, iter, hint)
 	}
 }
 
 func (tr *BTree) DescendHintMut(pivot any, iter func(item any) bool,
 	hint *PathHint,
 ) {
 	if pivot == nil {
 		tr.base.ReverseMut(iter)
 	} else {
 		tr.base.DescendHintMut(pivot, iter, hint)
 	}
 }
 
 // Load is for bulk loading pre-sorted items
 // If the load replaces and existing item then the value for the replaced item
 // is returned.
@@ -317,6  357,10 @@ func (iter *Iter) Seek(key any) bool {
 	return iter.base.Seek(key)
 }
 
 func (iter *Iter) SeekHint(key any, hint *PathHint) bool {
 	return iter.base.SeekHint(key, hint)
 }
 
 // First moves iterator to first item in tree.
 // Returns false if the tree is empty.
 func (iter *Iter) First() bool {

diff --git a/btreeg.go b/btreeg.go
@@ -601,19  601,30 @@ func (tr *BTreeG[T]) nodeRebalance(n *node[T], i int) {
 // Pass nil for pivot to scan all item in ascending order
 // Return false to stop iterating
 func (tr *BTreeG[T]) Ascend(pivot T, iter func(item T) bool) {
-	tr.ascend(pivot, iter, false)
 	tr.ascend(pivot, iter, false, nil)
 }
 func (tr *BTreeG[T]) AscendMut(pivot T, iter func(item T) bool) {
-	tr.ascend(pivot, iter, true)
 	tr.ascend(pivot, iter, true, nil)
 }
-func (tr *BTreeG[T]) ascend(pivot T, iter func(item T) bool, mut bool) {
 func (tr *BTreeG[T]) ascend(pivot T, iter func(item T) bool, mut bool,
 	hint *PathHint,
 ) {
 	if tr.lock(mut) {
 		defer tr.unlock(mut)
 	}
 	if tr.root == nil {
 		return
 	}
-	tr.nodeAscend(&tr.root, pivot, nil, 0, iter, mut)
 	tr.nodeAscend(&tr.root, pivot, hint, 0, iter, mut)
 }
 func (tr *BTreeG[T]) AscendHint(pivot T, iter func(item T) bool, hint *PathHint,
 ) {
 	tr.ascend(pivot, iter, false, hint)
 }
 func (tr *BTreeG[T]) AscendHintMut(pivot T, iter func(item T) bool,
 	hint *PathHint,
 ) {
 	tr.ascend(pivot, iter, true, hint)
 }
 
 // The return value of this function determines whether we should keep iterating
@@ -693,19  704,32 @@ func (tr *BTreeG[T]) nodeReverse(cn **node[T], iter func(item T) bool, mut bool,
 // Pass nil for pivot to scan all item in descending order
 // Return false to stop iterating
 func (tr *BTreeG[T]) Descend(pivot T, iter func(item T) bool) {
-	tr.descend(pivot, iter, false)
 	tr.descend(pivot, iter, false, nil)
 }
 func (tr *BTreeG[T]) DescendMut(pivot T, iter func(item T) bool) {
-	tr.descend(pivot, iter, true)
 	tr.descend(pivot, iter, true, nil)
 }
-func (tr *BTreeG[T]) descend(pivot T, iter func(item T) bool, mut bool) {
 func (tr *BTreeG[T]) descend(pivot T, iter func(item T) bool, mut bool,
 	hint *PathHint,
 ) {
 	if tr.lock(mut) {
 		defer tr.unlock(mut)
 	}
 	if tr.root == nil {
 		return
 	}
-	tr.nodeDescend(&tr.root, pivot, nil, 0, iter, mut)
 	tr.nodeDescend(&tr.root, pivot, hint, 0, iter, mut)
 }
 
 func (tr *BTreeG[T]) DescendHint(pivot T, iter func(item T) bool,
 	hint *PathHint,
 ) {
 	tr.descend(pivot, iter, false, hint)
 }
 func (tr *BTreeG[T]) DescendHintMut(pivot T, iter func(item T) bool,
 	hint *PathHint,
 ) {
 	tr.descend(pivot, iter, true, hint)
 }
 
 func (tr *BTreeG[T]) nodeDescend(cn **node[T], pivot T, hint *PathHint,
@@ -1104,6  1128,7 @@ type IterG[T any] struct {
 	seeked  bool
 	atstart bool
 	atend   bool
 	stack0  [4]iterStackItemG[T]
 	stack   []iterStackItemG[T]
 	item    T
 }
@@ -1128,12  1153,21 @@ func (tr *BTreeG[T]) iter(mut bool) IterG[T] {
 	iter.tr = tr
 	iter.mut = mut
 	iter.locked = tr.lock(iter.mut)
 	iter.stack = iter.stack0[:0]
 	return iter
 }
 
 // Seek to item greater-or-equal-to key.
 // Returns false if there was no item found.
 func (iter *IterG[T]) Seek(key T) bool {
 	return iter.seek(key, nil)
 }
 
 func (iter *IterG[T]) SeekHint(key T, hint *PathHint) bool {
 	return iter.seek(key, hint)
 }
 
 func (iter *IterG[T]) seek(key T, hint *PathHint) bool {
 	if iter.tr == nil {
 		return false
 	}
@@ -1143,8  1177,9 @@ func (iter *IterG[T]) Seek(key T) bool {
 		return false
 	}
 	n := iter.tr.isoLoad(&iter.tr.root, iter.mut)
 	var depth int
 	for {
-		i, found := iter.tr.find(n, key, nil, 0)
 		i, found := iter.tr.find(n, key, hint, depth)
 		iter.stack = append(iter.stack, iterStackItemG[T]{n, i})
 		if found {
 			iter.item = n.items[i]
@@ -1155,6  1190,7 @@ func (iter *IterG[T]) Seek(key T) bool {
 			return iter.Next()
 		}
 		n = iter.tr.isoLoad(&(*n.children)[i], iter.mut)
 		depth  
 	}
 }
 

diff --git a/btreeg_test.go b/btreeg_test.go
@@ -146,6  146,66 @@ func TestGenericDescend(t *testing.T) {
 	}
 }
 
 func TestGenericDescendHint(t *testing.T) {
 	tr := testNewBTree()
 	var count int
 	var hint PathHint
 	tr.DescendHint(testMakeItem(rand.Int()), func(item testKind) bool {
 		count  
 		return true
 	}, &hint)
 	if count > 0 {
 		t.Fatalf("expected 0, got %v", count)
 	}
 	var keys []testKind
 	for i := 0; i < 1000; i  = 10 {
 		keys = append(keys, testMakeItem(i))
 		tr.Set(keys[len(keys)-1])
 	}
 	var exp []testKind
 	tr.Reverse(func(item testKind) bool {
 		exp = append(exp, item)
 		return true
 	})
 	for i := 999; i >= 0; i-- {
 		key := testMakeItem(i)
 		var all []testKind
 		tr.DescendHint(key, func(item testKind) bool {
 			all = append(all, item)
 			return true
 		}, &hint)
 		for len(exp) > 0 && tr.Less(key, exp[0]) {
 			exp = exp[1:]
 		}
 		var count int
 		tr.DescendHint(key, func(item testKind) bool {
 			if count == (i 1)%tr.max {
 				return false
 			}
 			count  
 			return true
 		}, &hint)
 		if count > len(exp) {
 			t.Fatalf("expected 1, got %v", count)
 		}
 		if !kindsAreEqual(exp, all) {
 			fmt.Printf("exp: %v\n", exp)
 			fmt.Printf("all: %v\n", all)
 			t.Fatal("mismatch")
 		}
 		for j := 0; j < tr.Len(); j   {
 			count = 0
 			tr.DescendHint(key, func(item testKind) bool {
 				if count == j {
 					return false
 				}
 				count  
 				return true
 			}, &hint)
 		}
 	}
 }
 
 func TestGenericAscend(t *testing.T) {
 	tr := testNewBTree()
 	var count int
@@ -190,6  250,51 @@ func TestGenericAscend(t *testing.T) {
 	}
 }
 
 func TestGenericAscendHint(t *testing.T) {
 	tr := testNewBTree()
 	var hint PathHint
 	var count int
 	tr.AscendHint(testMakeItem(1), func(item testKind) bool {
 		count  
 		return true
 	}, &hint)
 	if count > 0 {
 		t.Fatalf("expected 0, got %v", count)
 	}
 	var keys []testKind
 	for i := 0; i < 1000; i  = 10 {
 		keys = append(keys, testMakeItem(i))
 		tr.Set(keys[len(keys)-1])
 		tr.sane()
 	}
 	exp := keys
 	for i := -1; i < 1000; i   {
 		key := testMakeItem(i)
 		var all []testKind
 		tr.AscendHint(key, func(item testKind) bool {
 			all = append(all, item)
 			return true
 		}, &hint)
 		for len(exp) > 0 && tr.Less(exp[0], key) {
 			exp = exp[1:]
 		}
 		var count int
 		tr.AscendHint(key, func(item testKind) bool {
 			if count == (i 1)%tr.max {
 				return false
 			}
 			count  
 			return true
 		}, &hint)
 		if count > len(exp) {
 			t.Fatalf("expected 1, got %v", count)
 		}
 		if !kindsAreEqual(exp, all) {
 			t.Fatal("mismatch")
 		}
 	}
 }
 
 func TestGenericItems(t *testing.T) {
 	tr := testNewBTree()
 	if len(tr.Items()) != 0 {
@@ -1328,6  1433,37 @@ func TestGenericIterSeek(t *testing.T) {
 	}
 }
 
 func TestGenericIterSeekHint(t *testing.T) {
 	tr := NewBTreeG(func(a, b int) bool {
 		return a < b
 	})
 	var all []int
 	for i := 0; i < 10000; i   {
 		tr.Set(i * 2)
 		all = append(all, i)
 	}
 	var hint PathHint
 	_ = all
 	{
 		iter := tr.Iter()
 		var vals []int
 		for ok := iter.SeekHint(501, &hint); ok; ok = iter.Next() {
 			vals = append(vals, iter.Item())
 		}
 		iter.Release()
 		assert(vals[0] == 502 && vals[1] == 504)
 	}
 	{
 		iter := tr.Iter()
 		var vals []int
 		for ok := iter.SeekHint(501, &hint); ok; ok = iter.Prev() {
 			vals = append(vals, iter.Item())
 		}
 		iter.Release()
 		assert(vals[0] == 502 && vals[1] == 500)
 	}
 }
 
 func TestGenericIterSeekPrefix(t *testing.T) {
 	tr := NewBTreeG(func(a, b int) bool {
 		return a < b