2 files changed, 46 insertions, 14 deletions
diff --git a/nizk/stage1/stage1.go b/nizk/stage1/stage1.go
index 0bd3e29..5c729c8 100644
--- a/nizk/stage1/stage1.go
+++ b/nizk/stage1/stage1.go
@@ -10,12 +10,12 @@ import (
 // for given Z, X, Y, R, C, A and B
 
 type Statement struct {
-	x    *Scalar
-	y    *Scalar
-	r    *Scalar
-	α    *Scalar
-	β    *Scalar
-	plus bool
+	x      *Scalar
+	y      *Scalar
+	r      *Scalar
+	α      *Scalar
+	β      *Scalar
+	bitSet bool
 	*Commitment
 }
 
@@ -29,22 +29,30 @@ type Commitment struct {
 	Z *Point
 }
 
-func NewStatement(x, y, r, α, β *Scalar, plus bool) *Statement {
+func NewStatement(bitSet bool) *Statement {
+	var x [5]*Scalar
+	for i := range x {
+		x[i] = Curve.RandomScalar()
+	}
+	return NewStatementFromScalars(bitSet, x[0], x[1], x[2], x[3], x[4])
+}
+
+func NewStatementFromScalars(bitSet bool, x, y, r, α, β *Scalar) *Statement {
 	return &Statement{
 		x:          x,
 		y:          y,
 		r:          r,
 		α:          α,
 		β:          β,
-		plus:       plus,
-		Commitment: commitment(x, y, r, α, β, plus),
+		bitSet:     bitSet,
+		Commitment: commitment(x, y, r, α, β, bitSet),
 	}
 }
 
-func commitment(x, y, r, α, β *Scalar, plus bool) *Commitment {
+func commitment(x, y, r, α, β *Scalar, bitSet bool) *Commitment {
 	var Z *Point
 	φ := α.Mul(β)
-	if plus {
+	if bitSet {
 		Z = G.Exp(x.Mul(r))
 		φ = φ.Add(One)
 	} else {
@@ -78,7 +86,7 @@ func (s *Statement) Proof() *Proof {
 		*s = Curve.RandomScalar()
 	}
 
-	if s.plus {
+	if s.bitSet {
 		ε[0][0] = G.Exp(r1).Mul(s.X.Exp(ω))
 		ε[0][1] = G.Exp(r2).Mul(s.A.Exp(ω))
 		ε[0][2] = s.Y.Exp(r1).Mul(s.Z.Exp(ω))
@@ -109,7 +117,7 @@ func (s *Statement) Proof() *Proof {
 	ch := Challenge(p...)
 	pr := &Proof{}
 
-	if s.plus {
+	if s.bitSet {
 		pr.Ch[0] = ω
 		pr.Ch[1] = ch.Sub(ω)
 		pr.Rho[0][0] = r1
diff --git a/nizk/stage1/stage1_test.go b/nizk/stage1/stage1_test.go
index 828dd0a..df93cdb 100644
--- a/nizk/stage1/stage1_test.go
+++ b/nizk/stage1/stage1_test.go
@@ -7,12 +7,32 @@ import (
 )
 
 func TestStatement(t *testing.T) {
+	st1 := NewStatement(true)
+	st2 := NewStatement(false)
+
+	c1, c2 := st1.Commit(), st2.Commit()
+	pr1, pr2 := st1.Proof(), st2.Proof()
+	if !c1.Verify(pr1) {
+		t.Fatal("Could not verify st1 with c1 and pr1, plus=true case")
+	}
+	if !c2.Verify(pr2) {
+		t.Fatal("Could not verify st2 with c2 and pr2, plus=false case")
+	}
+	// Wrong proof test
+	if c1.Verify(pr2) {
+		t.Fatal("Shouldn't be able to verify c1 with pr2")
+	}
+}
+
+func TestStatementFromScalars(t *testing.T) {
 	var x, y, r, α, β *Scalar
 	for _, s := range []**Scalar{&x, &y, &r, &α, &β} {
 		*s = Curve.RandomScalar()
 	}
 
-	st1, st2 := NewStatement(x, y, r, α, β, true), NewStatement(x, y, r, α, β, false)
+	st1 := NewStatementFromScalars(true, x, y, r, α, β)
+	st2 := NewStatementFromScalars(false, x, y, r, α, β)
+
 	c1, c2 := st1.Commit(), st2.Commit()
 	pr1, pr2 := st1.Proof(), st2.Proof()
 	if !c1.Verify(pr1) {
@@ -21,4 +41,8 @@ func TestStatement(t *testing.T) {
 	if !c2.Verify(pr2) {
 		t.Fatal("Could not verify st2 with c2 and pr2, plus=false case")
 	}
+	// Wrong proof test
+	if c1.Verify(pr2) {
+		t.Fatal("Shouldn't be able to verify c1 with pr2")
+	}
 }