refactor: make Bit and Stage1 more composable

4 files changed, 157 insertions, 120 deletions

diff --git a/nizk/commit.go b/nizk/commit.go
index 5b703d9..957e4a9 100644
--- a/nizk/commit.go
+++ b/nizk/commit.go
@@ -5,71 +5,97 @@ import (
 	"kesim.org/seal/nizk/schnorr"
 )
 
-// This is a construction of a proof of a statement of the form
-//    [(C = g^(αβ))   && (A = g^α) && (Β = g^β)]
-// || [(C = g^(αβ+1)) && (A = g^α) && (Β = g^β)]
-//
-// for given C, A and B
-
 type Bit struct {
-	bitSet bool
-	α      *Scalar
-	β      *Scalar
+	id  Bytes
+	bit bool
+	α   *Scalar
+	β   *Scalar
+
+	com *Commitment
+	prf *Proof
 }
 
 type Commitment struct {
 	A *Point // g^α
 	B *Point // g^β
-	C *Point // g^(ab)g^(bitSet)
+	C *Point // g^(ab)g^(bit)
+}
+
+// This is a construction of a proof of a statement of the form
+//
+//	[(C = g^(αβ))   && (A = g^α) && (Β = g^β)]
+//	 || [(C = g^(αβ+1)) && (A = g^α) && (Β = g^β)]
+//
+// for given C, A and B
+type Proof struct {
+	A *schnorr.Proof // Proof for knowledge of α in A = G^α
+	B *schnorr.Proof // Proof for knowledge of β in B = G^β
+	C struct {       // Proof for knowledge of statement above
+		Ch [2]*Scalar
+		R  [2]*Scalar
+	}
 }
 
