example from paper runs as test

3 files changed, 107 insertions, 12 deletions

diff --git a/nizk/stage1.go b/nizk/stage1.go
index 453d683..e7ed44d 100644
--- a/nizk/stage1.go
+++ b/nizk/stage1.go
@@ -10,6 +10,7 @@ type Stage struct {
 
 	*StageCommitment
 	*StageReveal
+	Sent bool
 }
 
 type StageCommitment struct {
@@ -94,6 +95,7 @@ func (b *Bit) reveal(prev_true bool, Xs ...*Point) (r *StageReveal) {
 
 	if prev_true && b.IsSet() {
 		r.Z = s.R.Exp(s.x)
+		s.Sent = true
 	} else {
 		r.Z = Y.Exp(s.x)
 	}
diff --git a/nizk/stage2.go b/nizk/stage2.go
index 733a172..d5e8070 100644
--- a/nizk/stage2.go
+++ b/nizk/stage2.go
@@ -43,7 +43,7 @@ func (b *Bit) RevealStage2(lost bool, prev *Bit, Xs ...*Point) (rv2 *StageReveal
 	c1 := prev.StageCommitment
 	c2 := s.StageCommitment
 	rv1 := prev.StageReveal
-	b.StageReveal = b.reveal(prev.IsSet(), Xs...)
+	b.StageReveal = b.reveal(prev.Sent, Xs...)
 	rv2 = b.StageReveal
 
 	if lost {
diff --git a/nizk/stage2_test.go b/nizk/stage2_test.go
index 3aeb894..4e6232e 100644
--- a/nizk/stage2_test.go
+++ b/nizk/stage2_test.go
@@ -49,10 +49,11 @@ func TestStage2Simple1(t *testing.T) {
 	}
 }
 
-func uint2bits(bid int) [4]*Bit {
+func uint2bits(bid int) []*Bit {
 	id := Curve.RandomScalar()
 
-	return [4]*Bit{
+	return []*Bit{
+		NewBit(id, (bid>>4)&1 != 0),
 		NewBit(id, (bid>>3)&1 != 0),
 		NewBit(id, (bid>>2)&1 != 0),
 		NewBit(id, (bid>>1)&1 != 0),
@@ -65,7 +66,7 @@ var Id = Curve.Identity()
 func TestStage2Complex(t *testing.T) {
 	bid1 := 0b0101
 	bid2 := 0b0010
-	t.Logf("testing bid1: %04b vs. bid2: %04b", bid1, bid2)
+	t.Logf("testing bid1: %05b vs. bid2: %05b", bid1, bid2)
 
 	bits1 := uint2bits(bid1)
 	bits2 := uint2bits(bid2)
@@ -73,7 +74,7 @@ func TestStage2Complex(t *testing.T) {
 	lost1 := false
 	lost2 := false
 
-	if len(bits1) != len(bits2) || len(bits1) != 4 {
+	if len(bits1) != len(bits2) || len(bits1) != 5 {
 		t.Fatalf("oops")
 	}
 
@@ -116,7 +117,7 @@ func TestStage2Complex(t *testing.T) {
 					t.Logf("setting lost2 to true")
 					lost2 = true
 				}
-				result |= 1 << (3 - c)
+				result |= 1 << (4 - c)
 			} else {
 				t.Logf("Z[%d] == Id, staying in stage1", c)
 			}
@@ -150,18 +151,110 @@ func TestStage2Complex(t *testing.T) {
 					t.Logf("setting lost2 to true")
 					lost2 = true
 				}
-				result |= 1 << (3 - c)
+				result |= 1 << (4 - c)
 			}
 		}
 	}
 	if result != bid1 {
-		t.Fatalf("wrong result: %04b", result)
+		t.Fatalf("wrong result: %05b", result)
 	}
 }
 
 func TestFromPaper(t *testing.T) {
-	bid1 := 0b01010
-	bid2 := 0b01001
-	bid3 := 0b00111
-	t.Logf("testing\n\tbits1: %04b\n\tbits2: %04b\n\tbits3: %04b", bid1, bid2, bid3)
+	vals := []int{
+		0b01010,
+		0b01001,
+		0b00111,
+	}
+
+	t.Logf("testing bits: %05b, %05b, %05b", vals[0], vals[1], vals[2])
+
+	var bits = [3][]*Bit{}
+	for i, b := range vals {
+		bits[i] = uint2bits(b)
+	}
+
+	var lost = [3]bool{}
+	instage1 := true
+	junction := -1
+	result := 0
+
+	for idx := 0; idx < 5; idx++ {
+		var c = [3]*StageCommitment{}
+		var r = [3]*StageReveal{}
+
+		for i, b := range bits {
+			c[i] = b[idx].StageCommit()
+		}
+
+		if instage1 {
+			t.Logf("bit[%d] b1 = %t vs b2 = %t vs b3 = %t", idx,
+				bits[0][idx].IsSet(), bits[1][idx].IsSet(), bits[2][idx].IsSet())
+
+			var p = [3]*Stage1Proof{}
+			for i := range bits {
+				r[i], p[i] = bits[i][idx].RevealStage1(c[0].X, c[1].X, c[2].X)
+				t.Logf("bits[%d][%d] has sent %t", i, idx, bits[i][idx].Sent)
+				if !bits[i][idx].Commitment.VerifyStage1(c[i], r[i], p[i]) {
+					t.Fatalf("bits[%d][%d] commitment failed to verify in stage1", i, idx)
+				}
+			}
+
+			Z := Curve.Product(r[0].Z, r[1].Z, r[2].Z)
+			if !Id.Equal(Z) {
+				t.Logf("Aha! Z[%d] != Id, switch to stage2", idx)
+				junction = idx
+				instage1 = false
+
+				for i := range bits {
+					if !lost[i] && !bits[i][idx].IsSet() {
+						lost[i] = true
+						t.Logf("bit %d, set lost[%d] to true, so far: %v", idx, i, lost)
+					}
+				}
+				result |= 1 << (4 - idx)
+			} else {
+				t.Logf("Z[%d] == Id, staying in stage1", idx)
+			}
+		} else {
+			t.Logf("\nTesing bit[%d]:\n"+
+				"set  0: %t\t1: %t\t2: %t\n"+
+				"lost 0: %t\t1  %t\t2: %t\n",
+				idx,
+				bits[0][idx].IsSet(), bits[1][idx].IsSet(), bits[2][idx].IsSet(),
+				lost[0], lost[1], lost[2])
+
+			var bj = [3]*Bit{}
+			for i := range bits {
+				bj[i] = bits[i][junction]
+			}
+
+			var p = [3]*Stage2Proof{}
+			for i := range bits {
+				r[i], p[i] = bits[i][idx].RevealStage2(lost[i], bj[i], c[0].X, c[1].X, c[2].X)
+				t.Logf("bits[%d][%d] has sent %t", i, idx, bits[i][idx].Sent)
+				if !bits[i][idx].Commitment.VerifyStage2(bj[i].StageCommitment, c[i], bj[i].StageReveal, r[i], p[i]) {
+					t.Fatalf("bits[%d][%d] commitment failed to verify in stage2, result so far: %05b", i, idx, result)
+				}
+			}
+
+			Z := Curve.Product(r[0].Z, r[1].Z, r[2].Z)
+			if !Id.Equal(Z) {
+				t.Logf("Aha! Z[%d] != Id, new junction!", idx)
+				junction = idx
+
+				for i := range bits {
+					if !lost[i] && !bits[i][idx].IsSet() {
+						lost[i] = true
+						t.Logf("bits[%d][%d], set lost[%d] to true, so far: %v", i, idx, i, lost)
+					}
+				}
+				result |= 1 << (4 - idx)
+			}
+		}
+	}
+	if result != vals[0] {
+		t.Fatalf("wrong result: %05b", result)
+	}
+
 }