From f4e030793c60890abebcbd3add6b3e2288d39dfe Mon Sep 17 00:00:00 2001
From: Georgios Konstantopoulos <me@gakonst.com>
Date: Mon, 23 Mar 2020 11:07:06 +0200
Subject: [PATCH] test: ensure the Counter and the System match

---
 r1cs-std/src/test_constraint_counter.rs | 62 +++++++++++++++++++++++--
 1 file changed, 57 insertions(+), 5 deletions(-)

diff --git a/r1cs-std/src/test_constraint_counter.rs b/r1cs-std/src/test_constraint_counter.rs
index 78e6883..fc42157 100644
--- a/r1cs-std/src/test_constraint_counter.rs
+++ b/r1cs-std/src/test_constraint_counter.rs
@@ -4,8 +4,8 @@ use r1cs_core::{ConstraintSystem, Index, LinearCombination, SynthesisError, Vari
 
 /// Constraint counter for testing purposes.
 pub struct TestConstraintCounter {
-    pub num_inputs:      usize,
-    pub num_aux:         usize,
+    pub num_inputs: usize,
+    pub num_aux: usize,
     pub num_constraints: usize,
 }
 
@@ -64,9 +64,10 @@ impl<ConstraintF: Field> ConstraintSystem<ConstraintF> for TestConstraintCounter
     where
         NR: Into<String>,
         N: FnOnce() -> NR,
-    {}
+    {
+    }
 
-    fn pop_namespace(&mut self) { }
+    fn pop_namespace(&mut self) {}
 
     fn get_root(&mut self) -> &mut Self::Root {
         self
@@ -75,4 +76,55 @@ impl<ConstraintF: Field> ConstraintSystem<ConstraintF> for TestConstraintCounter
     fn num_constraints(&self) -> usize {
         self.num_constraints
     }
-}
\ No newline at end of file
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::test_constraint_system::TestConstraintSystem;
+    use algebra::{bls12_381::Fq, Field};
+    use r1cs_core::{ConstraintSynthesizer, SynthesisError};
+
+    // circuit proving knowledge of a square root
+    #[derive(Clone, Debug)]
+    struct TestCircuit<F>(Option<F>);
+
+    impl<F: Field> ConstraintSynthesizer<F> for TestCircuit<F> {
+        fn generate_constraints<CS: ConstraintSystem<F>>(
+            self,
+            cs: &mut CS,
+        ) -> Result<(), SynthesisError> {
+            let x = cs.alloc(|| "x", || self.0.ok_or(SynthesisError::AssignmentMissing))?;
+            // 1 input!
+            let out = cs.alloc_input(
+                || "square",
+                || {
+                    self.0
+                        .map(|x| x.square())
+                        .ok_or(SynthesisError::AssignmentMissing)
+                },
+            )?;
+            cs.enforce(|| "x * x = x^2", |lc| lc + x, |lc| lc + x, |lc| lc + out);
+            Ok(())
+        }
+    }
+
+    #[test]
+    fn test_constraints_counter() {
+        let empty_circuit = TestCircuit::<Fq>(None);
+        let populated_circuit = TestCircuit(Some(Fq::from(10u32)));
+
+        let mut counter = TestConstraintCounter::new();
+        let mut cs = TestConstraintSystem::new();
+
+        empty_circuit
+            .clone()
+            .generate_constraints(&mut counter)
+            .unwrap();
+        // an empty circuit cannot be used with TestConstraintSystem
+        empty_circuit.generate_constraints(&mut cs).unwrap_err();
+        populated_circuit.generate_constraints(&mut cs).unwrap();
+
+        assert_eq!(counter.num_constraints(), cs.num_constraints())
+    }
+}