Allow creation of bitset with 0 size (#116)

* allow creation of bitset with 0 size Signed-off-by: Daniel Rammer <daniel@union.ai> * removed unecessary math import Signed-off-by: Daniel Rammer <daniel@union.ai> * added unit test and fixed lint issue Signed-off-by: Daniel Rammer <daniel@union.ai>
flyteorg · Feb 23, 2022 · d304944 · d304944
1 parent 9f812e9
commit d304944
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/bitarray/bitset.go b/bitarray/bitset.go
@@ -67,7 +67,10 @@ func (s BitSet) DeepCopy() BitSet {
 
 // Initializes a new BitSet of the specified size.
 func NewBitSet(desiredCap uint) *BitSet {
-	// Create enough blocks to contain the number of intended bits.
-	a := make(BitSet, ((desiredCap-1)/blockSize)+1)
+	size := desiredCap / blockSize
+	if desiredCap%blockSize != 0 {
+		size++
+	}
+	a := make(BitSet, size)
 	return &a
 }
diff --git a/bitarray/bitset_test.go b/bitarray/bitset_test.go
@@ -36,6 +36,14 @@ func TestBitSet_Set(t *testing.T) {
 }
 
 func TestNewBitSet(t *testing.T) {
+	t.Run("Empty", func(t *testing.T) {
+		b := new(BitSet)
+		assert.Equal(t, 0, b.BlockCount())
+
+		b = NewBitSet(0)
+		assert.Equal(t, 0, b.BlockCount())
+	})
+
 	t.Run("Block size", func(t *testing.T) {
 		b := NewBitSet(63)
 		assert.Equal(t, 2, b.BlockCount())