refactor rejection sampling

algorand · Apr 28, 2022 · ab2a2ff · ab2a2ff
1 parent b566381
commit ab2a2ff
Showing 1 changed file with 10 additions and 10 deletions.
diff --git a/crypto/compactcert/coinGenerator.go b/crypto/compactcert/coinGenerator.go
@@ -61,7 +61,7 @@ func (cc *coinChoiceSeed) ToBeHashed() (protocol.HashID, []byte) {
 // coinGenerator is used for extracting "randomized" 64 bits for coin flips
 type coinGenerator struct {
 	shkContext   sha3.ShakeHash
-	signedWeight *big.Int
+	signedWeight uint64
 	threshold    *big.Int
 }
 
@@ -75,12 +75,12 @@ func makeCoinGenerator(choice *coinChoiceSeed) coinGenerator {
 	shk := sha3.NewShake256()
 	shk.Write(rep)
 
-	threshold, signedWt := prepareRejectionSamplingValues(choice.signedWeight)
-	return coinGenerator{shkContext: shk, signedWeight: signedWt, threshold: threshold}
+	threshold := prepareRejectionSamplingThreshold(choice.signedWeight)
+	return coinGenerator{shkContext: shk, signedWeight: choice.signedWeight, threshold: threshold}
 
 }
 
-func prepareRejectionSamplingValues(signedWeight uint64) (*big.Int, *big.Int) {
+func prepareRejectionSamplingThreshold(signedWeight uint64) *big.Int {
 	// we use rejection sampling in order to have a uniform random coin in [0,signedWeight).
 	// use b bits (b=64) per attempt.
 	// define k = roundDown( 2^b / signedWeight )  implemented as (2^b div signedWeight)
@@ -98,26 +98,26 @@ func prepareRejectionSamplingValues(signedWeight uint64) (*big.Int, *big.Int) {
 	threshold.Div(threshold, signedWt)
 
 	threshold.Mul(threshold, signedWt)
-	return threshold, signedWt
+	return threshold
 }
 
 // getNextCoin returns the next 64bits integer which represents a number between [0,signedWeight)
 func (cg *coinGenerator) getNextCoin() uint64 {
 	// take b bits from the XOF and generate an integer z.
 	// we accept the sample if z < threshold
 	// else, we reject the sample and repeat the process.
-	z := &big.Int{}
+	var randNumFromXof uint64
 	for true {
 		var shakeDigest [8]byte
 		cg.shkContext.Read(shakeDigest[:])
-		randNumFromXof := binary.LittleEndian.Uint64(shakeDigest[:])
+		randNumFromXof = binary.LittleEndian.Uint64(shakeDigest[:])
 
+		z := &big.Int{}
 		z.SetUint64(randNumFromXof)
-		if comp := z.Cmp(cg.threshold); comp == -1 {
+		if z.Cmp(cg.threshold) == -1 {
 			break
 		}
 	}
 
-	z.Mod(z, cg.signedWeight)
-	return z.Uint64()
+	return randNumFromXof % cg.signedWeight
 }