Fix random distribution in bitArray.PickRandom (#2534)

* Fix random distribution in bitArray.PickRandom

Previously it was very biased. 63 "_" followed by a single "x" had
much greater odds of being chosen. Additionally, the last element was
skewed. This fixes that by first preproccessing the set of all true
indices, and then randomly selecting a single element from there.

This commit also makes the code here significantly simpler, and
improves test cases.

* unlock mtx right after we select true indices

CHANGELOG_PENDING.md

@@ -45,3 +45,4 @@ timeoutPrecommit before starting next round
 - [evidence] \#2515 fix db iter leak (@goolAdapter)
 - [common/bit_array] Fixed a bug in the `Or` function
 - [common/bit_array] Fixed a bug in the `Sub` function (@bradyjoestar)
+- [common] \#2534 make bit array's PickRandom choose uniformly from true bits

libs/common/bit_array.go

@@ -234,49 +234,53 @@ func (bA *BitArray) IsFull() bool {
 	return (lastElem+1)&((uint64(1)<<uint(lastElemBits))-1) == 0
 }
 
-// PickRandom returns a random index in the bit array, and its value.
+// PickRandom returns a random index for a set bit in the bit array.
+// If there is no such value, it returns 0, false.
 // It uses the global randomness in `random.go` to get this index.
 func (bA *BitArray) PickRandom() (int, bool) {
 	if bA == nil {
 		return 0, false
 	}
+
 	bA.mtx.Lock()
-	defer bA.mtx.Unlock()
+	trueIndices := bA.getTrueIndices()
+	bA.mtx.Unlock()
 
-	length := len(bA.Elems)
-	if length == 0 {
+	if len(trueIndices) == 0 { // no bits set to true
 		return 0, false
 	}
-	randElemStart := RandIntn(length)
-	for i := 0; i < length; i++ {
-		elemIdx := ((i + randElemStart) % length)
-		if elemIdx < length-1 {
-			if bA.Elems[elemIdx] > 0 {
-				randBitStart := RandIntn(64)
-				for j := 0; j < 64; j++ {
-					bitIdx := ((j + randBitStart) % 64)
-					if (bA.Elems[elemIdx] & (uint64(1) << uint(bitIdx))) > 0 {
-						return 64*elemIdx + bitIdx, true
-					}
-				}
-				PanicSanity("should not happen")
-			}
-		} else {
-			// Special case for last elem, to ignore straggler bits
-			elemBits := bA.Bits % 64
-			if elemBits == 0 {
-				elemBits = 64
-			}
-			randBitStart := RandIntn(elemBits)
-			for j := 0; j < elemBits; j++ {
-				bitIdx := ((j + randBitStart) % elemBits)
-				if (bA.Elems[elemIdx] & (uint64(1) << uint(bitIdx))) > 0 {
-					return 64*elemIdx + bitIdx, true
-				}
+
+	return trueIndices[RandIntn(len(trueIndices))], true
+}
+
+func (bA *BitArray) getTrueIndices() []int {
+	trueIndices := make([]int, 0, bA.Bits)
+	curBit := 0
+	numElems := len(bA.Elems)
+	// set all true indices
+	for i := 0; i < numElems-1; i++ {
+		elem := bA.Elems[i]
+		if elem == 0 {
+			curBit += 64
+			continue
+		}
+		for j := 0; j < 64; j++ {
+			if (elem & (uint64(1) << uint64(j))) > 0 {
+				trueIndices = append(trueIndices, curBit)
 			}
+			curBit++
+		}
+	}
+	// handle last element
+	lastElem := bA.Elems[numElems-1]
+	numFinalBits := bA.Bits - curBit
+	for i := 0; i < numFinalBits; i++ {
+		if (lastElem & (uint64(1) << uint64(i))) > 0 {
+			trueIndices = append(trueIndices, curBit)
 		}
+		curBit++
 	}
-	return 0, false
+	return trueIndices
 }
 
 // String returns a string representation of BitArray: BA{<bit-string>},

libs/common/bit_array_test.go

@@ -107,16 +107,29 @@ func TestSub(t *testing.T) {
 }
 
 func TestPickRandom(t *testing.T) {
-	for idx := 0; idx < 123; idx++ {
-		bA1 := NewBitArray(123)
-		bA1.SetIndex(idx, true)
-		index, ok := bA1.PickRandom()
-		if !ok {
-			t.Fatal("Expected to pick element but got none")
-		}
-		if index != idx {
-			t.Fatalf("Expected to pick element at %v but got wrong index", idx)
-		}
+	empty16Bits := "________________"
+	empty64Bits := empty16Bits + empty16Bits + empty16Bits + empty16Bits
+	testCases := []struct {
+		bA string
+		ok bool
+	}{
+		{`null`, false},
+		{`"x"`, true},
+		{`"` + empty16Bits + `"`, false},
+		{`"x` + empty16Bits + `"`, true},
+		{`"` + empty16Bits + `x"`, true},
+		{`"x` + empty16Bits + `x"`, true},
+		{`"` + empty64Bits + `"`, false},
+		{`"x` + empty64Bits + `"`, true},
+		{`"` + empty64Bits + `x"`, true},
+		{`"x` + empty64Bits + `x"`, true},
+	}
+	for _, tc := range testCases {
+		var bitArr *BitArray
+		err := json.Unmarshal([]byte(tc.bA), &bitArr)
+		require.NoError(t, err)
+		_, ok := bitArr.PickRandom()
+		require.Equal(t, tc.ok, ok, "PickRandom got an unexpected result on input %s", tc.bA)
 	}
 }