[ready for review] btreeg: add support for no alloc iter

btreeg.go

@@ -1690,8 +1690,31 @@ func (iter *IterG[T]) Release() {
 		iter.tr.unlock(iter.mut)
 		iter.locked = false
 	}
-	iter.stack = nil
-	iter.tr = nil
+
+	// 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
+}
+
+// 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
+
+	iter.tr = tr
+	iter.mut = mut
+
+	iter.locked = tr.lock(iter.mut)
+	if iter.stack == nil {
+		iter.stack = iter.stack0[:0]
+	} else {
+		iter.stack = iter.stack[:0]
+	}
 }
 
 // Next moves iterator to the next item in iterator.

btreeg_test.go

@@ -1872,3 +1872,113 @@ func TestContiguousDeleteRangeLoad(t *testing.T) {
 	)
 
 }
+
+type largeItem struct {
+	a uint64
+	b uint64
+	c uint64
+	d uint64
+}
+
+func useIterator(iter IterG[largeItem]) {
+	iter.Seek(largeItem{a: 0})
+	defer iter.Release()
+
+	// Iterate over 10 items beginning the seeked item.
+	assert(iter.Item().a == 0)
+}
+
+func useIteratorPointer(iter *IterG[largeItem]) {
+	iter.Seek(largeItem{a: 0})
+	defer iter.Release()
+
+	assert(iter.Item().a == 0)
+}
+
+// This benchmark proves that there exist cases where the iterator creation can
+// cause an allocation
+//
+// Run using: go test -run=^$ -bench ^BenchmarkIteratorCreationAlloc$ github.com/tidwall/btree
+func BenchmarkIteratorCreationAlloc(b *testing.B) {
+	tr := NewBTreeG(func(a, b largeItem) bool {
+		return a.a < b.a
+	})
+
+	for i := 0; i < 1; i++ {
+		tr.Set(largeItem{a: uint64(i * 2), b: uint64(i * 2), c: uint64(i * 2), d: uint64(i * 2)})
+	}
+
+	iter := tr.Iter()
+	iter.Seek(largeItem{a: 0})
+	assert(iter.Item().a == 0)
+
+	// The following will cause 1 allocation per op. Note that this allocation is not due
+	// to iter.stack slice allocating to grow larger. The benchmark is only performing a
+	// single seek, and the btree only has 1 item. The allocation is due to the iterator
+	// escaping to the heap.
+	b.Run("no reuse", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			iter := tr.Iter()
+			useIterator(iter)
+		}
+	})
+
+	reusableIter := tr.Iter()
+	reusableIterPointer := &reusableIter
+
+	// The following will cause 0 allocations per op, since a re-usable iterator is used.
+	b.Run("reuse", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			reusableIterPointer.Init(tr, false)
+			useIteratorPointer(reusableIterPointer)
+		}
+	})
+}
+
+// Ensure that the re-usable iterator works as expected even if the tree
+// is receiving new writes after every iteration.
+func TestBenchmarkIteratorReuseWorks(t *testing.T) {
+	tr := NewBTreeGOptions(func(a, b largeItem) bool {
+		return a.a < b.a
+	}, Options{
+		NoLocks: true,
+	})
+
+	tr.Set(largeItem{a: 0, b: 0, c: 0, d: 0})
+	iter := tr.Iter()
+	reusableIter := &iter
+
+	found := reusableIter.Seek(largeItem{a: 0})
+	assert(found)
+	assert(reusableIter.Item().a == 0)
+	reusableIter.Release()
+
+	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
+
+	for i := 1; i < 10000; i++ {
+		// Insert a single item. Reset iterator, randomly seek b/w 0 to i, and ensure that the
+		// iterator is seeing the expected items.
+		tr.Set(largeItem{a: uint64(i), b: uint64(i), c: uint64(i), d: uint64(i)})
+
+		seekTo := rng.Intn(i + 1)
+
+		// Reset the iterator.
+		reusableIter.Init(tr, false)
+
+		// Seek to the random position.
+		found = reusableIter.Seek(largeItem{a: uint64(seekTo)})
+		assert(found)
+
+		nextExpectedItem := uint64(seekTo) + 1
+		for iter.Next() {
+			assert(iter.Item().a == nextExpectedItem)
+			nextExpectedItem++
+		}
+
+		assert(nextExpectedItem == uint64(i)+1)
+
+		reusableIter.Release()
+	}
+}