Add test coverage for GetSetBits

Author: danielway-wk

bitarray/bitarray_test.go

@@ -20,6 +20,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestBitOperations(t *testing.T) {
@@ -195,6 +196,50 @@ func BenchmarkGetBit(b *testing.B) {
 	}
 }
 
+func TestGetSetBits(t *testing.T) {
+	ba := newBitArray(1000)
+	buf := make([]uint64, 0, 5)
+
+	require.NoError(t, ba.SetBit(1))
+	require.NoError(t, ba.SetBit(4))
+	require.NoError(t, ba.SetBit(8))
+	require.NoError(t, ba.SetBit(63))
+	require.NoError(t, ba.SetBit(64))
+	require.NoError(t, ba.SetBit(200))
+	require.NoError(t, ba.SetBit(1000))
+
+	assert.Equal(t, []uint64{1, 4, 8, 63, 64}, ba.GetSetBits(0, buf))
+	assert.Equal(t, []uint64{63, 64, 200, 1000}, ba.GetSetBits(10, buf))
+	assert.Equal(t, []uint64{63, 64, 200, 1000}, ba.GetSetBits(63, buf))
+	assert.Equal(t, []uint64{200, 1000}, ba.GetSetBits(128, buf))
+
+	require.NoError(t, ba.ClearBit(4))
+	require.NoError(t, ba.ClearBit(64))
+	assert.Equal(t, []uint64{1, 8, 63, 200, 1000}, ba.GetSetBits(0, buf))
+	assert.Empty(t, ba.GetSetBits(1001, buf))
+
+	ba.Reset()
+	assert.Empty(t, ba.GetSetBits(0, buf))
+}
+
+func BenchmarkGetSetBits(b *testing.B) {
+	numItems := uint64(168000)
+
+	ba := newBitArray(numItems)
+	for i := uint64(0); i < numItems; i++ {
+		if i%13 == 0 || i%5 == 0 {
+			require.NoError(b, ba.SetBit(i))
+		}
+	}
+
+	buf := make([]uint64, 0, ba.Capacity())
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ba.GetSetBits(0, buf)
+	}
+}
+
 func TestEquality(t *testing.T) {
 	ba := newBitArray(s + 1)
 	other := newBitArray(s + 1)

bitarray/sparse_bitarray_test.go

@@ -20,6 +20,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestGetCompressedBit(t *testing.T) {
@@ -76,6 +77,48 @@ func BenchmarkSetCompressedBit(b *testing.B) {
 	}
 }
 
+func TestGetSetCompressedBits(t *testing.T) {
+	ba := newSparseBitArray()
+	buf := make([]uint64, 0, 5)
+
+	require.NoError(t, ba.SetBit(1))
+	require.NoError(t, ba.SetBit(4))
+	require.NoError(t, ba.SetBit(8))
+	require.NoError(t, ba.SetBit(63))
+	require.NoError(t, ba.SetBit(64))
+	require.NoError(t, ba.SetBit(200))
+	require.NoError(t, ba.SetBit(1000))
+
+	assert.Equal(t, []uint64{1, 4, 8, 63, 64}, ba.GetSetBits(0, buf))
+	assert.Equal(t, []uint64{63, 64, 200, 1000}, ba.GetSetBits(10, buf))
+	assert.Equal(t, []uint64{63, 64, 200, 1000}, ba.GetSetBits(63, buf))
+	assert.Equal(t, []uint64{200, 1000}, ba.GetSetBits(128, buf))
+
+	require.NoError(t, ba.ClearBit(4))
+	require.NoError(t, ba.ClearBit(64))
+	assert.Equal(t, []uint64{1, 8, 63, 200, 1000}, ba.GetSetBits(0, buf))
+	assert.Empty(t, ba.GetSetBits(1001, buf))
+
+	ba.Reset()
+	assert.Empty(t, ba.GetSetBits(0, buf))
+}
+
+func BenchmarkGetSetCompressedBits(b *testing.B) {
+	ba := newSparseBitArray()
+	for i := uint64(0); i < 168000; i++ {
+		if i%13 == 0 || i%5 == 0 {
+			require.NoError(b, ba.SetBit(i))
+		}
+	}
+
+	buf := make([]uint64, 0, ba.Capacity())
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ba.GetSetBits(0, buf)
+	}
+}
+
 func TestClearCompressedBit(t *testing.T) {
 	ba := newSparseBitArray()
 	ba.SetBit(5)