1 files changed, 26 insertions, 46 deletions

diff --git a/nizk/bench_test.go b/nizk/bench_test.go
index 9d1cc4b..3ee30a8 100644
--- a/nizk/bench_test.go
+++ b/nizk/bench_test.go
@@ -117,7 +117,7 @@ func runSeal(n int, bitlength int, tb testing.TB) {
 	var bj = make([]*Bit, n)
 	var p2 = make([]*Stage2Proof, n)
 	var lost = make([]bool, n)
-	instage1 := true
+	stage := stage1
 	junction := -1
 	result := 0
 
@@ -125,64 +125,44 @@ func runSeal(n int, bitlength int, tb testing.TB) {
 		for i := range n {
 			c[i] = bits[i][idx].StageCommit()
 			Xs[i] = c[i].X
+			if stage == stage2 {
+				bj[i] = bits[i][junction]
+
+			}
 		}
-		if instage1 {
-			var wg sync.WaitGroup
-			wg.Add(n)
+		var wg sync.WaitGroup
+		wg.Add(n)
 
-			for i := range n {
-				go func() {
+		for i := range n {
+			go func() {
+				defer wg.Done()
+
+				if stage == stage1 {
 					r[i], p1[i] = bits[i][idx].RevealStage1(Xs...)
 					if !bits[i][idx].Commitment.VerifyStage1(c[i], r[i], p1[i]) {
 						tb.Fatalf("bits[%d][%d] commitment failed to verify in stage1", i, idx)
 					}
-					Zs[i] = r[i].Z
-					wg.Done()
-				}()
-			}
-			wg.Wait()
-
-			Z := Curve.Product(Zs...)
-			if !Id.Equal(Z) {
-				junction = idx
-				instage1 = false
-				for i := range bits {
-					if !lost[i] && !bits[i][idx].IsSet() {
-						lost[i] = true
-					}
-				}
-				result |= 1 << (bitlength - 1 - idx)
-			}
-		} else {
-			for i := range bits {
-				bj[i] = bits[i][junction]
-			}
-
-			var wg sync.WaitGroup
-			wg.Add(n)
-			for i := range n {
-				go func() {
+				} else {
 					r[i], p2[i] = bits[i][idx].RevealStage2(lost[i], bj[i], Xs...)
 					if !bits[i][idx].Commitment.VerifyStage2(bj[i].StageCommitment, c[i], bj[i].StageReveal, r[i], p2[i]) {
 						tb.Fatalf("bits[%d][%d] commitment failed to verify in stage2, result so far: %05b", i, idx, result)
 					}
-					Zs[i] = r[i].Z
-					wg.Done()
-				}()
-			}
-			wg.Wait()
-
-			Z := Curve.Product(Zs...)
-			if !Id.Equal(Z) {
-				junction = idx
+				}
+				Zs[i] = r[i].Z
+			}()
+		}
+		wg.Wait()
 
-				for i := range n {
-					if !lost[i] && !bits[i][idx].IsSet() {
-						lost[i] = true
-					}
+		Z := Curve.Product(Zs...)
+		if !Id.Equal(Z) {
+			junction = idx
+			stage = stage2
+			for i := range bits {
+				if !lost[i] && !bits[i][idx].IsSet() {
+					lost[i] = true
 				}
-				result |= 1 << (bitlength - 1 - idx)
 			}
+			result |= 1 << (bitlength - 1 - idx)
 		}
 	}
 	if result != max {