refactor bit and stages: common stages type for both statges

4 files changed, 58 insertions, 56 deletions

diff --git a/auction.go b/auction.go
index cf46acc..a4e9272 100644
--- a/auction.go
+++ b/auction.go
@@ -117,7 +117,7 @@ type auction struct {
 	bidder_ids []string
 
 	// Stage 1 data per round
-	stage1 []*nizk.Stage1
+	stage1 []*nizk.Stage
 
 	log *slog.Logger
 }
diff --git a/nizk/stage1.go b/nizk/stage1.go
index dd4a896..228970b 100644
--- a/nizk/stage1.go
+++ b/nizk/stage1.go
@@ -2,18 +2,19 @@ package nizk
 
 import . "kesim.org/seal/common"
 
-type Stage1 struct {
+type Stage struct {
 	x *Scalar
 	y *Scalar
 	r *Scalar
 
-	com *Stage1Commitment
-	prf *Stage1Proof
+	com  *StageCommitment
+	prf1 *Stage1Proof
+	prf2 *Stage2Proof
 
 	bit *Bit
 }
 
-type Stage1Commitment struct {
+type StageCommitment struct {
 	R *Point
 	X *Point
 	Y *Point
@@ -31,16 +32,8 @@ type Stage1Proof struct {
 	Rho [2][2]*Scalar
 }
 
-func (b *Bit) Stage1() *Stage1 {
-	var x [3]*Scalar
-	for i := range x {
-		x[i] = Curve.RandomScalar()
-	}
-	return b.Stage1FromScalars(x[0], x[1], x[2])
-}
-
-func (b *Bit) Stage1FromScalars(x, y, r *Scalar) *Stage1 {
-	return &Stage1{
+func (b *Bit) stage(x, y, r *Scalar) *Stage {
+	return &Stage{
 		x: x,
 		y: y,
 		r: r,
@@ -49,7 +42,20 @@ func (b *Bit) Stage1FromScalars(x, y, r *Scalar) *Stage1 {
 	}
 }
 
-func (s *Stage1) commit() *Stage1Commitment {
+func (b *Bit) CommitStage1() (c *Stage, s *StageCommitment, p *Stage1Proof) {
+	var x [3]*Scalar
+	for i := range x {
+		x[i] = Curve.RandomScalar()
+	}
+	return b.CommitStage1FromScalars(x[0], x[1], x[2])
+}
+
+func (b *Bit) CommitStage1FromScalars(x, y, r *Scalar) (s *Stage, c *StageCommitment, p *Stage1Proof) {
+	s = b.stage(x, y, r)
+	return s, s.commit1(), s.proof1()
+}
+
+func (s *Stage) commit1() *StageCommitment {
 	if s.com != nil {
 		return s.com
 	}
@@ -60,7 +66,7 @@ func (s *Stage1) commit() *Stage1Commitment {
 		Z = G.Exp(s.x.Mul(s.y))
 	}
 
-	s.com = &Stage1Commitment{
+	s.com = &StageCommitment{
 		Z: Z,
 		X: G.Exp(s.x),
 		Y: G.Exp(s.y),
@@ -69,13 +75,13 @@ func (s *Stage1) commit() *Stage1Commitment {
 	return s.com
 }
 
-func (s *Stage1) proof() *Stage1Proof {
+func (s *Stage) proof1() *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()
+	c := s.commit1()
 	bc, _ := s.bit.Commit()
 
 	if s.bit.IsSet() {
@@ -126,26 +132,23 @@ func (s *Stage1) proof() *Stage1Proof {
 		pr.Rho[1][1] = ρ2
 	}
 
+	s.prf1 = pr
 	return pr
 }
 
-func (s *Stage1) Commit() (*Stage1Commitment, *Stage1Proof) {
-	return s.commit(), s.proof()
-}
-
-func (c1 *Stage1Commitment) Verify(c *Commitment, p *Stage1Proof) bool {
+func (c *Commitment) VerifyStage1(sc *StageCommitment, p *Stage1Proof) bool {
 	var ε [2][4]*Point
 
-	ε[0][0] = G.Exp(p.Rho[0][0]).Mul(c1.X.Exp(p.Ch[0]))
+	ε[0][0] = G.Exp(p.Rho[0][0]).Mul(sc.X.Exp(p.Ch[0]))
 	ε[0][1] = G.Exp(p.Rho[0][1]).Mul(c.A.Exp(p.Ch[0]))
-	ε[0][2] = c1.Y.Exp(p.Rho[0][0]).Mul(c1.Z.Exp(p.Ch[0]))
+	ε[0][2] = sc.Y.Exp(p.Rho[0][0]).Mul(sc.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(c1.X.Exp(p.Ch[1]))
+	ε[1][0] = G.Exp(p.Rho[1][0]).Mul(sc.X.Exp(p.Ch[1]))
 	ε[1][1] = G.Exp(p.Rho[1][1]).Mul(c.A.Exp(p.Ch[1]))
-	ε[1][2] = c1.R.Exp(p.Rho[1][0]).Mul(c1.Z.Exp(p.Ch[1]))
+	ε[1][2] = sc.R.Exp(p.Rho[1][0]).Mul(sc.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, c1.R, c1.X, c1.Y, c1.Z}
+	points := []Bytes{G, c.A, c.B, c.C, sc.R, sc.X, sc.Y, sc.Z}
 	for _, e := range ε[0] {
 		points = append(points, e)
 	}
diff --git a/nizk/stage1_test.go b/nizk/stage1_test.go
index a1b7327..516cb12 100644
--- a/nizk/stage1_test.go
+++ b/nizk/stage1_test.go
@@ -12,19 +12,17 @@ func TestStage1(t *testing.T) {
 	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(bc1, pr1) {
+	_, c1, pr1 := b1.CommitStage1()
+	_, c2, pr2 := b2.CommitStage1()
+	if !bc1.VerifyStage1(c1, pr1) {
 		t.Fatal("Could not verify st1 with c1 and pr1, plus=true case")
 	}
-	if !c2.Verify(bc2, pr2) {
+	if !bc2.VerifyStage1(c2, pr2) {
 		t.Fatal("Could not verify st2 with c2 and pr2, plus=false case")
 	}
 	// Wrong proof test
-	if c1.Verify(bc1, pr2) {
+	if bc1.VerifyStage1(c1, pr2) {
 		t.Fatal("Shouldn't be able to verify c1 with pr2")
 	}
 }
@@ -37,21 +35,19 @@ func TestStage1FromScalars(t *testing.T) {
 
 	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(bc1, pr1) {
+
+	_, c1, pr1 := b1.CommitStage1()
+	_, c2, pr2 := b2.CommitStage1()
+	if !bc1.VerifyStage1(c1, pr1) {
 		t.Fatal("Could not verify st1 with c1 and pr1, plus=true case")
 	}
-	if !c2.Verify(bc2, pr2) {
+	if !bc2.VerifyStage1(c2, pr2) {
 		t.Fatal("Could not verify st2 with c2 and pr2, plus=false case")
 	}
 	// Wrong proof test
-	if c1.Verify(bc2, pr2) {
+	if bc1.VerifyStage1(c2, pr2) {
 		t.Fatal("Shouldn't be able to verify c1 with pr2")
 	}
 }
diff --git a/nizk/stage2.go b/nizk/stage2.go
index be3d21f..d791ef8 100644
--- a/nizk/stage2.go
+++ b/nizk/stage2.go
@@ -26,15 +26,18 @@ type Stage2 struct {
 	a   *Scalar
 	b   *Scalar
 
-	// Stage1
+	// Private data from previous stage1 or stage2
+	x_ *Scalar
+	y_ *Scalar
+	r_ *Scalar
+
+	// New stage2 private data
 	r *Scalar
 	x *Scalar
 	y *Scalar
 
-	// Stage2
-	r_ *Scalar
-	x_ *Scalar
-	y_ *Scalar
+	com *Stage2Commitment
+	prf *Stage2Proof
 }
 
 type Stage2Commitment struct {
@@ -43,17 +46,17 @@ type Stage2Commitment struct {
 	B *Point
 	C *Point
 
-	// Stage1Commitment
-	R *Point
-	X *Point
-	Y *Point
-	Z *Point
-
-	// New
+	// Previous Commitment
 	R_ *Point
 	X_ *Point
 	Y_ *Point
 	Z_ *Point
+
+	// Stage2Commitment
+	R *Point
+	X *Point
+	Y *Point
+	Z *Point
 }
 
 func NewStage2(typ Type) *Stage2 {