1use std::{collections::HashMap, f32::consts::E};
  2
  3use ark_ff::PrimeField;
  4use ark_relations::r1cs::{
  5    self, ConstraintSynthesizer, ConstraintSystemRef, LinearCombination, SynthesisError, Variable,
  6};
  7
  8use color_eyre::Result;
  9
 10use super::R1CS;
 11
 12/// `self.public_inputs_indexes` stores already allocated public variables. We assume that:
 13///      1. `self.public_inputs_indexes` is sorted in the same order as circom's r1cs public inputs
 14///      2. the first element of `self.public_inputs_indexes` is the first element *after* the last non yet allocated public input
 15/// example:
 16///      if circom public values are [1, out1, out2, in1, in2] (where out and in denote public signals only)
 17///      and `self.public_inputs_indexes` is [Variable(2), Variable(3)]
 18///      then we will allocate Variable(1), Variable(out1), Variable(out2) and consider that
 19///      Variable(in1) and Variable(in2) are already allocated as Variable(2), Variable(3)
 20#[derive(Clone, Debug)]
 21pub struct CircomCircuit<F: PrimeField> {
 22    pub r1cs: R1CS<F>,
 23    pub witness: Option<Vec<F>>,
 24    pub public_inputs_indexes: Vec<Variable>,
 25    pub allocate_inputs_as_witnesses: bool,
 26}
 27
 28impl<F: PrimeField> CircomCircuit<F> {
 29    pub fn get_public_inputs(&self) -> Option<Vec<F>> {
 30        match &self.witness {
 31            None => None,
 32            Some(w) => match &self.r1cs.wire_mapping {
 33                None => Some(w[1..self.r1cs.num_inputs].to_vec()),
 34                Some(m) => Some(m[1..self.r1cs.num_inputs].iter().map(|i| w[*i]).collect()),
 35            },
 36        }
 37    }
 38}
 39
 40impl<F: PrimeField> ConstraintSynthesizer<F> for CircomCircuit<F> {
 41    fn generate_constraints(self, cs: ConstraintSystemRef<F>) -> Result<(), SynthesisError> {
 42        let witness = &self.witness;
 43        let wire_mapping = &self.r1cs.wire_mapping;
 44
 45        // Since our cs might already have allocated constraints,
 46        // We store a mapping between circom's defined indexes and the newly obtained cs indexes
 47        let mut circom_index_to_cs_index = HashMap::new();
 48
 49        // constant 1 at idx 0 is already allocated by arkworks
 50        circom_index_to_cs_index.insert(0, Variable::One);
 51
 52        // length pub inputs - constant 1 (already allocated by arkworks) - length already allocated pub inputs
 53        let n_non_allocated_inputs = self.r1cs.num_inputs - 1 - self.public_inputs_indexes.len();
 54
 55        // allocate non-allocated inputs and update mapping
 56        for circom_public_input_index in 1..=n_non_allocated_inputs {
 57            let variable = match self.allocate_inputs_as_witnesses {
 58                true => {
 59                    let witness_var = Variable::Witness(cs.num_witness_variables());
 60                    cs.new_witness_variable(|| {
 61                        Ok(match witness {
 62                            None => F::from(1u32),
 63                            Some(w) => match wire_mapping {
 64                                Some(m) => w[m[circom_public_input_index]],
 65                                None => w[circom_public_input_index],
 66                            },
 67                        })
 68                    })?;
 69                    witness_var
 70                }
 71                false => {
 72                    let instance_var = Variable::Instance(cs.num_instance_variables());
 73                    cs.new_input_variable(|| {
 74                        Ok(match witness {
 75                            None => F::from(1u32),
 76                            Some(w) => match wire_mapping {
 77                                Some(m) => w[m[circom_public_input_index]],
 78                                None => w[circom_public_input_index],
 79                            },
 80                        })
 81                    })?;
 82                    instance_var
 83                }
 84            };
 85            circom_index_to_cs_index.insert(circom_public_input_index, variable);
 86        }
 87
 88        // update mapping with already allocated public inputs
 89        for circom_public_input_index in (n_non_allocated_inputs + 1)..self.r1cs.num_inputs {
 90            let access_input_index = circom_public_input_index - n_non_allocated_inputs - 1;
 91            circom_index_to_cs_index.insert(
 92                circom_public_input_index,
 93                self.public_inputs_indexes[access_input_index],
 94            );
 95        }
 96
 97        match (
 98            circom_index_to_cs_index.get(&0),
 99            circom_index_to_cs_index.len(),
100        ) == (Some(&Variable::One), self.r1cs.num_inputs)
101        {
102            true => Ok(()),
103            false => Err(SynthesisError::Unsatisfiable),
104        }?;
105
106        for i in 0..self.r1cs.num_aux {
107            let circom_defined_r1cs_index = i + self.r1cs.num_inputs;
108            circom_index_to_cs_index.insert(
109                circom_defined_r1cs_index,
110                Variable::Witness(cs.num_witness_variables()),
111            );
112            cs.new_witness_variable(|| {
113                Ok(match witness {
114                    None => F::from(1u32),
115                    Some(w) => match wire_mapping {
116                        Some(m) => w[m[i + self.r1cs.num_inputs]],
117                        None => w[circom_defined_r1cs_index],
118                    },
119                })
120            })?;
121        }
122
123        let make_index = |index| {
124            let new_index = match circom_index_to_cs_index.get(&index) {
125                Some(i) => Ok(*i),
126                None => Err(SynthesisError::AssignmentMissing),
127            };
128            Ok(new_index?)
129        };
130
131        let make_lc = |lc_data: &[(usize, F)]| {
132            lc_data.iter().try_fold(
133                LinearCombination::<F>::zero(),
134                |lc: LinearCombination<F>, (index, coeff)| {
135                    let r1cs_index = make_index(*index);
136                    match r1cs_index {
137                        Ok(r1cs_index) => Ok(lc + (*coeff, r1cs_index)),
138                        Err(e) => Err(e),
139                    }
140                },
141            )
142        };
143
144        for constraint in self.r1cs.constraints.iter() {
145            let a = make_lc(&constraint.0)?;
146            let b = make_lc(&constraint.1)?;
147            let c = make_lc(&constraint.2)?;
148            let res = cs.enforce_constraint(a, b, c);
149            res?
150        }
151
152        Ok(())
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159    use crate::{CircomBuilder, CircomConfig};
160    use ark_bn254::Fr;
161    use ark_relations::r1cs::ConstraintSystem;
162    use num_bigint::BigInt;
163
164    #[test]
165    fn satisfied() {
166        let cfg = CircomConfig::<Fr>::new(
167            "./test-vectors/mycircuit.wasm",
168            "./test-vectors/mycircuit.r1cs",
169        )
170        .unwrap();
171        let mut builder = CircomBuilder::new(cfg);
172        builder.push_input("a", 3);
173        builder.push_input("b", 11);
174
175        let circom = builder.build().unwrap();
176        let cs = ConstraintSystem::<Fr>::new_ref();
177        circom.generate_constraints(cs.clone()).unwrap();
178        assert!(cs.is_satisfied().unwrap());
179    }
180
181    #[test]
182    fn no_public_inputs_allocated() {
183        let mut cfg = CircomConfig::<Fr>::new(
184            "./test-vectors/mycircuit2.wasm",
185            "./test-vectors/mycircuit2.r1cs",
186        )
187        .unwrap();
188
189        let inputs = vec![
190            ("a".to_string(), vec![BigInt::from(3)]),
191            ("b".to_string(), vec![BigInt::from(3)]),
192            ("c".to_string(), vec![BigInt::from(3)]),
193        ];
194
195        let witness = cfg.wtns.calculate_witness_element(inputs, true).unwrap();
196
197        // satisfy with inputs as instance variables
198        let mut circom = CircomCircuit {
199            r1cs: cfg.r1cs.clone(),
200            witness: Some(witness.clone()),
201            public_inputs_indexes: vec![],
202            allocate_inputs_as_witnesses: false,
203        };
204
205        let cs = ConstraintSystem::<Fr>::new_ref();
206        circom.clone().generate_constraints(cs.clone()).unwrap();
207        assert_eq!(cs.num_witness_variables(), 2);
208        assert_eq!(cs.num_instance_variables(), 5);
209        assert!(cs.is_satisfied().unwrap());
210
211        // satisfy with inputs as witness variables
212        let cs = ConstraintSystem::<Fr>::new_ref();
213        circom.allocate_inputs_as_witnesses = true;
214        circom.generate_constraints(cs.clone()).unwrap();
215
216        assert_eq!(cs.num_witness_variables(), 6);
217        assert_eq!(cs.num_instance_variables(), 1);
218        assert!(cs.is_satisfied().unwrap());
219    }
220
221    #[test]
222    fn all_public_inputs_allocated() {
223        let mut cfg = CircomConfig::<Fr>::new(
224            "./test-vectors/mycircuit2.wasm",
225            "./test-vectors/mycircuit2.r1cs",
226        )
227        .unwrap();
228
229        let inputs = vec![
230            ("a".to_string(), vec![BigInt::from(3)]),
231            ("b".to_string(), vec![BigInt::from(3)]),
232            ("c".to_string(), vec![BigInt::from(3)]),
233        ];
234
235        let witness = cfg.wtns.calculate_witness_element(inputs, true).unwrap();
236
237        // satisfy with inputs as instance variables
238        let mut circom = CircomCircuit {
239            r1cs: cfg.r1cs.clone(),
240            witness: Some(witness.clone()),
241            public_inputs_indexes: vec![],
242            allocate_inputs_as_witnesses: false,
243        };
244
245        let cs = ConstraintSystem::<Fr>::new_ref();
246
247        // allocate all public input variables (except 1) as instance variables and push their indexes
248        for i in 1..cfg.r1cs.num_inputs {
249            circom
250                .public_inputs_indexes
251                .push(Variable::Instance(cs.num_instance_variables()));
252            cs.new_input_variable(|| Ok(witness[i])).unwrap();
253        }
254
255        circom.clone().generate_constraints(cs.clone()).unwrap();
256
257        assert_eq!(cs.num_witness_variables(), 2);
258        assert_eq!(cs.num_instance_variables(), 5);
259        assert!(cs.is_satisfied().unwrap());
260    }
261
262    #[test]
263    fn some_public_inputs_allocated() {
264        let mut cfg = CircomConfig::<Fr>::new(
265            "./test-vectors/mycircuit2.wasm",
266            "./test-vectors/mycircuit2.r1cs",
267        )
268        .unwrap();
269
270        let inputs = vec![
271            ("a".to_string(), vec![BigInt::from(3)]),
272            ("b".to_string(), vec![BigInt::from(3)]),
273            ("c".to_string(), vec![BigInt::from(3)]),
274        ];
275
276        let witness = cfg.wtns.calculate_witness_element(inputs, true).unwrap();
277
278        // satisfy with inputs as instance variables
279        let mut circom = CircomCircuit {
280            r1cs: cfg.r1cs.clone(),
281            witness: Some(witness.clone()),
282            public_inputs_indexes: vec![],
283            allocate_inputs_as_witnesses: false,
284        };
285
286        let cs = ConstraintSystem::<Fr>::new_ref();
287
288        // allocate all some input variables as instance variables and push their indexes
289        // recall that it should start with the first non-allocated public input
290        for i in 3..cfg.r1cs.num_inputs {
291            circom
292                .public_inputs_indexes
293                .push(Variable::Instance(cs.num_instance_variables()));
294            cs.new_input_variable(|| Ok(witness[i])).unwrap();
295        }
296
297        circom.clone().generate_constraints(cs.clone()).unwrap();
298
299        assert_eq!(cs.num_witness_variables(), 2);
300        assert_eq!(cs.num_instance_variables(), 5);
301        assert!(cs.is_satisfied().unwrap());
302
303        // re-init and now satisfy with inputs as witness variables
304        let mut circom = CircomCircuit {
305            r1cs: cfg.r1cs.clone(),
306            witness: Some(witness.clone()),
307            public_inputs_indexes: vec![],
308            allocate_inputs_as_witnesses: false,
309        };
310
311        let cs = ConstraintSystem::<Fr>::new_ref();
312        for i in 3..cfg.r1cs.num_inputs {
313            circom
314                .public_inputs_indexes
315                .push(Variable::Instance(cs.num_instance_variables()));
316            cs.new_input_variable(|| Ok(witness[i])).unwrap();
317        }
318        circom.allocate_inputs_as_witnesses = true;
319        circom.generate_constraints(cs.clone()).unwrap();
320
321        assert_eq!(cs.num_witness_variables(), 4);
322        assert_eq!(cs.num_instance_variables(), 3);
323    }
324}