refactor: make Bit and Stage1 more composable

1 files changed, 17 insertions, 12 deletions

diff --git a/nizk/stage1_test.go b/nizk/stage1_test.go
index 3a1fac3..a1b7327 100644
--- a/nizk/stage1_test.go
+++ b/nizk/stage1_test.go
@@ -7,46 +7,51 @@ import (
 )
 
 func TestStage1(t *testing.T) {
-	b1 := NewBit(true)
-	b2 := NewBit(false)
+	id := Curve.RandomScalar()
+	b1 := NewBit(id, true)
+	b2 := NewBit(id, false)
+	bc1, _ := b1.Commit()
+	bc2, _ := b2.Commit()
 	st1 := b1.Stage1()
 	st2 := b2.Stage1()
 
 	c1, pr1 := st1.Commit()
 	c2, pr2 := st2.Commit()
-	if !c1.Verify(pr1) {
+	if !c1.Verify(bc1, pr1) {
 		t.Fatal("Could not verify st1 with c1 and pr1, plus=true case")
 	}
-	if !c2.Verify(pr2) {
+	if !c2.Verify(bc2, pr2) {
 		t.Fatal("Could not verify st2 with c2 and pr2, plus=false case")
 	}
 	// Wrong proof test
-	if c1.Verify(pr2) {
+	if c1.Verify(bc1, pr2) {
 		t.Fatal("Shouldn't be able to verify c1 with pr2")
 	}
 }
 
 func TestStage1FromScalars(t *testing.T) {
-	var x, y, r, α, β *Scalar
-	for _, s := range []**Scalar{&x, &y, &r, &α, &β} {
+	var id, x, y, r, α, β *Scalar
+	for _, s := range []**Scalar{&id, &x, &y, &r, &α, &β} {
 		*s = Curve.RandomScalar()
 	}
 
-	b1 := NewBitFromScalars(true, α, β)
-	b2 := NewBitFromScalars(false, α, β)
+	b1 := NewBitFromScalars(id, true, α, β)
+	b2 := NewBitFromScalars(id, false, α, β)
 	st1 := b1.Stage1FromScalars(x, y, r)
 	st2 := b2.Stage1FromScalars(x, y, r)
 
+	bc1, _ := b1.Commit()
+	bc2, _ := b2.Commit()
 	c1, pr1 := st1.Commit()
 	c2, pr2 := st2.Commit()
-	if !c1.Verify(pr1) {
+	if !c1.Verify(bc1, pr1) {
 		t.Fatal("Could not verify st1 with c1 and pr1, plus=true case")
 	}
-	if !c2.Verify(pr2) {
+	if !c2.Verify(bc2, pr2) {
 		t.Fatal("Could not verify st2 with c2 and pr2, plus=false case")
 	}
 	// Wrong proof test
-	if c1.Verify(pr2) {
+	if c1.Verify(bc2, pr2) {
 		t.Fatal("Shouldn't be able to verify c1 with pr2")
 	}
 }