-func NewBit(bitSet bool) *Bit {
+func NewBit(id Bytes, bit bool) *Bit {
 	α, β := Curve.RandomScalar(), Curve.RandomScalar()
-	return NewBitFromScalars(bitSet, α, β)
+	return NewBitFromScalars(id, bit, α, β)
 }
 
-func NewBitFromScalars(bitSet bool, α, β *Scalar) *Bit {
+func NewBitFromScalars(id Bytes, bit bool, α, β *Scalar) *Bit {
 	return &Bit{
-		α:      α,
-		β:      β,
-		bitSet: bitSet,
+		id:  id,
+		bit: bit,
+		α:   α,
+		β:   β,
 	}
 }
 
-func (b *Bit) commitment() *Commitment {
+func (b *Bit) IsSet() bool {
+	return b.bit
+}
+
+func (b *Bit) Id() Bytes {
+	return b.id
+}
+
+func (b *Bit) Scalars() (α *Scalar, β *Scalar) {
+	return b.α, b.β
+}
+
+func (b *Bit) commit() *Commitment {
+	if b.com != nil {
+		return b.com
+	}
+
 	var C *Point
 	c := b.α.Mul(b.β)
 
-	if b.bitSet {
+	if b.bit {
 		C = G.Exp(c.Add(One))
 	} else {
 		C = G.Exp(c)
 	}
-	return &Commitment{
+	b.com = &Commitment{
 		C: C,
 		A: G.Exp(b.α),
 		B: G.Exp(b.β),
 	}
+	return b.com
 }
 
-type Proof struct {
-	Id Bytes
-	A  *schnorr.Proof // Proof for knowledge of α in A = G^α
-	B  *schnorr.Proof // Proof for knowledge of β in B = G^β
-	C  struct {       // Proof for knowledge of statement above
-		Ch [2]*Scalar
-		R  [2]*Scalar
+func (s *Bit) proof() *Proof {
+	if s.prf != nil {
+		return s.prf
 	}
-}
 
-func (s *Bit) proof(id Bytes, c *Commitment) *Proof {
 	var e [2][2]*Point
 	var r1, r2, w *Scalar
 	r1 = Curve.RandomScalar()
 	r2 = Curve.RandomScalar()
 	w = Curve.RandomScalar()
+	c := s.commit()
 
-	if s.bitSet {
+	if s.bit {
 		e[0][0] = G.Exp(r1)
 		e[0][1] = c.B.Exp(r1).Mul(G.Exp(w))
 		e[1][0] = G.Exp(r2)
@@ -81,10 +107,10 @@ func (s *Bit) proof(id Bytes, c *Commitment) *Proof {
 		e[1][1] = c.B.Exp(r2).Mul(c.C.Div(G).Exp(w))
 	}
 
-	ch := Challenge(G, c.C, c.A, c.B, e[0][0], e[0][1], e[1][0], e[1][1], id)
-	pr := &Proof{Id: id}
+	ch := Challenge(G, c.C, c.A, c.B, e[0][0], e[0][1], e[1][0], e[1][1], s.id)
+	pr := &Proof{}
 
-	if s.bitSet {
+	if s.bit {
 		pr.C.Ch[0] = w
 		pr.C.Ch[1] = ch.Sub(w)
 		pr.C.R[0] = r1.Sub(s.α.Mul(pr.C.Ch[0]))
@@ -95,26 +121,26 @@ func (s *Bit) proof(id Bytes, c *Commitment) *Proof {
 		pr.C.R[0] = r1.Sub(s.α.Mul(pr.C.Ch[0]))
 		pr.C.R[1] = r2
 	}
-	pr.A = (*schnorr.Statement)(s.α).Proof(id)
-	pr.B = (*schnorr.Statement)(s.β).Proof(id)
+	pr.A = (*schnorr.Statement)(s.α).Proof(s.id)
+	pr.B = (*schnorr.Statement)(s.β).Proof(s.id)
 
+	s.prf = pr
 	return pr
 }
 
-func (s *Bit) Commit(id Bytes) (*Commitment, *Proof) {
-	c := s.commitment()
-	return c, s.proof(id, c)
+func (s *Bit) Commit() (*Commitment, *Proof) {
+	return s.commit(), s.proof()
 }
 
-func (c *Commitment) Verify(p *Proof) bool {
+func (c *Commitment) Verify(id Bytes, p *Proof) bool {
 	var e [2][2]*Point
 
 	e[0][0] = G.Exp(p.C.R[0]).Mul(c.A.Exp(p.C.Ch[0]))
 	e[0][1] = c.B.Exp(p.C.R[0]).Mul(c.C.Exp(p.C.Ch[0]))
 	e[1][0] = G.Exp(p.C.R[1]).Mul(c.A.Exp(p.C.Ch[1]))
 	e[1][1] = c.B.Exp(p.C.R[1]).Mul(c.C.Div(G).Exp(p.C.Ch[1]))
-	ch := Challenge(G, c.C, c.A, c.B, e[0][0], e[0][1], e[1][0], e[1][1], p.Id)
+	ch := Challenge(G, c.C, c.A, c.B, e[0][0], e[0][1], e[1][0], e[1][1], id)
 	return p.C.Ch[0].Add(p.C.Ch[1]).Equal(ch) &&
-		(*schnorr.Commitment)(c.A).Verify(p.A, p.Id) &&
-		(*schnorr.Commitment)(c.B).Verify(p.B, p.Id)
+		(*schnorr.Commitment)(c.A).Verify(p.A, id) &&
+		(*schnorr.Commitment)(c.B).Verify(p.B, id)
 }
diff --git a/nizk/commit_test.go b/nizk/commit_test.go
index 061a772..3d65aa4 100644
--- a/nizk/commit_test.go
+++ b/nizk/commit_test.go
@@ -8,41 +8,50 @@ import (
 
 func TestStatement(t *testing.T) {
 	id := Curve.RandomScalar()
-	Id := G.Exp(id)
 
-	st1, st2 := NewBit(true), NewBit(false)
-	c1, p1 := st1.Commit(Id)
-	c2, p2 := st2.Commit(Id)
-	if !c1.Verify(p1) {
+	st1, st2 := NewBit(id, true), NewBit(id, false)
+	c1, p1 := st1.Commit()
+	c2, p2 := st2.Commit()
+	if !c1.Verify(id, p1) {
 		t.Fatal("Could not verify st1 with c1, plus=true case")
 	}
-	if !c2.Verify(p2) {
+	if !c2.Verify(id, p2) {
 		t.Fatal("Could not verify st2 with c2, plus=false case")
 	}
 
 	// Use the wrong proof
-	if c2.Verify(p1) {
+	if c2.Verify(id, p1) {
 		t.Fatal("Verify with wrong proof should have failed!")
 	}
+
+	// Use wrong id
+	x := Curve.RandomScalar()
+	if c1.Verify(x, p1) || c2.Verify(x, p2) {
+		t.Fatal("Verify with wrong id should have failed!")
+	}
 }
 
 func TestStatementFromScalar(t *testing.T) {
 	var α, β, id = Curve.RandomScalar(), Curve.RandomScalar(), Curve.RandomScalar()
 
-	Id := G.Exp(id)
-
-	st1, st2 := NewBitFromScalars(true, α, β), NewBitFromScalars(false, α, β)
-	c1, p1 := st1.Commit(Id)
-	c2, p2 := st2.Commit(Id)
-	if !c1.Verify(p1) {
+	st1, st2 := NewBitFromScalars(id, true, α, β), NewBitFromScalars(id, false, α, β)
+	c1, p1 := st1.Commit()
+	c2, p2 := st2.Commit()
+	if !c1.Verify(id, p1) {
 		t.Fatal("Could not verify st1 with c1, plus=true case")
 	}
-	if !c2.Verify(p2) {
+	if !c2.Verify(id, p2) {
 		t.Fatal("Could not verify st2 with c2, plus=false case")
 	}
 
 	// Use the wrong proof
-	if c2.Verify(p1) {
+	if c2.Verify(id, p1) {
 		t.Fatal("Verify with wrong proof should have failed!")
 	}
+
+	// Use the wrong Id
+	x := Curve.RandomScalar()
+	if c1.Verify(x, p1) || c2.Verify(x, p2) {
+		t.Fatal("Verify with wrong id should have failed!")
+	}
 }
diff --git a/nizk/stage1.go b/nizk/stage1.go
index ff42e7f..dd4a896 100644
--- a/nizk/stage1.go
+++ b/nizk/stage1.go
@@ -2,32 +2,35 @@ package nizk
 
 import . "kesim.org/seal/common"
 
-// Implements the proof and verification of statements of the following form:
-//     [ Z=g^(xy) && X=g^x && Y=g^y && C=g^(αβ)   && A=g^α && B=g^β ]
-//	|| [ Z=g^(xr) && X=g^x && R=g^r && C=g^(αβ+1) && A=g^α && B=g^β ]
-// for given Z, X, Y, R, C, A and B
-
 type Stage1 struct {
-	// Original Bit
-	*Bit
-
-	// New stage 1 private data
 	x *Scalar
 	y *Scalar
 	r *Scalar
+
+	com *Stage1Commitment
+	prf *Stage1Proof
+
+	bit *Bit
 }
 
 type Stage1Commitment struct {
-	// Original Commitment
-	*Commitment
-
-	// New
 	R *Point
 	X *Point
 	Y *Point
 	Z *Point
 }
 
+// Represents the proof of statements of the following form:
+//
+//	    [ Z=g^(xy) && X=g^x && Y=g^y && C=g^(αβ)   && A=g^α && B=g^β ]
+//		|| [ Z=g^(xr) && X=g^x && R=g^r && C=g^(αβ+1) && A=g^α && B=g^β ]
+//
+// for given Z, X, Y, R, C, A and B
+type Stage1Proof struct {
+	Ch  [2]*Scalar
+	Rho [2][2]*Scalar
+}
+
 func (b *Bit) Stage1() *Stage1 {
 	var x [3]*Scalar
 	for i := range x {
@@ -38,70 +41,64 @@ func (b *Bit) Stage1() *Stage1 {
 
 func (b *Bit) Stage1FromScalars(x, y, r *Scalar) *Stage1 {
 	return &Stage1{
-		x:   x,
-		y:   y,
-		r:   r,
-		Bit: b,
+		x: x,
+		y: y,
+		r: r,
+
+		bit: b,
 	}
 }
 
 func (s *Stage1) commit() *Stage1Commitment {
+	if s.com != nil {
+		return s.com
+	}
 	var Z *Point
-	φ := s.α.Mul(s.β)
-	if s.bitSet {
+	if s.bit.IsSet() {
 		Z = G.Exp(s.x.Mul(s.r))
-		φ = φ.Add(One)
 	} else {
 		Z = G.Exp(s.x.Mul(s.y))
 	}
 
-	return &Stage1Commitment{
+	s.com = &Stage1Commitment{
 		Z: Z,
 		X: G.Exp(s.x),
 		Y: G.Exp(s.y),
 		R: G.Exp(s.r),
-
-		Commitment: &Commitment{
-			A: G.Exp(s.α),
-			B: G.Exp(s.β),
-			C: G.Exp(φ),
-		},
 	}
+	return s.com
 }
 
-type Stage1Proof struct {
-	Ch  [2]*Scalar
-	Rho [2][2]*Scalar
-}
-
-func (s *Stage1) proof(c *Stage1Commitment) *Stage1Proof {
+func (s *Stage1) proof() *Stage1Proof {
 	var ε [2][4]*Point
 	var r1, r2, ρ1, ρ2, ω *Scalar
 	for _, s := range []**Scalar{&r1, &r2, &ρ1, &ρ2, &ω} {
 		*s = Curve.RandomScalar()
 	}
+	c := s.commit()
+	bc, _ := s.bit.Commit()
 
-	if s.bitSet {
+	if s.bit.IsSet() {
 		ε[0][0] = G.Exp(r1).Mul(c.X.Exp(ω))
-		ε[0][1] = G.Exp(r2).Mul(c.A.Exp(ω))
+		ε[0][1] = G.Exp(r2).Mul(bc.A.Exp(ω))
 		ε[0][2] = c.Y.Exp(r1).Mul(c.Z.Exp(ω))
-		ε[0][3] = c.B.Exp(r2).Mul(c.C.Exp(ω))
+		ε[0][3] = bc.B.Exp(r2).Mul(bc.C.Exp(ω))
 		ε[1][0] = G.Exp(ρ1)
 		ε[1][1] = G.Exp(ρ2)
 		ε[1][2] = c.R.Exp(ρ1)
-		ε[1][3] = c.B.Exp(ρ2)
+		ε[1][3] = bc.B.Exp(ρ2)
 	} else {
 		ε[0][0] = G.Exp(r1)
 		ε[0][1] = G.Exp(r2)
 		ε[0][2] = c.Y.Exp(r1)
-		ε[0][3] = c.B.Exp(r2)
+		ε[0][3] = bc.B.Exp(r2)
 		ε[1][0] = G.Exp(ρ1).Mul(c.X.Exp(ω))
-		ε[1][1] = G.Exp(ρ2).Mul(c.A.Exp(ω))
+		ε[1][1] = G.Exp(ρ2).Mul(bc.A.Exp(ω))
 		ε[1][2] = c.R.Exp(ρ1).Mul(c.Z.Exp(ω))
-		ε[1][3] = c.B.Exp(ρ2).Mul(c.C.Div(G).Exp(ω))
+		ε[1][3] = bc.B.Exp(ρ2).Mul(bc.C.Div(G).Exp(ω))
 	}
 
-	p := []Bytes{G, c.A, c.B, c.C, c.R, c.X, c.Y, c.Z}
+	p := []Bytes{G, bc.A, bc.B, bc.C, c.R, c.X, c.Y, c.Z}
 	for _, e := range ε[0] {
 		p = append(p, e)
 	}
@@ -111,19 +108,20 @@ func (s *Stage1) proof(c *Stage1Commitment) *Stage1Proof {
 
 	ch := Challenge(p...)
 	pr := &Stage1Proof{}
+	α, _ := s.bit.Scalars()
 
-	if s.bitSet {
+	if s.bit.IsSet() {
 		pr.Ch[0] = ω
 		pr.Ch[1] = ch.Sub(ω)
 		pr.Rho[0][0] = r1
 		pr.Rho[0][1] = r2
 		pr.Rho[1][0] = ρ1.Sub(s.x.Mul(pr.Ch[1]))
-		pr.Rho[1][1] = ρ2.Sub(s.α.Mul(pr.Ch[1]))
+		pr.Rho[1][1] = ρ2.Sub(α.Mul(pr.Ch[1]))
 	} else {
 		pr.Ch[0] = ch.Sub(ω)
 		pr.Ch[1] = ω
 		pr.Rho[0][0] = r1.Sub(s.x.Mul(pr.Ch[0]))
-		pr.Rho[0][1] = r2.Sub(s.α.Mul(pr.Ch[0]))
+		pr.Rho[0][1] = r2.Sub(α.Mul(pr.Ch[0]))
 		pr.Rho[1][0] = ρ1
 		pr.Rho[1][1] = ρ2
 	}
@@ -132,23 +130,22 @@ func (s *Stage1) proof(c *Stage1Commitment) *Stage1Proof {
 }
 
 func (s *Stage1) Commit() (*Stage1Commitment, *Stage1Proof) {
-	c := s.commit()
-	return c, s.proof(c)
+	return s.commit(), s.proof()
 }
 
-func (c *Stage1Commitment) Verify(p *Stage1Proof) bool {
+func (c1 *Stage1Commitment) Verify(c *Commitment, p *Stage1Proof) bool {
 	var ε [2][4]*Point
 
-	ε[0][0] = G.Exp(p.Rho[0][0]).Mul(c.X.Exp(p.Ch[0]))
+	ε[0][0] = G.Exp(p.Rho[0][0]).Mul(c1.X.Exp(p.Ch[0]))
 	ε[0][1] = G.Exp(p.Rho[0][1]).Mul(c.A.Exp(p.Ch[0]))
-	ε[0][2] = c.Y.Exp(p.Rho[0][0]).Mul(c.Z.Exp(p.Ch[0]))
+	ε[0][2] = c1.Y.Exp(p.Rho[0][0]).Mul(c1.Z.Exp(p.Ch[0]))
 	ε[0][3] = c.B.Exp(p.Rho[0][1]).Mul(c.C.Exp(p.Ch[0]))
-	ε[1][0] = G.Exp(p.Rho[1][0]).Mul(c.X.Exp(p.Ch[1]))
+	ε[1][0] = G.Exp(p.Rho[1][0]).Mul(c1.X.Exp(p.Ch[1]))
 	ε[1][1] = G.Exp(p.Rho[1][1]).Mul(c.A.Exp(p.Ch[1]))
-	ε[1][2] = c.R.Exp(p.Rho[1][0]).Mul(c.Z.Exp(p.Ch[1]))
+	ε[1][2] = c1.R.Exp(p.Rho[1][0]).Mul(c1.Z.Exp(p.Ch[1]))
 	ε[1][3] = c.B.Exp(p.Rho[1][1]).Mul(c.C.Div(G).Exp(p.Ch[1]))
 
-	points := []Bytes{G, c.A, c.B, c.C, c.R, c.X, c.Y, c.Z}
+	points := []Bytes{G, c.A, c.B, c.C, c1.R, c1.X, c1.Y, c1.Z}
 	for _, e := range ε[0] {
 		points = append(points, e)
 	}
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")
 	}
 }