Fix random distribution in bitArray.PickRandom

CHANGELOG_PENDING.md

@@ -43,3 +43,4 @@ BUG FIXES:
 - [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)
 	}
 }