cometbft · melekes · May 2, 2024 · May 2, 2024 · May 2, 2024 · May 2, 2024
@@ -0,0 +1,2 @@
+- `[internal/bits]` 10x speedup creating initialized bitArrays, which speedsup extendedCommit.BitArray(). This is used in consensus vote gossip.
+  ([\#2959](https://github.com/cometbft/cometbft/pull/2841)).
@@ -32,7 +32,27 @@ func NewBitArray(bits int) *BitArray {
 	}
 }
 
-// Size returns the number of bits in the bitarray
+// NewBitArrayFromFn returns a new bit array.
+// It returns nil if the number of bits is zero.
+// It initializes the `i`th bit to the value of `fn(i)`.
+func NewBitArrayFromFn(bits int, fn func(int) bool) *BitArray {
+	if bits <= 0 {
+		return nil
+	}
+	bA := &BitArray{
+		Bits:  bits,
+		Elems: make([]uint64, (bits+63)/64),
+	}
+	for i := 0; i < bits; i++ {
+		v := fn(i)
+		if v {
+			bA.Elems[i/64] |= (uint64(1) << uint(i%64))
+		}
+	}
+	return bA
+}
+
+// Size returns the number of bits in the bitarray.
 func (bA *BitArray) Size() int {
 	if bA == nil {
 		return 0

@@ -19,24 +19,17 @@ var (
 	full64bits  = full16bits + full16bits + full16bits + full16bits
 )
 
-func randBitArray(bits int) (*BitArray, []byte) {
+func randBitArray(bits int) *BitArray {
 	src := cmtrand.Bytes((bits + 7) / 8)
-	bA := NewBitArray(bits)
-	for i := 0; i < len(src); i++ {
-		for j := 0; j < 8; j++ {
-			if i*8+j >= bits {
-				return bA, src
-			}
-			setBit := src[i]&(1<<uint(j)) > 0
-			bA.SetIndex(i*8+j, setBit)
-		}
+	srcIndexToBit := func(i int) bool {
+		return src[i/8]&(1<<uint(i%8)) > 0
 	}
-	return bA, src
+	return NewBitArrayFromFn(bits, srcIndexToBit)
 }
 
 func TestAnd(t *testing.T) {
-	bA1, _ := randBitArray(51)
-	bA2, _ := randBitArray(31)
+	bA1 := randBitArray(51)
+	bA2 := randBitArray(31)
 	bA3 := bA1.And(bA2)
 
 	var bNil *BitArray
@@ -59,16 +52,16 @@ func TestAnd(t *testing.T) {
 }
 
 func TestOr(t *testing.T) {
-	bA1, _ := randBitArray(51)
-	bA2, _ := randBitArray(31)
+	bA1 := randBitArray(57)
+	bA2 := randBitArray(31)
 	bA3 := bA1.Or(bA2)
 
 	bNil := (*BitArray)(nil)
 	require.Equal(t, bNil.Or(bA1), bA1)
 	require.Equal(t, bA1.Or(nil), bA1)
 	require.Equal(t, bNil.Or(nil), (*BitArray)(nil))
 
-	if bA3.Bits != 51 {
+	if bA3.Bits != 57 {
 		t.Error("Expected max bits")
 	}
 	if len(bA3.Elems) != len(bA1.Elems) {
@@ -80,6 +73,10 @@ func TestOr(t *testing.T) {
 			t.Error("Wrong bit from bA3", i, bA1.GetIndex(i), bA2.GetIndex(i), bA3.GetIndex(i))
 		}
 	}
+	if bA3.getNumTrueIndices() == 0 {
+		t.Error("Expected at least one true bit. " +
+			"This has a false positive rate that is less than 1 in 2^80 (cryptographically improbable).")
+	}
 }
 
 func TestSub(t *testing.T) {
@@ -276,7 +273,7 @@ func TestEmptyFull(t *testing.T) {
 
 func TestUpdateNeverPanics(_ *testing.T) {
 	newRandBitArray := func(n int) *BitArray {
-		ba, _ := randBitArray(n)
+		ba := randBitArray(n)
 		return ba
 	}
 	pairs := []struct {

@@ -1179,12 +1179,12 @@ func (ec *ExtendedCommit) Size() int {
 // Implements VoteSetReader.
 func (ec *ExtendedCommit) BitArray() *bits.BitArray {
 	if ec.bitArray == nil {
-		ec.bitArray = bits.NewBitArray(len(ec.ExtendedSignatures))
-		for i, extCommitSig := range ec.ExtendedSignatures {
+		initialBitFn := func(i int) bool {
 			// TODO: need to check the BlockID otherwise we could be counting conflicts,
 			//       not just the one with +2/3 !
-			ec.bitArray.SetIndex(i, extCommitSig.BlockIDFlag != BlockIDFlagAbsent)
+			return ec.ExtendedSignatures[i].BlockIDFlag != BlockIDFlagAbsent
 		}
+		ec.bitArray = bits.NewBitArrayFromFn(len(ec.ExtendedSignatures), initialBitFn)
 	}
 	return ec.bitArray
 }