fix(bits): handle last element in PickRandom (backport cometbft#2899) (…

…cometbft#2902) (#35) (#36)

Follow-up to cometbft#2841

---

#### PR checklist

- [x] Tests written/updated
- [ ] ~~Changelog entry added in `.changelog` (we use
[unclog](https://github.com/informalsystems/unclog) to manage our
changelog)~~
- [ ] ~~Updated relevant documentation (`docs/` or `spec/`) and code
comments~~
- [x] Title follows the [Conventional
Commits](https://www.conventionalcommits.org/en/v1.0.0/) spec
<hr>This is an automatic backport of pull request cometbft#2899 done by
[Mergify](https://mergify.com).

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Anton Kaliaev <[email protected]>
(cherry picked from commit 08efc4c)

Co-authored-by: Adam Tucker <[email protected]>

libs/bits/bit_array.go

@@ -264,9 +264,17 @@ func (bA *BitArray) PickRandom() (int, bool) {
 func (bA *BitArray) getNumTrueIndices() int {
 	count := 0
 	numElems := len(bA.Elems)
-	for i := 0; i < numElems; i++ {
+	// handle all elements except the last one
+	for i := 0; i < numElems-1; i++ {
 		count += bits.OnesCount64(bA.Elems[i])
 	}
+	// handle last element
+	numFinalBits := bA.Bits - (numElems-1)*64
+	for i := 0; i < numFinalBits; i++ {
+		if (bA.Elems[numElems-1] & (uint64(1) << uint64(i))) > 0 {
+			count++
+		}
+	}
 	return count
 }
 

libs/bits/bit_array_test.go

@@ -173,6 +173,8 @@ func TestGetNumTrueIndices(t *testing.T) {
 		require.NoError(t, err)
 		result := bitArr.getNumTrueIndices()
 		require.Equal(t, tc.ExpectedResult, result, "for input %s, expected %d, got %d", tc.Input, tc.ExpectedResult, result)
+		result = bitArr.Not().getNumTrueIndices()
+		require.Equal(t, bitArr.Bits-result, bitArr.getNumTrueIndices())
 	}
 }