1use ark_circom::{
  2    circom::{r1cs_reader, R1CS},
  3    WitnessCalculator,
  4};
  5use ark_ff::{BigInteger, PrimeField};
  6use color_eyre::Result;
  7use num_bigint::{BigInt, Sign};
  8use std::{fs::File, io::BufReader, marker::PhantomData, path::PathBuf};
  9
 10use crate::Error;
 11
 12// A struct that wraps Circom functionalities, allowing for extraction of R1CS and witnesses
 13// based on file paths to Circom's .r1cs and .wasm.
 14#[derive(Clone, Debug)]
 15pub struct CircomWrapper<F: PrimeField> {
 16    r1cs_filepath: PathBuf,
 17    wasm_filepath: PathBuf,
 18    _marker: PhantomData<F>,
 19}
 20
 21impl<F: PrimeField> CircomWrapper<F> {
 22    // Creates a new instance of the CircomWrapper with the file paths.
 23    pub fn new(r1cs_filepath: PathBuf, wasm_filepath: PathBuf) -> Self {
 24        CircomWrapper {
 25            r1cs_filepath,
 26            wasm_filepath,
 27            _marker: PhantomData,
 28        }
 29    }
 30
 31    // Aggregated function to obtain R1CS and witness from Circom.
 32    pub fn extract_r1cs_and_witness(
 33        &self,
 34        inputs: &[(String, Vec<BigInt>)],
 35    ) -> Result<(R1CS<F>, Option<Vec<F>>), Error> {
 36        // Extracts the R1CS
 37        let file = File::open(&self.r1cs_filepath)?;
 38        let reader = BufReader::new(file);
 39        let r1cs_file = r1cs_reader::R1CSFile::<F>::new(reader)?;
 40        let r1cs = r1cs_reader::R1CS::<F>::from(r1cs_file);
 41
 42        // Extracts the witness vector
 43        let witness_vec = self.extract_witness(inputs)?;
 44
 45        Ok((r1cs, Some(witness_vec)))
 46    }
 47
 48    pub fn extract_r1cs(&self) -> Result<R1CS<F>, Error> {
 49        let file = File::open(&self.r1cs_filepath)?;
 50        let reader = BufReader::new(file);
 51        let r1cs_file = r1cs_reader::R1CSFile::<F>::new(reader)?;
 52        let mut r1cs = r1cs_reader::R1CS::<F>::from(r1cs_file);
 53        r1cs.wire_mapping = None;
 54        Ok(r1cs)
 55    }
 56
 57    // Extracts the witness vector as a vector of PrimeField elements.
 58    pub fn extract_witness(&self, inputs: &[(String, Vec<BigInt>)]) -> Result<Vec<F>, Error> {
 59        let witness_bigint = self.calculate_witness(inputs)?;
 60
 61        witness_bigint
 62            .iter()
 63            .map(|big_int| {
 64                self.num_bigint_to_ark_bigint(big_int)
 65                    .and_then(|ark_big_int| {
 66                        F::from_bigint(ark_big_int)
 67                            .ok_or_else(|| Error::Other("could not get F from bigint".to_string()))
 68                    })
 69            })
 70            .collect()
 71    }
 72
 73    // Calculates the witness given the Wasm filepath and inputs.
 74    pub fn calculate_witness(
 75        &self,
 76        inputs: &[(String, Vec<BigInt>)],
 77    ) -> Result<Vec<BigInt>, Error> {
 78        let mut calculator = WitnessCalculator::new(&self.wasm_filepath).map_err(|e| {
 79            Error::WitnessCalculationError(format!("Failed to create WitnessCalculator: {}", e))
 80        })?;
 81        calculator
 82            .calculate_witness(inputs.iter().cloned(), true)
 83            .map_err(|e| {
 84                Error::WitnessCalculationError(format!("Failed to calculate witness: {}", e))
 85            })
 86    }
 87
 88    // Converts a num_bigint::BigInt to a PrimeField::BigInt.
 89    pub fn num_bigint_to_ark_bigint(&self, value: &BigInt) -> Result<F::BigInt, Error> {
 90        let big_uint = value
 91            .to_biguint()
 92            .ok_or_else(|| Error::BigIntConversionError("BigInt is negative".to_string()))?;
 93        F::BigInt::try_from(big_uint).map_err(|_| {
 94            Error::BigIntConversionError("Failed to convert to PrimeField::BigInt".to_string())
 95        })
 96    }
 97
 98    // Converts a PrimeField element to a num_bigint::BigInt representation.
 99    pub fn ark_primefield_to_num_bigint(&self, value: F) -> BigInt {
100        let primefield_bigint: F::BigInt = value.into_bigint();
101        let bytes = primefield_bigint.to_bytes_be();
102        BigInt::from_bytes_be(Sign::Plus, &bytes)
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use ark_bn254::Fr;
110    use ark_circom::circom::{CircomBuilder, CircomConfig};
111    use ark_circom::CircomCircuit;
112    use ark_relations::r1cs::{ConstraintSynthesizer, ConstraintSystem};
113
114    //To generate .r1cs and .wasm files, run the below command in the terminal.
115    //bash ./folding-schemes/src/frontend/circom/test_folder/compile.sh
116
117    // Test the satisfication by using the CircomBuilder of circom-compat
118    #[test]
119    fn test_circombuilder_satisfied() {
120        let cfg = CircomConfig::<Fr>::new(
121            "./src/frontend/circom/test_folder/cubic_circuit_js/cubic_circuit.wasm",
122            "./src/frontend/circom/test_folder/cubic_circuit.r1cs",
123        )
124        .unwrap();
125        let mut builder = CircomBuilder::new(cfg);
126        builder.push_input("ivc_input", 3);
127
128        let circom = builder.build().unwrap();
129        let cs = ConstraintSystem::<Fr>::new_ref();
130        circom.generate_constraints(cs.clone()).unwrap();
131        assert!(cs.is_satisfied().unwrap());
132    }
133
134    // Test the satisfication by using the CircomWrapper
135    #[test]
136    fn test_extract_r1cs_and_witness() {
137        let r1cs_path = PathBuf::from("./src/frontend/circom/test_folder/cubic_circuit.r1cs");
138        let wasm_path =
139            PathBuf::from("./src/frontend/circom/test_folder/cubic_circuit_js/cubic_circuit.wasm");
140
141        let inputs = vec![("ivc_input".to_string(), vec![BigInt::from(3)])];
142        let wrapper = CircomWrapper::<Fr>::new(r1cs_path, wasm_path);
143
144        let (r1cs, witness) = wrapper.extract_r1cs_and_witness(&inputs).unwrap();
145
146        let cs = ConstraintSystem::<Fr>::new_ref();
147
148        let circom_circuit = CircomCircuit {
149            r1cs,
150            witness,
151            public_inputs_indexes: vec![],
152            allocate_inputs_as_witnesses: false,
153        };
154
155        circom_circuit.generate_constraints(cs.clone()).unwrap();
156        assert!(cs.is_satisfied().unwrap());
157    }
158}