From 86146ddad7af24a65d9827a71662c489b21cd66d Mon Sep 17 00:00:00 2001
From: Sean Bowe <ewillbefull@gmail.com>
Date: Sat, 26 Dec 2015 10:18:44 -0700
Subject: [PATCH] Rearranged tests, added XOR example.

---
 src/main.rs           | 17 +++++++++---
 tinysnark/src/lib.rs  | 63 ++++++++++++++++++++++++++++++++++---------
 tinysnark/src/r1cs.rs |  2 +-
 3 files changed, 65 insertions(+), 17 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index fbea99e..1528dc0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,11 +6,20 @@ use tinysnark::{Proof, Keypair, FieldT, LinearTerm, ConstraintSystem};
 fn main() {
     tinysnark::init();
 
-    let mut cs = ConstraintSystem::new(1, 2);
+    let mut cs = ConstraintSystem::new(2, 1);
+    // xor
+    // (2*b) * c = b+c - a
     cs.add_constraint(
-        &[LinearTerm{coeff: FieldT::one(), index: 2}],
+        &[LinearTerm{coeff: FieldT::from(2), index: 2}],
         &[LinearTerm{coeff: FieldT::one(), index: 3}],
-        &[LinearTerm{coeff: FieldT::one(), index: 1}]
+        &[LinearTerm{coeff: FieldT::one(), index: 2},
+          LinearTerm{coeff: FieldT::one(), index: 3},
+          LinearTerm{coeff: -FieldT::one(), index: 1}]
     );
-    assert!(cs.test(&[100.into()], &[10.into(), 10.into()]));
+    let prompt = [0.into(), 1.into()];
+    let solution = [1.into()];
+    assert!(cs.test(&prompt, &solution));
+    let kp = Keypair::new(&cs);
+    let proof = Proof::new(&kp, &prompt, &solution);
+    assert!(proof.verify(&kp, &prompt));
 }
diff --git a/tinysnark/src/lib.rs b/tinysnark/src/lib.rs
index 39d9fb9..00a5582 100644
--- a/tinysnark/src/lib.rs
+++ b/tinysnark/src/lib.rs
@@ -37,6 +37,24 @@ mod tests {
 
     #[test]
     fn test_zk() {
+        fn test_cs_and_prove<N: Into<FieldT> + Copy>(cs: &ConstraintSystem, primary: &[N], aux: &[N]) -> bool
+        {
+            let primary: Vec<FieldT> = primary.iter().map(|n| (*n).into()).collect();
+            let aux: Vec<FieldT> = aux.iter().map(|n| (*n).into()).collect();
+
+            if !cs.test(&primary, &aux) {
+                return false;
+            }
+
+            let kp = Keypair::new(cs);
+            let proof = Proof::new(&kp, &primary, &aux);
+            // If we construct a proof, it should be impossible
+            // that it doesn't verify.
+            assert!(proof.verify(&kp, &primary));
+
+            return true;
+        }
+
         init();
         {
             let mut cs = ConstraintSystem::new(1, 2);
@@ -46,28 +64,49 @@ mod tests {
                 &[LinearTerm{coeff: FieldT::one(), index: 3}],
                 &[LinearTerm{coeff: FieldT::one(), index: 1}]
             );
-            assert!(cs.test(&[10.into()], &[5.into(), 2.into()]));
-            assert!(!cs.test(&[10.into()], &[6.into(), 2.into()]));
 
-            let kp = Keypair::new(&cs);
-            let proof = Proof::new(&kp, &[10.into()], &[5.into(), 2.into()]);
-            assert!(proof.verify(&kp, &[10.into()]));
+            assert!(test_cs_and_prove(&cs, &[1], &[1, 1]));
+            assert!(test_cs_and_prove(&cs, &[0], &[0, 1]));
+            assert!(test_cs_and_prove(&cs, &[10], &[5, 2]));
+            assert!(!test_cs_and_prove(&cs, &[10], &[6, 2]));
         }
         {
             let mut cs = ConstraintSystem::new(0, 1);
             // simple boolean constraint
+            // (1-x) * x = 0
             cs.add_constraint(
-                &[LinearTerm{coeff: FieldT::one(), index: 0}, LinearTerm{coeff: -FieldT::one(), index: 1}],
+                &[LinearTerm{coeff: FieldT::one(), index: 0},
+                  LinearTerm{coeff: -FieldT::one(), index: 1}],
                 &[LinearTerm{coeff: FieldT::one(), index: 1}],
                 &[LinearTerm{coeff: FieldT::zero(), index: 0}]
             );
-            assert!(cs.test(&[], &[1.into()]));
-            assert!(cs.test(&[], &[0.into()]));
-            assert!(!cs.test(&[], &[2.into()]));
 
-            let kp = Keypair::new(&cs);
-            let proof = Proof::new(&kp, &[], &[1.into()]);
-            assert!(proof.verify(&kp, &[]));
+            assert!(test_cs_and_prove(&cs, &[], &[0]));
+            assert!(test_cs_and_prove(&cs, &[], &[1]));
+            assert!(!test_cs_and_prove(&cs, &[], &[2]));
+        }
+        {
+            let mut cs = ConstraintSystem::new(2, 1);
+            // boolean + xor
+            cs.add_constraint(
+                &[LinearTerm{coeff: FieldT::one(), index: 0},
+                  LinearTerm{coeff: -FieldT::one(), index: 3}],
+                &[LinearTerm{coeff: FieldT::one(), index: 3}],
+                &[LinearTerm{coeff: FieldT::zero(), index: 0}]
+            );
+            cs.add_constraint(
+                &[LinearTerm{coeff: FieldT::from(2), index: 2}],
+                &[LinearTerm{coeff: FieldT::one(), index: 3}],
+                &[LinearTerm{coeff: FieldT::one(), index: 2},
+                  LinearTerm{coeff: FieldT::one(), index: 3},
+                  LinearTerm{coeff: -FieldT::one(), index: 1}]
+            );
+
+            assert!(test_cs_and_prove(&cs, &[0, 0], &[0]));
+            assert!(test_cs_and_prove(&cs, &[1, 1], &[0]));
+            assert!(test_cs_and_prove(&cs, &[1, 0], &[1]));
+            assert!(test_cs_and_prove(&cs, &[0, 1], &[1]));
+            assert!(!test_cs_and_prove(&cs, &[0, 1], &[100]));
         }
     }
 
diff --git a/tinysnark/src/r1cs.rs b/tinysnark/src/r1cs.rs
index f7db5b4..fffa77e 100644
--- a/tinysnark/src/r1cs.rs
+++ b/tinysnark/src/r1cs.rs
@@ -60,7 +60,7 @@ impl ConstraintSystem {
         }
     }
 
-    pub fn test(&mut self, primary: &[FieldT], aux: &[FieldT]) -> bool
+    pub fn test(&self, primary: &[FieldT], aux: &[FieldT]) -> bool
     {
         assert_eq!(primary.len(), self.primary_size);
         assert_eq!(aux.len(), self.aux_size);