simplify API and structs

5 files changed, 104 insertions, 98 deletions

diff --git a/nizk/commit.go b/nizk/commit.go
index 957e4a9..5634b16 100644
--- a/nizk/commit.go
+++ b/nizk/commit.go
@@ -7,18 +7,20 @@ import (
 
 type Bit struct {
 	id  Bytes
-	bit bool
+	set bool
 	α   *Scalar
 	β   *Scalar
 
-	com *Commitment
+	*Commitment
 	prf *Proof
+
+	*Stage
 }
 
 type Commitment struct {
 	A *Point // g^α
 	B *Point // g^β
-	C *Point // g^(ab)g^(bit)
+	C *Point // g^(ab)g^(set)
 }
 
 // This is a construction of a proof of a statement of the form
@@ -36,51 +38,43 @@ type Proof struct {
 	}
 }
 
-func NewBit(id Bytes, bit bool) *Bit {
+func NewBit(id Bytes, set bool) *Bit {
 	α, β := Curve.RandomScalar(), Curve.RandomScalar()
-	return NewBitFromScalars(id, bit, α, β)
+	return NewBitFromScalars(id, set, α, β)
 }
 
-func NewBitFromScalars(id Bytes, bit bool, α, β *Scalar) *Bit {
+func NewBitFromScalars(id Bytes, set bool, α, β *Scalar) *Bit {
 	return &Bit{
 		id:  id,
-		bit: bit,
+		set: set,
 		α:   α,
 		β:   β,
 	}
 }
 
 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.β
+	return b.set
 }
 
 func (b *Bit) commit() *Commitment {
-	if b.com != nil {
-		return b.com
+	if b.Commitment != nil {
+		return b.Commitment
 	}
 
 	var C *Point
 	c := b.α.Mul(b.β)
 
-	if b.bit {
+	if b.set {
 		C = G.Exp(c.Add(One))
 	} else {
 		C = G.Exp(c)
 	}
-	b.com = &Commitment{
+	b.Commitment = &Commitment{
 		C: C,
 		A: G.Exp(b.α),
 		B: G.Exp(b.β),
 	}
-	return b.com
+	return b.Commitment
 }
 
 func (s *Bit) proof() *Proof {
@@ -95,7 +89,7 @@ func (s *Bit) proof() *Proof {
 	w = Curve.RandomScalar()
 	c := s.commit()
 
-	if s.bit {
+	if s.set {
 		e[0][0] = G.Exp(r1)
 		e[0][1] = c.B.Exp(r1).Mul(G.Exp(w))
 		e[1][0] = G.Exp(r2)
@@ -110,7 +104,7 @@ func (s *Bit) proof() *Proof {
 	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.bit {
+	if s.set {
 		pr.C.Ch[0] = w
 		pr.C.Ch[1] = ch.Sub(w)
 		pr.C.R[0] = r1.Sub(s.α.Mul(pr.C.Ch[0]))
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
 }
 
diff --git a/nizk/stage1_test.go b/nizk/stage1_test.go
index 9c6d957..954f356 100644
--- a/nizk/stage1_test.go
+++ b/nizk/stage1_test.go
@@ -13,10 +13,10 @@ func TestStage1Simple(t *testing.T) {
 	bc1, _ := b1.Commit()
 	bc2, _ := b2.Commit()
 
-	s1, c1 := b1.Stage()
-	s2, c2 := b2.Stage()
-	r1, pr1, _ := s1.RevealStage1() // Note: no Xs.
-	r2, pr2, _ := s2.RevealStage1() // Note: no Xs.
+	c1 := b1.StageCommit()
+	c2 := b2.StageCommit()
+	r1, pr1, _ := b1.RevealStage1() // Note: no Xs.
+	r2, pr2, _ := b2.RevealStage1() // Note: no Xs.
 	if !bc1.VerifyStage1(c1, r1, pr1) {
 		t.Fatal("Could not verify st1 with c1 and pr1, plus=true case")
 	}
@@ -40,10 +40,10 @@ func TestStage1FromScalars(t *testing.T) {
 	bc1, _ := b1.Commit()
 	bc2, _ := b2.Commit()
 
-	s1, c1 := b1.StageFromScalars(r, x)
-	s2, c2 := b2.StageFromScalars(x, r)
-	r1, pr1, _ := s1.RevealStage1() // Note: no Xs
-	r2, pr2, _ := s2.RevealStage1() // Note: no Xs
+	c1 := b1.StageFromScalars(r, x)
+	c2 := b2.StageFromScalars(x, r)
+	r1, pr1, _ := b1.RevealStage1() // Note: no Xs
+	r2, pr2, _ := b2.RevealStage1() // Note: no Xs
 	if !bc1.VerifyStage1(c1, r1, pr1) {
 		t.Fatal("Could not verify st1 with c1 and pr1, plus=true case")
 	}
diff --git a/nizk/stage2.go b/nizk/stage2.go
index 37cd8c1..e6e6b34 100644
--- a/nizk/stage2.go
+++ b/nizk/stage2.go
@@ -1,6 +1,8 @@
 package nizk
 
 import (
+	"fmt"
+
 	. "kesim.org/seal/common"
 )
 
@@ -19,7 +21,11 @@ type Stage2Proof struct {
 	R3 [2]*Scalar
 }
 
-func (s *Stage) RevealStage2(prev_true bool, prev *Stage, Xs ...*Point) (rv2 *StageReveal, pr *Stage2Proof, e error) {
+func (b *Bit) RevealStage2(prev *Bit, Xs ...*Point) (rv2 *StageReveal, pr *Stage2Proof, e error) {
+	s := b.Stage
+	if s == nil {
+		return nil, nil, fmt.Errorf("stage not ready")
+	}
 	var (
 		ε1, ε1_ [3]Bytes
 		ε2, ε2_ [3]Bytes
@@ -36,31 +42,30 @@ func (s *Stage) RevealStage2(prev_true bool, prev *Stage, Xs ...*Point) (rv2 *St
 		}
 	}
 
-	bc := s.bit.com
-	c1 := prev.com
-	c2 := s.com
-	rv1 := prev.rev
-	rv2, e = s.reveal(prev_true, Xs...)
+	c1 := prev.StageCommitment
+	c2 := s.StageCommitment
+	rv1 := prev.StageReveal
+	rv2, e = b.reveal(prev.IsSet(), Xs...)
 	if e != nil {
 		return nil, nil, e
 	}
 
-	if !prev_true {
+	if !prev.IsSet() {
 		ε1[0] = G.Exp(ρ1[0]).Mul(c2.X.Exp(ω[0]))
 		ε1[1] = G.Exp(ρ1[1]).Mul(c1.X.Exp(ω[0]))
-		ε1[2] = G.Exp(ρ1[2]).Mul(bc.A.Exp(ω[0]))
+		ε1[2] = G.Exp(ρ1[2]).Mul(b.A.Exp(ω[0]))
 
 		ε1_[0] = c2.R.Exp(ρ1[0]).Mul(rv2.Z.Exp(ω[0]))
 		ε1_[1] = c1.R.Exp(ρ1[1]).Mul(rv1.Z.Exp(ω[0]))
-		ε1_[2] = bc.B.Exp(ρ1[2]).Mul(bc.C.Div(G).Exp(ω[0]))
+		ε1_[2] = b.B.Exp(ρ1[2]).Mul(b.C.Div(G).Exp(ω[0]))
 
 		ε2[0] = G.Exp(ρ2[0]).Mul(c2.X.Exp(ω[1]))
 		ε2[1] = G.Exp(ρ2[1]).Mul(c1.X.Exp(ω[1]))
-		ε2[2] = G.Exp(ρ2[2]).Mul(bc.A.Exp(ω[1]))
+		ε2[2] = G.Exp(ρ2[2]).Mul(b.A.Exp(ω[1]))
 
 		ε2_[0] = rv2.Y.Exp(ρ2[0]).Mul(rv2.Z.Exp(ω[1]))
 		ε2_[1] = c1.R.Exp(ρ2[1]).Mul(rv1.Z.Exp(ω[1]))
-		ε2_[2] = bc.B.Exp(ρ2[2]).Mul(bc.C.Exp(ω[1]))
+		ε2_[2] = b.B.Exp(ρ2[2]).Mul(b.C.Exp(ω[1]))
 
 		ε3[0] = G.Exp(ρ3[0])
 		ε3[1] = G.Exp(ρ3[1])
@@ -68,22 +73,22 @@ func (s *Stage) RevealStage2(prev_true bool, prev *Stage, Xs ...*Point) (rv2 *St
 		ε3_[0] = rv2.Y.Exp(ρ3[0])
 		ε3_[1] = rv1.Y.Exp(ρ3[1])
 	} else {
-		if s.bit.IsSet() {
+		if b.IsSet() {
 			ε1[0] = G.Exp(ρ1[0])
 			ε1[1] = G.Exp(ρ1[1])
 			ε1[2] = G.Exp(ρ1[2])
 
 			ε1_[0] = c2.R.Exp(ρ1[0])
 			ε1_[1] = c1.R.Exp(ρ1[1])
-			ε1_[2] = bc.B.Exp(ρ1[2])
+			ε1_[2] = b.B.Exp(ρ1[2])
 
 			ε2[0] = G.Exp(ρ2[0]).Mul(c2.X.Exp(ω[0]))
 			ε2[1] = G.Exp(ρ2[1]).Mul(c1.X.Exp(ω[0]))
-			ε2[2] = G.Exp(ρ2[2]).Mul(bc.A.Exp(ω[0]))
+			ε2[2] = G.Exp(ρ2[2]).Mul(b.A.Exp(ω[0]))
 
 			ε2_[0] = rv2.Y.Exp(ρ2[0]).Mul(rv2.Z.Exp(ω[0]))
 			ε2_[1] = c1.R.Exp(ρ2[1]).Mul(rv1.Z.Exp(ω[0]))
-			ε2_[2] = bc.B.Exp(ρ2[2]).Mul(bc.C.Exp(ω[0]))
+			ε2_[2] = b.B.Exp(ρ2[2]).Mul(b.C.Exp(ω[0]))
 
 			ε3[0] = G.Exp(ρ3[0]).Mul(c2.X.Exp(ω[1]))
 			ε3[1] = G.Exp(ρ3[1]).Mul(c1.X.Exp(ω[1]))
@@ -93,11 +98,11 @@ func (s *Stage) RevealStage2(prev_true bool, prev *Stage, Xs ...*Point) (rv2 *St
 		} else {
 			ε1[0] = G.Exp(ρ1[0]).Mul(c2.X.Exp(ω[0]))
 			ε1[1] = G.Exp(ρ1[1]).Mul(c1.X.Exp(ω[0]))
-			ε1[2] = G.Exp(ρ1[2]).Mul(bc.A.Exp(ω[0]))
+			ε1[2] = G.Exp(ρ1[2]).Mul(b.A.Exp(ω[0]))
 
 			ε1_[0] = c2.R.Exp(ρ1[0]).Mul(rv2.Z.Exp(ω[0]))
 			ε1_[1] = c1.R.Exp(ρ1[1]).Mul(rv1.Z.Exp(ω[0]))
-			ε1_[2] = bc.B.Exp(ρ1[2]).Mul(bc.C.Div(G).Exp(ω[0]))
+			ε1_[2] = b.B.Exp(ρ1[2]).Mul(b.C.Div(G).Exp(ω[0]))
 
 			ε2[0] = G.Exp(ρ2[0])
 			ε2[1] = G.Exp(ρ2[1])
@@ -105,7 +110,7 @@ func (s *Stage) RevealStage2(prev_true bool, prev *Stage, Xs ...*Point) (rv2 *St
 
 			ε2_[0] = rv2.Y.Exp(ρ2[0])
 			ε2_[1] = c1.R.Exp(ρ2[1])
-			ε2_[2] = bc.B.Exp(ρ2[2])
+			ε2_[2] = b.B.Exp(ρ2[2])
 
 			ε3[0] = G.Exp(ρ3[0]).Mul(c2.X.Exp(ω[1]))
 			ε3[1] = G.Exp(ρ3[1]).Mul(c1.X.Exp(ω[1]))
@@ -115,7 +120,7 @@ func (s *Stage) RevealStage2(prev_true bool, prev *Stage, Xs ...*Point) (rv2 *St
 		}
 	}
 
-	points := []Bytes{G, bc.A, bc.B, bc.C, c2.R, c2.X, rv2.Y, rv2.Z, c1.R, c1.X, rv1.Y, rv1.Z}
+	points := []Bytes{G, b.A, b.B, b.C, c2.R, c2.X, rv2.Y, rv2.Z, c1.R, c1.X, rv1.Y, rv1.Z}
 	points = append(points, ε1[:]...)
 	points = append(points, ε2[:]...)
 	points = append(points, ε3[:]...)
@@ -126,7 +131,7 @@ func (s *Stage) RevealStage2(prev_true bool, prev *Stage, Xs ...*Point) (rv2 *St
 	ch := Challenge(points...)
 	pr = &Stage2Proof{}
 
-	if !prev_true {
+	if !prev.IsSet() {
 		pr.Ch[0] = ω[0]
 		pr.Ch[1] = ω[1]
 		pr.Ch[2] = ch.Sub(ω[0]).Sub(ω[1])
@@ -142,14 +147,14 @@ func (s *Stage) RevealStage2(prev_true bool, prev *Stage, Xs ...*Point) (rv2 *St
 		pr.R3[0] = ρ3[0].Sub(s.x.Mul(pr.Ch[2]))
 		pr.R3[1] = ρ3[1].Sub(prev.x.Mul(pr.Ch[2]))
 	} else {
-		if s.bit.IsSet() {
+		if b.IsSet() {
 			pr.Ch[0] = ch.Sub(ω[0]).Sub(ω[1])
 			pr.Ch[1] = ω[0]
 			pr.Ch[2] = ω[1]
 
 			pr.R1[0] = ρ1[0].Sub(s.x.Mul(pr.Ch[0]))
 			pr.R1[1] = ρ1[1].Sub(prev.x.Mul(pr.Ch[0]))
-			pr.R1[2] = ρ1[2].Sub(s.bit.α.Mul(pr.Ch[0]))
+			pr.R1[2] = ρ1[2].Sub(b.α.Mul(pr.Ch[0]))
 
 			pr.R2[0] = ρ2[0]
 			pr.R2[1] = ρ2[1]
@@ -168,7 +173,7 @@ func (s *Stage) RevealStage2(prev_true bool, prev *Stage, Xs ...*Point) (rv2 *St
 
 			pr.R2[0] = ρ2[0].Sub(s.x.Mul(pr.Ch[1]))
 			pr.R2[1] = ρ2[1].Sub(prev.x.Mul(pr.Ch[1]))
-			pr.R2[2] = ρ2[2].Sub(s.bit.α.Mul(pr.Ch[1]))
+			pr.R2[2] = ρ2[2].Sub(b.α.Mul(pr.Ch[1]))
 
 			pr.R3[0] = ρ3[0]
 			pr.R3[1] = ρ3[1]
diff --git a/nizk/stage2_test.go b/nizk/stage2_test.go
index c0a5656..e5b44c6 100644
--- a/nizk/stage2_test.go
+++ b/nizk/stage2_test.go
@@ -8,16 +8,15 @@ import (
 
 func TestStage2Simple(t *testing.T) {
 	id1 := Curve.RandomScalar()
-	id2 := Curve.RandomScalar()
-	b1 := NewBit(id1, true)
-	b2 := NewBit(id2, false)
+	b1 := NewBit(id1, false)
+	b2 := NewBit(id1, true)
 
-	s1, c1 := b1.Stage()
-	r1, _, _ := s1.RevealStage1()
+	c1 := b1.StageCommit()
+	r1, _, _ := b1.RevealStage1()
 
 	bc2, _ := b2.Commit()
-	s2, c2 := b2.Stage()
-	r2, p2, e := s2.RevealStage2(true, s1)
+	c2 := b2.StageCommit()
+	r2, p2, e := b2.RevealStage2(b1)
 	if e != nil {
 		t.Fatalf("e: %v", e)
 	}