1use ark_ff::PrimeField;
  2use ark_relations::r1cs::{
  3    ConstraintSynthesizer, ConstraintSystemRef, LinearCombination, SynthesisError, Variable,
  4};
  5
  6use super::R1CS;
  7
  8use color_eyre::Result;
  9
 10#[derive(Clone, Debug)]
 11pub struct CircomCircuit<F: PrimeField> {
 12    pub r1cs: R1CS<F>,
 13    pub witness: Option<Vec<F>>,
 14    pub inputs_already_allocated: bool,
 15}
 16
 17impl<F: PrimeField> CircomCircuit<F> {
 18    pub fn get_public_inputs(&self) -> Option<Vec<F>> {
 19        match &self.witness {
 20            None => None,
 21            Some(w) => match &self.r1cs.wire_mapping {
 22                None => Some(w[1..self.r1cs.num_inputs].to_vec()),
 23                Some(m) => Some(m[1..self.r1cs.num_inputs].iter().map(|i| w[*i]).collect()),
 24            },
 25        }
 26    }
 27}
 28
 29impl<F: PrimeField> ConstraintSynthesizer<F> for CircomCircuit<F> {
 30    fn generate_constraints(self, cs: ConstraintSystemRef<F>) -> Result<(), SynthesisError> {
 31        let witness = &self.witness;
 32        let wire_mapping = &self.r1cs.wire_mapping;
 33
 34        // Start from 1 because Arkworks implicitly allocates One for the first input
 35        if !self.inputs_already_allocated {
 36            for i in 1..self.r1cs.num_inputs {
 37                cs.new_input_variable(|| {
 38                    Ok(match witness {
 39                        None => F::from(1u32),
 40                        Some(w) => match wire_mapping {
 41                            Some(m) => w[m[i]],
 42                            None => w[i],
 43                        },
 44                    })
 45                })?;
 46            }
 47        }
 48
 49        for i in 0..self.r1cs.num_aux {
 50            cs.new_witness_variable(|| {
 51                Ok(match witness {
 52                    None => F::from(1u32),
 53                    Some(w) => match wire_mapping {
 54                        Some(m) => w[m[i + self.r1cs.num_inputs]],
 55                        None => w[i + self.r1cs.num_inputs],
 56                    },
 57                })
 58            })?;
 59        }
 60
 61        let make_index = |index| {
 62            if index < self.r1cs.num_inputs {
 63                Variable::Instance(index)
 64            } else {
 65                Variable::Witness(index - self.r1cs.num_inputs)
 66            }
 67        };
 68        let make_lc = |lc_data: &[(usize, F)]| {
 69            lc_data.iter().fold(
 70                LinearCombination::<F>::zero(),
 71                |lc: LinearCombination<F>, (index, coeff)| lc + (*coeff, make_index(*index)),
 72            )
 73        };
 74
 75        let skip = if self.inputs_already_allocated {
 76            self.r1cs.num_inputs
 77        } else {
 78            0
 79        };
 80        for constraint in self.r1cs.constraints.iter().skip(skip) {
 81            cs.enforce_constraint(
 82                make_lc(&constraint.0),
 83                make_lc(&constraint.1),
 84                make_lc(&constraint.2),
 85            )?;
 86        }
 87
 88        Ok(())
 89    }
 90}
 91
 92#[cfg(test)]
 93mod tests {
 94    use super::*;
 95    use crate::{CircomBuilder, CircomConfig};
 96    use ark_bn254::Fr;
 97    use ark_relations::r1cs::ConstraintSystem;
 98
 99    #[test]
100    fn satisfied() {
101        let cfg = CircomConfig::<Fr>::new(
102            "./test-vectors/mycircuit.wasm",
103            "./test-vectors/mycircuit.r1cs",
104        )
105        .unwrap();
106        let mut builder = CircomBuilder::new(cfg);
107        builder.push_input("a", 3);
108        builder.push_input("b", 11);
109
110        let circom = builder.build().unwrap();
111        let cs = ConstraintSystem::<Fr>::new_ref();
112        circom.generate_constraints(cs.clone()).unwrap();
113        assert!(cs.is_satisfied().unwrap());
114    }
115}