diff --git a/cardinality/cardinality_test.go b/cardinality/cardinality_test.go index bb6959be..07473631 100644 --- a/cardinality/cardinality_test.go +++ b/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 +} diff --git a/cardinality/roaring64.go b/cardinality/roaring64.go index 4beaad0f..f453d46a 100644 --- a/cardinality/roaring64.go +++ b/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] { diff --git a/cardinality/roaring64_test.go b/cardinality/roaring64_test.go new file mode 100644 index 00000000..33ff8309 --- /dev/null +++ b/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()) +} diff --git a/go.mod b/go.mod index 49cef52b..add3959a 100644 --- a/go.mod +++ b/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