diff --git a/README.md b/README.md index 6574328..b3572f9 100644 --- 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 index 3483cc6..9e2327e 100644 --- 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 index 86b8990..b83055d 100644 --- 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 index cf61aa9..255bff0 100644 --- 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