perf: reuse pooled roaring64 iterator in bitmap64.Each BED-9922

cardinality/cardinality_test.go

@@ -1,10 +1,13 @@
 package cardinality_test
 
 import (
+	"sort"
+	"sync"
 	"testing"
 
 	"github.com/specterops/dawgs/cardinality"
 	"github.com/specterops/dawgs/graph"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -44,3 +47,80 @@ func TestNodeSetToDuplex(t *testing.T) {
 	require.True(t, duplex.Contains(1))
 	require.True(t, duplex.Contains(2))
 }
+
+// collectEach drains a Duplex via Each into a sorted slice.
+func collectEach(duplex cardinality.Duplex[uint64]) []uint64 {
+	var collected []uint64
+
+	duplex.Each(func(value uint64) bool {
+		collected = append(collected, value)
+		return true
+	})
+
+	sort.Slice(collected, func(i, j int) bool { return collected[i] < collected[j] })
+
+	return collected
+}
+
+func TestBitmap64Each(t *testing.T) {
+	// Empty bitmap yields nothing.
+	require.Empty(t, collectEach(cardinality.NewBitmap64()))
+
+	// Values spanning multiple roaring containers are all returned, and repeated
+	// calls return identical results (the iterator is pooled and reset per call).
+	values := []uint64{0, 1, 5, 1 << 16, (1 << 16) + 3, 1 << 32, (1 << 32) + 7}
+	duplex := cardinality.NewBitmap64With(values...)
+
+	require.Equal(t, values, collectEach(duplex))
+	require.Equal(t, values, collectEach(duplex))
+
+	// Early termination stops iteration.
+	var seen int
+	duplex.Each(func(uint64) bool {
+		seen++
+		return false
+	})
+	require.Equal(t, 1, seen)
+}
+
+func TestBitmap64EachConcurrent(t *testing.T) {
+	values := []uint64{0, 2, 4, 1 << 16, 1 << 32, (1 << 32) + 9}
+	duplex := cardinality.NewBitmap64With(values...)
+
+	var waitGroup sync.WaitGroup
+	for worker := 0; worker < 16; worker++ {
+		waitGroup.Add(1)
+		go func() {
+			defer waitGroup.Done()
+			for iteration := 0; iteration < 100; iteration++ {
+				assert.Equal(t, values, collectEach(duplex))
+			}
+		}()
+	}
+
+	waitGroup.Wait()
+}
+
+func BenchmarkBitmap64Each(b *testing.B) {
+	values := make([]uint64, 0, 4096)
+	for value := uint64(0); value < 4096; value++ {
+		values = append(values, value*3)
+	}
+
+	duplex := cardinality.NewBitmap64With(values...)
+
+	var sink uint64
+	delegate := func(value uint64) bool {
+		sink += value
+		return true
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for iteration := 0; iteration < b.N; iteration++ {
+		duplex.Each(delegate)
+	}
+
+	_ = sink
+}

cardinality/roaring64.go

@@ -1,9 +1,24 @@
 package cardinality
 
 import (
+	"sync"
+
 	"github.com/RoaringBitmap/roaring/v2/roaring64"
 )
 
+// bitmap64IteratorPool reuses roaring64 iterators across Each calls. roaring64.IntIterator64
+// is resettable via Initialize and embeds its container iterator, so a pooled iterator can be
+// re-pointed at each bitmap without allocating a fresh iterator (and its per-container iterator)
+// on every traversal. The pool is safe for concurrent use, and reentrant iteration is safe
+// because each active traversal holds its own iterator until it returns to the pool. Iterators
+// are zeroed before being returned so that pooled iterators do not keep the last iterated
+// bitmap reachable.
+var bitmap64IteratorPool = sync.Pool{
+	New: func() any {
+		return &roaring64.IntIterator64{}
+	},
+}
+
 type bitmap64Iterator struct {
 	iterator roaring64.IntPeekable64
 }
@@ -42,11 +57,17 @@ func (s bitmap64) Clear() {
 }
 
 func (s bitmap64) Each(delegate func(nextValue uint64) bool) {
-	for itr := s.bitmap.Iterator(); itr.HasNext(); {
+	itr := bitmap64IteratorPool.Get().(*roaring64.IntIterator64)
+	itr.Initialize(s.bitmap)
+
+	for itr.HasNext() {
 		if ok := delegate(itr.Next()); !ok {
 			break
 		}
 	}
+
+	*itr = roaring64.IntIterator64{}
+	bitmap64IteratorPool.Put(itr)
 }
 
 func (s bitmap64) Iterator() Iterator[uint64] {

cardinality/roaring64_test.go

@@ -0,0 +1,32 @@
+package cardinality
+
+import (
+	"runtime"
+	"testing"
+	"weak"
+
+	"github.com/RoaringBitmap/roaring/v2/roaring64"
+	"github.com/stretchr/testify/require"
+)
+
+// iterateAndRelease builds a bitmap, traverses it with Each and returns only a weak reference so that the caller holds
+// no strong reference to the bitmap.
+func iterateAndRelease(delegate func(uint64) bool) weak.Pointer[roaring64.Bitmap] {
+	duplex := NewBitmap64With(0, 1, 1<<16, 1<<32, (1<<32)+1).(bitmap64)
+	duplex.Each(delegate)
+
+	return weak.Make(duplex.bitmap)
+}
+
+func TestBitmap64EachDoesNotRetainBitmap(t *testing.T) {
+	// Full traversal and early termination both return the iterator to the pool.
+	fullTraversal := iterateAndRelease(func(uint64) bool { return true })
+	earlyTermination := iterateAndRelease(func(uint64) bool { return false })
+
+	// A single GC cycle must reclaim the bitmaps. A pooled iterator that still references a bitmap keeps it reachable
+	// until the pool's victim cache is cleared on a second GC cycle.
+	runtime.GC()
+
+	require.Nil(t, fullTraversal.Value())
+	require.Nil(t, earlyTermination.Value())
+}

go.mod

@@ -2,7 +2,6 @@ module github.com/specterops/dawgs
 
 go 1.27.2
 
-// Dawgrun requirements
 require (
 	cuelang.org/go v0.17.0
 	github.com/RoaringBitmap/roaring/v2 v2.19.0