Fix allocated public inputs indexes (#1)

* fix: manage already allocated public inputs with a custom mapping

* fix: initialize CircomCircuit with empty allocated public input variables

* feat: add ability to choose to allocate inputs as witnesses or instances

* chore: remove outdated comment

* test: add relevant tests and asserts for generate_constraints

* chore: add files used for tests

* chore: replace assert_eq with match statement

src/circom/builder.rs

@@ -58,7 +58,8 @@ impl CircomBuilder {
         let mut circom = CircomCircuit {
             r1cs: self.cfg.r1cs.clone(),
             witness: None,
-            inputs_already_allocated: false,
+            public_inputs_indexes: vec![],
+            allocate_inputs_as_witnesses: false,
         };
 
         // Disable the wire mapping

src/circom/circuit.rs

@@ -1,17 +1,28 @@
+use std::{collections::HashMap, f32::consts::E};
+
 use ark_ff::PrimeField;
 use ark_relations::r1cs::{
-    ConstraintSynthesizer, ConstraintSystemRef, LinearCombination, SynthesisError, Variable,
+    self, ConstraintSynthesizer, ConstraintSystemRef, LinearCombination, SynthesisError, Variable,
 };
 
-use super::R1CS;
-
 use color_eyre::Result;
 
+use super::R1CS;
+
+/// `self.public_inputs_indexes` stores already allocated public variables. We assume that:
+///      1. `self.public_inputs_indexes` is sorted in the same order as circom's r1cs public inputs
+///      2. the first element of `self.public_inputs_indexes` is the first element *after* the last non yet allocated public input
+/// example:
+///      if circom public values are [1, out1, out2, in1, in2] (where out and in denote public signals only)
+///      and `self.public_inputs_indexes` is [Variable(2), Variable(3)]
+///      then we will allocate Variable(1), Variable(out1), Variable(out2) and consider that
+///      Variable(in1) and Variable(in2) are already allocated as Variable(2), Variable(3)
 #[derive(Clone, Debug)]
 pub struct CircomCircuit<F: PrimeField> {
     pub r1cs: R1CS<F>,
     pub witness: Option<Vec<F>>,
-    pub inputs_already_allocated: bool,
+    pub public_inputs_indexes: Vec<Variable>,
+    pub allocate_inputs_as_witnesses: bool,
 }
 
 impl<F: PrimeField> CircomCircuit<F> {
@@ -31,58 +42,111 @@ impl ConstraintSynthesizer for CircomCircuit {
         let witness = &self.witness;
         let wire_mapping = &self.r1cs.wire_mapping;
 
-        // Start from 1 because Arkworks implicitly allocates One for the first input
-        if !self.inputs_already_allocated {
-            for i in 1..self.r1cs.num_inputs {
-                cs.new_input_variable(|| {
-                    Ok(match witness {
-                        None => F::from(1u32),
-                        Some(w) => match wire_mapping {
-                            Some(m) => w[m[i]],
-                            None => w[i],
-                        },
-                    })
-                })?;
-            }
+        // Since our cs might already have allocated constraints,
+        // We store a mapping between circom's defined indexes and the newly obtained cs indexes
+        let mut circom_index_to_cs_index = HashMap::new();
+
+        // constant 1 at idx 0 is already allocated by arkworks
+        circom_index_to_cs_index.insert(0, Variable::One);
+
+        // length pub inputs - constant 1 (already allocated by arkworks) - length already allocated pub inputs
+        let n_non_allocated_inputs = self.r1cs.num_inputs - 1 - self.public_inputs_indexes.len();
+
+        // allocate non-allocated inputs and update mapping
+        for circom_public_input_index in 1..=n_non_allocated_inputs {
+            let variable = match self.allocate_inputs_as_witnesses {
+                true => {
+                    let witness_var = Variable::Witness(cs.num_witness_variables());
+                    cs.new_witness_variable(|| {
+                        Ok(match witness {
+                            None => F::from(1u32),
+                            Some(w) => match wire_mapping {
+                                Some(m) => w[m[circom_public_input_index]],
+                                None => w[circom_public_input_index],
+                            },
+                        })
+                    })?;
+                    witness_var
+                }
+                false => {
+                    let instance_var = Variable::Instance(cs.num_instance_variables());
+                    cs.new_input_variable(|| {
+                        Ok(match witness {
+                            None => F::from(1u32),
+                            Some(w) => match wire_mapping {
+                                Some(m) => w[m[circom_public_input_index]],
+                                None => w[circom_public_input_index],
+                            },
+                        })
+                    })?;
+                    instance_var
+                }
+            };
+            circom_index_to_cs_index.insert(circom_public_input_index, variable);
+        }
+
+        // update mapping with already allocated public inputs
+        for circom_public_input_index in (n_non_allocated_inputs + 1)..self.r1cs.num_inputs {
+            let access_input_index = circom_public_input_index - n_non_allocated_inputs - 1;
+            circom_index_to_cs_index.insert(
+                circom_public_input_index,
+                self.public_inputs_indexes[access_input_index],
+            );
         }
 
+        match (
+            circom_index_to_cs_index.get(&0),
+            circom_index_to_cs_index.len(),
+        ) == (Some(&Variable::One), self.r1cs.num_inputs)
+        {
+            true => Ok(()),
+            false => Err(SynthesisError::Unsatisfiable),
+        }?;
+
         for i in 0..self.r1cs.num_aux {
+            let circom_defined_r1cs_index = i + self.r1cs.num_inputs;
+            circom_index_to_cs_index.insert(
+                circom_defined_r1cs_index,
+                Variable::Witness(cs.num_witness_variables()),
+            );
             cs.new_witness_variable(|| {
                 Ok(match witness {
                     None => F::from(1u32),
                     Some(w) => match wire_mapping {
                         Some(m) => w[m[i + self.r1cs.num_inputs]],
-                        None => w[i + self.r1cs.num_inputs],
+                        None => w[circom_defined_r1cs_index],
                     },
                 })
             })?;
         }
 
         let make_index = |index| {
-            if index < self.r1cs.num_inputs {
-                Variable::Instance(index)
-            } else {
-                Variable::Witness(index - self.r1cs.num_inputs)
-            }
+            let new_index = match circom_index_to_cs_index.get(&index) {
+                Some(i) => Ok(*i),
+                None => Err(SynthesisError::AssignmentMissing),
+            };
+            Ok(new_index?)
         };
+
         let make_lc = |lc_data: &[(usize, F)]| {
-            lc_data.iter().fold(
+            lc_data.iter().try_fold(
                 LinearCombination::<F>::zero(),
-                |lc: LinearCombination<F>, (index, coeff)| lc + (*coeff, make_index(*index)),
+                |lc: LinearCombination<F>, (index, coeff)| {
+                    let r1cs_index = make_index(*index);
+                    match r1cs_index {
+                        Ok(r1cs_index) => Ok(lc + (*coeff, r1cs_index)),
+                        Err(e) => Err(e),
+                    }
+                },
             )
         };
 
-        let skip = if self.inputs_already_allocated {
-            self.r1cs.num_inputs
-        } else {
-            0
-        };
-        for constraint in self.r1cs.constraints.iter().skip(skip) {
-            cs.enforce_constraint(
-                make_lc(&constraint.0),
-                make_lc(&constraint.1),
-                make_lc(&constraint.2),
-            )?;
+        for constraint in self.r1cs.constraints.iter() {
+            let a = make_lc(&constraint.0)?;
+            let b = make_lc(&constraint.1)?;
+            let c = make_lc(&constraint.2)?;
+            let res = cs.enforce_constraint(a, b, c);
+            res?
         }
 
         Ok(())
@@ -95,6 +159,7 @@ mod tests {
     use crate::{CircomBuilder, CircomConfig};
     use ark_bn254::Fr;
     use ark_relations::r1cs::ConstraintSystem;
+    use num_bigint::BigInt;
 
     #[test]
     fn satisfied() {
@@ -112,4 +177,148 @@ mod tests {
         circom.generate_constraints(cs.clone()).unwrap();
         assert!(cs.is_satisfied().unwrap());
     }
+
+    #[test]
+    fn no_public_inputs_allocated() {
+        let mut cfg = CircomConfig::<Fr>::new(
+            "./test-vectors/mycircuit2.wasm",
+            "./test-vectors/mycircuit2.r1cs",
+        )
+        .unwrap();
+
+        let inputs = vec![
+            ("a".to_string(), vec![BigInt::from(3)]),
+            ("b".to_string(), vec![BigInt::from(3)]),
+            ("c".to_string(), vec![BigInt::from(3)]),
+        ];
+
+        let witness = cfg.wtns.calculate_witness_element(inputs, true).unwrap();
+
+        // satisfy with inputs as instance variables
+        let mut circom = CircomCircuit {
+            r1cs: cfg.r1cs.clone(),
+            witness: Some(witness.clone()),
+            public_inputs_indexes: vec![],
+            allocate_inputs_as_witnesses: false,
+        };
+
+        let cs = ConstraintSystem::<Fr>::new_ref();
+        circom.clone().generate_constraints(cs.clone()).unwrap();
+        assert_eq!(cs.num_witness_variables(), 2);
+        assert_eq!(cs.num_instance_variables(), 5);
+        assert!(cs.is_satisfied().unwrap());
+
+        // satisfy with inputs as witness variables
+        let cs = ConstraintSystem::<Fr>::new_ref();
+        circom.allocate_inputs_as_witnesses = true;
+        circom.generate_constraints(cs.clone()).unwrap();
+
+        assert_eq!(cs.num_witness_variables(), 6);
+        assert_eq!(cs.num_instance_variables(), 1);
+        assert!(cs.is_satisfied().unwrap());
+    }
+
+    #[test]
+    fn all_public_inputs_allocated() {
+        let mut cfg = CircomConfig::<Fr>::new(
+            "./test-vectors/mycircuit2.wasm",
+            "./test-vectors/mycircuit2.r1cs",
+        )
+        .unwrap();
+
+        let inputs = vec![
+            ("a".to_string(), vec![BigInt::from(3)]),
+            ("b".to_string(), vec![BigInt::from(3)]),
+            ("c".to_string(), vec![BigInt::from(3)]),
+        ];
+
+        let witness = cfg.wtns.calculate_witness_element(inputs, true).unwrap();
+
+        // satisfy with inputs as instance variables
+        let mut circom = CircomCircuit {
+            r1cs: cfg.r1cs.clone(),
+            witness: Some(witness.clone()),
+            public_inputs_indexes: vec![],
+            allocate_inputs_as_witnesses: false,
+        };
+
+        let cs = ConstraintSystem::<Fr>::new_ref();
+
+        // allocate all public input variables (except 1) as instance variables and push their indexes
+        for i in 1..cfg.r1cs.num_inputs {
+            circom
+                .public_inputs_indexes
+                .push(Variable::Instance(cs.num_instance_variables()));
+            cs.new_input_variable(|| Ok(witness[i])).unwrap();
+        }
+
+        circom.clone().generate_constraints(cs.clone()).unwrap();
+
+        assert_eq!(cs.num_witness_variables(), 2);
+        assert_eq!(cs.num_instance_variables(), 5);
+        assert!(cs.is_satisfied().unwrap());
+    }
+
+    #[test]
+    fn some_public_inputs_allocated() {
+        let mut cfg = CircomConfig::<Fr>::new(
+            "./test-vectors/mycircuit2.wasm",
+            "./test-vectors/mycircuit2.r1cs",
+        )
+        .unwrap();
+
+        let inputs = vec![
+            ("a".to_string(), vec![BigInt::from(3)]),
+            ("b".to_string(), vec![BigInt::from(3)]),
+            ("c".to_string(), vec![BigInt::from(3)]),
+        ];
+
+        let witness = cfg.wtns.calculate_witness_element(inputs, true).unwrap();
+
+        // satisfy with inputs as instance variables
+        let mut circom = CircomCircuit {
+            r1cs: cfg.r1cs.clone(),
+            witness: Some(witness.clone()),
+            public_inputs_indexes: vec![],
+            allocate_inputs_as_witnesses: false,
+        };
+
+        let cs = ConstraintSystem::<Fr>::new_ref();
+
+        // allocate all some input variables as instance variables and push their indexes
+        // recall that it should start with the first non-allocated public input
+        for i in 3..cfg.r1cs.num_inputs {
+            circom
+                .public_inputs_indexes
+                .push(Variable::Instance(cs.num_instance_variables()));
+            cs.new_input_variable(|| Ok(witness[i])).unwrap();
+        }
+
+        circom.clone().generate_constraints(cs.clone()).unwrap();
+
+        assert_eq!(cs.num_witness_variables(), 2);
+        assert_eq!(cs.num_instance_variables(), 5);
+        assert!(cs.is_satisfied().unwrap());
+
+        // re-init and now satisfy with inputs as witness variables
+        let mut circom = CircomCircuit {
+            r1cs: cfg.r1cs.clone(),
+            witness: Some(witness.clone()),
+            public_inputs_indexes: vec![],
+            allocate_inputs_as_witnesses: false,
+        };
+
+        let cs = ConstraintSystem::<Fr>::new_ref();
+        for i in 3..cfg.r1cs.num_inputs {
+            circom
+                .public_inputs_indexes
+                .push(Variable::Instance(cs.num_instance_variables()));
+            cs.new_input_variable(|| Ok(witness[i])).unwrap();
+        }
+        circom.allocate_inputs_as_witnesses = true;
+        circom.generate_constraints(cs.clone()).unwrap();
+
+        assert_eq!(cs.num_witness_variables(), 4);
+        assert_eq!(cs.num_instance_variables(), 3);
+    }
 }

test-vectors/mycircuit2.circom

@@ -0,0 +1,15 @@
+template Multiplier() {
+    signal input a; // 3
+    signal input b; // 3
+    signal input c; // 3
+    signal temp;
+    signal output d; // 27
+    signal output e; // 9
+
+    temp <== a*b;
+    d <== temp*c;
+    e <== b*c;
+}
+
+component main{public [a, b]} = Multiplier();
+

test-vectors/mycircuit2.sym

@@ -0,0 +1,6 @@
+1,1,0,main.d
+2,2,0,main.e
+3,3,0,main.a
+4,4,0,main.b
+5,5,0,main.c
+6,6,0,main.temp

test-vectors/mycircuit2_constraints.json

@@ -0,0 +1,7 @@
+{
+"constraints": [
+[{"3":"21888242871839275222246405745257275088548364400416034343698204186575808495616"},{"4":"1"},{"6":"21888242871839275222246405745257275088548364400416034343698204186575808495616"}],
+[{"6":"21888242871839275222246405745257275088548364400416034343698204186575808495616"},{"5":"1"},{"1":"21888242871839275222246405745257275088548364400416034343698204186575808495616"}],
+[{"4":"21888242871839275222246405745257275088548364400416034343698204186575808495616"},{"5":"1"},{"2":"21888242871839275222246405745257275088548364400416034343698204186575808495616"}]
+]
+}