diff --git a/btreeg.go b/btreeg.go index 218a2db..de1067e 100644 --- a/btreeg.go +++ b/btreeg.go @@ -1691,25 +1691,33 @@ func (iter *IterG[T]) Release() { iter.locked = false } - // Preserve the backing memory for the stack, so that the iterator can be re-used without - // allocating. - stack := iter.stack[:0] - *iter = IterG[T]{} - iter.stack = stack + // Preserve stack backing array for reuse. Clear fields manually to avoid heap escape. + iter.stack = iter.stack[:0] + iter.tr = nil + iter.mut = false + iter.seeked = false + iter.atstart = false + iter.atend = false + var empty T + iter.item = empty } // Init is used to initialize an existing iterator with a new tree. Release must've // been called on the iterator before re-using it using Init. func (iter *IterG[T]) Init(tr *BTreeG[T], mut bool) { - // Re-use the stack, but 0 out the rest of the memory. - stack := iter.stack[:0] - *iter = IterG[T]{} - iter.stack = stack - + // Reuse stack and clear fields manually to avoid struct copy overhead. + iter.stack = iter.stack[:0] iter.tr = tr iter.mut = mut - - iter.locked = tr.lock(iter.mut) + iter.seeked = false + iter.atstart = false + iter.atend = false + iter.locked = false + if tr != nil { + iter.locked = tr.lock(iter.mut) + } + var empty T + iter.item = empty if iter.stack == nil { iter.stack = iter.stack0[:0] } else { diff --git a/btreeg_test.go b/btreeg_test.go index 686758b..45e5bc8 100644 --- a/btreeg_test.go +++ b/btreeg_test.go @@ -4,6 +4,7 @@ import ( "fmt" "math/rand" "os" + "reflect" "runtime" "sort" "strconv" @@ -1895,6 +1896,28 @@ func useIteratorPointer(iter *IterG[largeItem]) { assert(iter.Item().a == 0) } +func assertIteratorCleared[T any](t *testing.T, iter IterG[T]) { + t.Helper() + v := reflect.ValueOf(iter) + vt := v.Type() + for i := 0; i < v.NumField(); i++ { + name := vt.Field(i).Name + switch name { + case "stack", "stack0": + continue + } + if !v.Field(i).IsZero() { + t.Fatalf("iterator field %s was not reset", name) + } + } + if len(iter.stack) != 0 { + t.Fatal("iterator stack length not reset") + } + if cap(iter.stack) == 0 { + t.Fatal("iterator stack capacity should be preserved for reuse") + } +} + // This benchmark proves that there exist cases where the iterator creation can // cause an allocation // @@ -1982,3 +2005,121 @@ func TestBenchmarkIteratorReuseWorks(t *testing.T) { reusableIter.Release() } } + +func TestGenericIteratorRelease(t *testing.T) { + tr := testNewBTree() + for i := 0; i < 100; i++ { + tr.Set(testMakeItem(i)) + } + iter := tr.Iter() + if !iter.First() || !tr.eq(iter.Item(), testMakeItem(0)) { + panic("!") + } + iter.Release() + assertIteratorCleared(t, iter) +} + +func TestGenericIteratorInit(t *testing.T) { + tr := testNewBTree() + for i := 0; i < 50; i++ { + tr.Set(testMakeItem(i * 2)) + } + iter := tr.Iter() + if !iter.First() || !tr.eq(iter.Item(), testMakeItem(0)) { + panic("!") + } + iter.Release() + iter.Init(tr, false) + if iter.tr == nil || iter.mut || iter.seeked || len(iter.stack) != 0 { + panic("!") + } + if !iter.First() || !tr.eq(iter.Item(), testMakeItem(0)) { + panic("!") + } + count := 1 + for iter.Next() { + if !tr.eq(iter.Item(), testMakeItem(count*2)) { + panic("!") + } + count++ + } + if count != 50 { + panic("!") + } + iter.Init(nil, false) + assertIteratorCleared(t, iter) +} + +func TestGenericIteratorReuse(t *testing.T) { + tr := testNewBTree() + for i := 0; i < 100; i++ { + tr.Set(testMakeItem(i)) + } + iter := tr.Iter() + for round := 0; round < 1000; round++ { + iter.Init(tr, false) + if !iter.First() { + panic("!") + } + count := 0 + for { + if !tr.eq(iter.Item(), testMakeItem(count)) { + panic("!") + } + count++ + if !iter.Next() { + break + } + } + if count != 100 { + panic("!") + } + iter.Release() + } +} + +func BenchmarkIteratorRelease(b *testing.B) { + tr := NewBTreeG(testLess) + for i := 0; i < 10000; i++ { + tr.Set(testMakeItem(i)) + } + iter := tr.Iter() + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + iter.First() + iter.Release() + iter.Init(tr, false) + } +} + +func BenchmarkIteratorReuse(b *testing.B) { + tr := NewBTreeG(testLess) + for i := 0; i < 1000; i++ { + tr.Set(testMakeItem(i)) + } + b.Run("Recreate", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + iter := tr.Iter() + iter.First() + for iter.Next() { + _ = iter.Item() + } + iter.Release() + } + }) + b.Run("Reuse", func(b *testing.B) { + iter := tr.Iter() + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + iter.Init(tr, false) + iter.First() + for iter.Next() { + _ = iter.Item() + } + iter.Release() + } + }) +}