From ddee77ae3757a6ad2b1013f1e21a8eee1f9a091c Mon Sep 17 00:00:00 2001
From: Adam Tucker <adam@osmosis.team>
Date: Thu, 2 May 2024 13:12:56 -0600
Subject: [PATCH] fix(bits): handle last element in PickRandom (backport #2899)
 (#2902) (#35)

Follow-up to #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 #2899 done by
[Mergify](https://mergify.com).

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Anton Kaliaev <anton.kalyaev@gmail.com>
(cherry picked from commit 08efc4c419196d62f0d6680eb546025dad6de322)
---
 libs/bits/bit_array.go      | 10 +++++++++-
 libs/bits/bit_array_test.go |  2 ++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/libs/bits/bit_array.go b/libs/bits/bit_array.go
index ad4efe31540..18daf50119c 100644
--- a/libs/bits/bit_array.go
+++ b/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
 }
 
diff --git a/libs/bits/bit_array_test.go b/libs/bits/bit_array_test.go
index dce587ca00f..ad6a3a9f00e 100644
--- a/libs/bits/bit_array_test.go
+++ b/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())
 	}
 }