Define BitArray.Count and add test coverage

Workiva · May 16, 2023 · d95ee5d · d95ee5d
1 parent 9eab2df
commit d95ee5d
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 0 deletions.
diff --git a/bitarray/bitarray_test.go b/bitarray/bitarray_test.go
@@ -143,6 +143,28 @@ func TestIsEmpty(t *testing.T) {
 	assert.False(t, ba.IsEmpty())
 }
 
+func TestCount(t *testing.T) {
+	ba := newBitArray(500)
+	assert.Equal(t, uint64(0), ba.Count())
+
+	require.NoError(t, ba.SetBit(0))
+	assert.Equal(t, uint64(1), ba.Count())
+
+	require.NoError(t, ba.SetBit(40))
+	require.NoError(t, ba.SetBit(64))
+	require.NoError(t, ba.SetBit(100))
+	require.NoError(t, ba.SetBit(200))
+	require.NoError(t, ba.SetBit(469))
+	require.NoError(t, ba.SetBit(500))
+	assert.Equal(t, uint64(7), ba.Count())
+
+	require.NoError(t, ba.ClearBit(200))
+	assert.Equal(t, uint64(6), ba.Count())
+
+	ba.Reset()
+	assert.Equal(t, uint64(0), ba.Count())
+}
+
 func TestClear(t *testing.T) {
 	ba := newBitArray(10)
 

diff --git a/bitarray/interface.go b/bitarray/interface.go
@@ -59,6 +59,8 @@ type BitArray interface {
 	// in the case of a dense bit array or the highest possible
 	// seen capacity of the sparse array.
 	Capacity() uint64
+	// Count returns the number of set bits in this array.
+	Count() uint64
 	// Or will bitwise or the two bitarrays and return a new bitarray
 	// representing the result.
 	Or(other BitArray) BitArray

diff --git a/bitarray/sparse_bitarray_test.go b/bitarray/sparse_bitarray_test.go
@@ -119,6 +119,28 @@ func BenchmarkGetSetCompressedBits(b *testing.B) {
 	}
 }
 
+func TestCompressedCount(t *testing.T) {
+	ba := newSparseBitArray()
+	assert.Equal(t, uint64(0), ba.Count())
+
+	require.NoError(t, ba.SetBit(0))
+	assert.Equal(t, uint64(1), ba.Count())
+
+	require.NoError(t, ba.SetBit(40))
+	require.NoError(t, ba.SetBit(64))
+	require.NoError(t, ba.SetBit(100))
+	require.NoError(t, ba.SetBit(200))
+	require.NoError(t, ba.SetBit(469))
+	require.NoError(t, ba.SetBit(500))
+	assert.Equal(t, uint64(7), ba.Count())
+
+	require.NoError(t, ba.ClearBit(200))
+	assert.Equal(t, uint64(6), ba.Count())
+
+	ba.Reset()
+	assert.Equal(t, uint64(0), ba.Count())
+}
+
 func TestClearCompressedBit(t *testing.T) {
 	ba := newSparseBitArray()
 	ba.SetBit(5)