1use ark_circom::{
  2    circom::{r1cs_reader, R1CS},
  3    WitnessCalculator,
  4};
  5use ark_ff::{BigInteger, PrimeField};
  6use ark_serialize::Read;
  7use num_bigint::{BigInt, Sign};
  8use std::{fs::File, io::Cursor, marker::PhantomData, path::PathBuf};
  9
 10use folding_schemes::{utils::PathOrBin, 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    r1csfile_bytes: Vec<u8>,
 17    wasmfile_bytes: Vec<u8>,
 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: PathOrBin, wasm: PathOrBin) -> Result<Self, Error> {
 24        match (r1cs, wasm) {
 25            (PathOrBin::Path(r1cs_path), PathOrBin::Path(wasm_path)) => {
 26                Self::new_from_path(r1cs_path, wasm_path)
 27            }
 28            (PathOrBin::Bin(r1cs_bin), PathOrBin::Bin(wasm_bin)) => Ok(Self {
 29                r1csfile_bytes: r1cs_bin,
 30                wasmfile_bytes: wasm_bin,
 31                _marker: PhantomData,
 32            }),
 33            _ => unreachable!("You should pass the same enum branch for both inputs"),
 34        }
 35    }
 36
 37    // Creates a new instance of the CircomWrapper with the file paths.
 38    fn new_from_path(r1cs_file_path: PathBuf, wasm_file_path: PathBuf) -> Result<Self, Error> {
 39        let mut file = File::open(r1cs_file_path)?;
 40        let metadata = File::metadata(&file)?;
 41        let mut r1csfile_bytes = vec![0; metadata.len() as usize];
 42        file.read_exact(&mut r1csfile_bytes)?;
 43
 44        let mut file = File::open(wasm_file_path)?;
 45        let metadata = File::metadata(&file)?;
 46        let mut wasmfile_bytes = vec![0; metadata.len() as usize];
 47        file.read_exact(&mut wasmfile_bytes)?;
 48
 49        Ok(CircomWrapper {
 50            r1csfile_bytes,
 51            wasmfile_bytes,
 52            _marker: PhantomData,
 53        })
 54    }
 55
 56    // Aggregated function to obtain R1CS and witness from Circom.
 57    pub fn extract_r1cs_and_witness(
 58        &self,
 59        inputs: &[(String, Vec<BigInt>)],
 60    ) -> Result<(R1CS<F>, Option<Vec<F>>), Error> {
 61        // Extracts the R1CS
 62        let r1cs_file = r1cs_reader::R1CSFile::<F>::new(Cursor::new(&self.r1csfile_bytes))?;
 63        let r1cs = r1cs_reader::R1CS::<F>::from(r1cs_file);
 64
 65        // Extracts the witness vector
 66        let witness_vec = self.extract_witness(inputs)?;
 67
 68        Ok((r1cs, Some(witness_vec)))
 69    }
 70
 71    pub fn extract_r1cs(&self) -> Result<R1CS<F>, Error> {
 72        let r1cs_file = r1cs_reader::R1CSFile::<F>::new(Cursor::new(&self.r1csfile_bytes))?;
 73        let mut r1cs = r1cs_reader::R1CS::<F>::from(r1cs_file);
 74        r1cs.wire_mapping = None;
 75        Ok(r1cs)
 76    }
 77
 78    // Extracts the witness vector as a vector of PrimeField elements.
 79    pub fn extract_witness(&self, inputs: &[(String, Vec<BigInt>)]) -> Result<Vec<F>, Error> {
 80        let witness_bigint = self.calculate_witness(inputs)?;
 81
 82        witness_bigint
 83            .iter()
 84            .map(|big_int| {
 85                self.num_bigint_to_ark_bigint(big_int)
 86                    .and_then(|ark_big_int| {
 87                        F::from_bigint(ark_big_int)
 88                            .ok_or_else(|| Error::Other("could not get F from bigint".to_string()))
 89                    })
 90            })
 91            .collect()
 92    }
 93
 94    // Calculates the witness given the Wasm filepath and inputs.
 95    pub fn calculate_witness(
 96        &self,
 97        inputs: &[(String, Vec<BigInt>)],
 98    ) -> Result<Vec<BigInt>, Error> {
 99        let mut calculator = WitnessCalculator::from_binary(&self.wasmfile_bytes).map_err(|e| {
100            Error::WitnessCalculationError(format!("Failed to create WitnessCalculator: {}", e))
101        })?;
102        calculator
103            .calculate_witness(inputs.iter().cloned(), true)
104            .map_err(|e| {
105                Error::WitnessCalculationError(format!("Failed to calculate witness: {}", e))
106            })
107    }
108
109    // Converts a num_bigint::BigInt to a PrimeField::BigInt.
110    pub fn num_bigint_to_ark_bigint(&self, value: &BigInt) -> Result<F::BigInt, Error> {
111        let big_uint = value.to_biguint().ok_or_else(|| {
112            Error::ConversionError(
113                "BigInt".into(),
114                "BigUint".into(),
115                "BigInt is negative".into(),
116            )
117        })?;
118        F::BigInt::try_from(big_uint).map_err(|_| {
119            Error::ConversionError(
120                "BigUint".into(),
121                "PrimeField::BigInt".into(),
122                "BigUint is too large to fit into PrimeField::BigInt".into(),
123            )
124        })
125    }
126
127    // Converts a PrimeField element to a num_bigint::BigInt representation.
128    pub fn ark_primefield_to_num_bigint(&self, value: F) -> BigInt {
129        let primefield_bigint: F::BigInt = value.into_bigint();
130        let bytes = primefield_bigint.to_bytes_be();
131        BigInt::from_bytes_be(Sign::Plus, &bytes)
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138    use ark_bn254::Fr;
139    use ark_circom::circom::{CircomBuilder, CircomConfig};
140    use ark_circom::CircomCircuit;
141    use ark_relations::r1cs::{ConstraintSynthesizer, ConstraintSystem};
142
143    //To generate .r1cs and .wasm files, run the below command in the terminal.
144    //bash ./frontends/src/circom/test_folder/compile.sh
145
146    // Test the satisfication by using the CircomBuilder of circom-compat
147    #[test]
148    fn test_circombuilder_satisfied() {
149        let cfg = CircomConfig::<Fr>::new(
150            "./src/circom/test_folder/cubic_circuit_js/cubic_circuit.wasm",
151            "./src/circom/test_folder/cubic_circuit.r1cs",
152        )
153        .unwrap();
154        let mut builder = CircomBuilder::new(cfg);
155        builder.push_input("ivc_input", 3);
156
157        let circom = builder.build().unwrap();
158        let cs = ConstraintSystem::<Fr>::new_ref();
159        circom.generate_constraints(cs.clone()).unwrap();
160        assert!(cs.is_satisfied().unwrap());
161    }
162
163    // Test the satisfication by using the CircomWrapper
164    #[test]
165    fn test_extract_r1cs_and_witness() {
166        let r1cs_path = PathBuf::from("./src/circom/test_folder/cubic_circuit.r1cs");
167        let wasm_path =
168            PathBuf::from("./src/circom/test_folder/cubic_circuit_js/cubic_circuit.wasm");
169
170        let inputs = vec![("ivc_input".to_string(), vec![BigInt::from(3)])];
171        let wrapper = CircomWrapper::<Fr>::new(r1cs_path.into(), wasm_path.into()).unwrap();
172
173        let (r1cs, witness) = wrapper.extract_r1cs_and_witness(&inputs).unwrap();
174
175        let cs = ConstraintSystem::<Fr>::new_ref();
176
177        let circom_circuit = CircomCircuit {
178            r1cs,
179            witness,
180            public_inputs_indexes: vec![],
181            allocate_inputs_as_witnesses: false,
182        };
183
184        circom_circuit.generate_constraints(cs.clone()).unwrap();
185        assert!(cs.is_satisfied().unwrap());
186    }
187}