simplify API and structs

1 files changed, 45 insertions, 37 deletions

diff --git a/nizk/stage1.go b/nizk/stage1.go
index 4cd9547..3ebed27 100644
--- a/nizk/stage1.go
+++ b/nizk/stage1.go
@@ -1,15 +1,17 @@
 package nizk
 
-import . "kesim.org/seal/common"
+import (
+	"fmt"
 
-type Stage struct {
-	bit *Bit
+	. "kesim.org/seal/common"
+)
 
+type Stage struct {
 	x *Scalar
 	r *Scalar
 
-	com *StageCommitment
-	rev *StageReveal
+	*StageCommitment
+	*StageReveal
 }
 
 type StageCommitment struct {
@@ -33,39 +35,43 @@ type Stage1Proof struct {
 	Rho [2][2]*Scalar
 }
 
-func (b *Bit) stage(x, r *Scalar) *Stage {
+func (b *Bit) stage(x, r *Scalar) {
 	b.Commit() // ensure non-null values for A, B, C
-	return &Stage{
-		x:   x,
-		r:   r,
-		bit: b,
+	b.Stage = &Stage{
+		x: x,
+		r: r,
 	}
 }
 
 func (s *Stage) commit(Xs ...*Point) *StageCommitment {
-	if s.com != nil {
-		return s.com
+	if s.StageCommitment != nil {
+		return s.StageCommitment
 	}
 
-	s.com = &StageCommitment{
+	s.StageCommitment = &StageCommitment{
 		X: G.Exp(s.x),
 		R: G.Exp(s.r),
 	}
-	return s.com
+	return s.StageCommitment
 }
 
-func (b *Bit) Stage(Xs ...*Point) (c *Stage, s *StageCommitment) {
+func (b *Bit) StageCommit(Xs ...*Point) (s *StageCommitment) {
 	x := Curve.RandomScalar()
 	r := Curve.RandomScalar()
 	return b.StageFromScalars(x, r, Xs...)
 }
 
-func (b *Bit) StageFromScalars(x, r *Scalar, Xs ...*Point) (s *Stage, c *StageCommitment) {
-	s = b.stage(x, r)
-	return s, s.commit(Xs...)
+func (b *Bit) StageFromScalars(x, r *Scalar, Xs ...*Point) (c *StageCommitment) {
+	b.stage(x, r)
+	return b.Stage.commit(Xs...)
 }
 
-func (s *Stage) reveal(prev_true bool, Xs ...*Point) (r *StageReveal, e error) {
+func (b *Bit) reveal(prev_true bool, Xs ...*Point) (r *StageReveal, e error) {
+	s := b.Stage
+	if s == nil {
+		return nil, fmt.Errorf("stage not ready")
+	}
+
 	// TODO: Calculate Y based on the Xs and our own X_i
 	// as Π_(i<k) X_k / Π_(i>k) X_k
 	// For now:
@@ -73,8 +79,8 @@ func (s *Stage) reveal(prev_true bool, Xs ...*Point) (r *StageReveal, e error) {
 
 	r = &StageReveal{Y: Y}
 
-	if prev_true && s.bit.IsSet() {
-		r.Z = s.com.R.Exp(s.x)
+	if prev_true && b.IsSet() {
+		r.Z = s.R.Exp(s.x)
 	} else {
 		r.Z = Y.Exp(s.x)
 	}
@@ -82,41 +88,44 @@ func (s *Stage) reveal(prev_true bool, Xs ...*Point) (r *StageReveal, e error) {
 	return r, e
 }
 
-func (s *Stage) RevealStage1(Xs ...*Point) (rev *StageReveal, pr *Stage1Proof, e error) {
+func (b *Bit) RevealStage1(Xs ...*Point) (rev *StageReveal, pr *Stage1Proof, e error) {
+	s := b.Stage
+	if s == nil {
+		return nil, nil, fmt.Errorf("stage not ready")
+	}
 	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.com
 
-	rev, e = s.reveal(true, Xs...)
+	rev, e = b.reveal(true, Xs...)
 	if e != nil {
 		return nil, nil, e
 	}
 
-	if s.bit.IsSet() {
+	if b.IsSet() {
 		ε[0][0] = G.Exp(r1).Mul(c.X.Exp(ω))
-		ε[0][1] = G.Exp(r2).Mul(bc.A.Exp(ω))
+		ε[0][1] = G.Exp(r2).Mul(b.A.Exp(ω))
 		ε[0][2] = rev.Y.Exp(r1).Mul(rev.Z.Exp(ω))
-		ε[0][3] = bc.B.Exp(r2).Mul(bc.C.Exp(ω))
+		ε[0][3] = b.B.Exp(r2).Mul(b.C.Exp(ω))
 		ε[1][0] = G.Exp(ρ1)
 		ε[1][1] = G.Exp(ρ2)
 		ε[1][2] = c.R.Exp(ρ1)
-		ε[1][3] = bc.B.Exp(ρ2)
+		ε[1][3] = b.B.Exp(ρ2)
 	} else {
 		ε[0][0] = G.Exp(r1)
 		ε[0][1] = G.Exp(r2)
 		ε[0][2] = rev.Y.Exp(r1)
-		ε[0][3] = bc.B.Exp(r2)
+		ε[0][3] = b.B.Exp(r2)
 		ε[1][0] = G.Exp(ρ1).Mul(c.X.Exp(ω))
-		ε[1][1] = G.Exp(ρ2).Mul(bc.A.Exp(ω))
+		ε[1][1] = G.Exp(ρ2).Mul(b.A.Exp(ω))
 		ε[1][2] = c.R.Exp(ρ1).Mul(rev.Z.Exp(ω))
-		ε[1][3] = bc.B.Exp(ρ2).Mul(bc.C.Div(G).Exp(ω))
+		ε[1][3] = b.B.Exp(ρ2).Mul(b.C.Div(G).Exp(ω))
 	}
 
-	p := []Bytes{G, bc.A, bc.B, bc.C, c.R, c.X, rev.Y, rev.Z}
+	p := []Bytes{G, b.A, b.B, b.C, c.R, c.X, rev.Y, rev.Z}
 	for _, e := range ε[0] {
 		p = append(p, e)
 	}
@@ -126,25 +135,24 @@ func (s *Stage) RevealStage1(Xs ...*Point) (rev *StageReveal, pr *Stage1Proof, e
 
 	ch := Challenge(p...)
 	pr = &Stage1Proof{}
-	α, _ := s.bit.Scalars()
 
-	if s.bit.IsSet() {
+	if b.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(α.Mul(pr.Ch[1]))
+		pr.Rho[1][1] = ρ2.Sub(b.α.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(α.Mul(pr.Ch[0]))
+		pr.Rho[0][1] = r2.Sub(b.α.Mul(pr.Ch[0]))
 		pr.Rho[1][0] = ρ1
 		pr.Rho[1][1] = ρ2
 	}
 
-	s.rev = rev
+	s.StageReveal = rev
 	return rev, pr, e
 }