feat (circom): allow to define the step_native in Rust (#105)

* create function pointer

* custom logic via function pointer

* fmt

* clippy

* rust-version

* update review code

* fmt

folding-schemes/src/frontend/circom/mod.rs

@@ -10,49 +10,53 @@ use ark_relations::r1cs::{ConstraintSynthesizer, ConstraintSystemRef, SynthesisE
 use ark_std::fmt::Debug;
 use num_bigint::BigInt;
 use std::path::PathBuf;
+use std::rc::Rc;
+use std::{fmt, usize};
 
 pub mod utils;
 use utils::CircomWrapper;
 
+type ClosurePointer<F> = Rc<dyn Fn(usize, Vec<F>, Vec<F>) -> Result<Vec<F>, Error>>;
+
+#[derive(Clone)]
+struct CustomStepNative<F: PrimeField> {
+    func: ClosurePointer<F>,
+}
+
+impl<F: PrimeField> fmt::Debug for CustomStepNative<F> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "Function pointer: {:?}",
+            std::any::type_name::<fn(usize, Vec<F>, Vec<F>) -> Result<Vec<F>, Error>>()
+        )
+    }
+}
+
 /// Define CircomFCircuit
 #[derive(Clone, Debug)]
 pub struct CircomFCircuit<F: PrimeField> {
     circom_wrapper: CircomWrapper<F>,
-    state_len: usize,
+    pub state_len: usize,
-    external_inputs_len: usize,
+    pub external_inputs_len: usize,
     r1cs: CircomR1CS<F>,
+    custom_step_native_code: Option<CustomStepNative<F>>,
 }
 
-impl<F: PrimeField> FCircuit<F> for CircomFCircuit<F> {
+impl<F: PrimeField> CircomFCircuit<F> {
-    /// (r1cs_path, wasm_path, state_len, external_inputs_len)
+    pub fn set_custom_step_native(&mut self, func: ClosurePointer<F>) {
-    type Params = (PathBuf, PathBuf, usize, usize);
+        self.custom_step_native_code = Some(CustomStepNative::<F> { func });
-
-    fn new(params: Self::Params) -> Result<Self, Error> {
-        let (r1cs_path, wasm_path, state_len, external_inputs_len) = params;
-        let circom_wrapper = CircomWrapper::new(r1cs_path, wasm_path);
-
-        let r1cs = circom_wrapper.extract_r1cs()?;
-        Ok(Self {
-            circom_wrapper,
-            state_len,
-            external_inputs_len,
-            r1cs,
-        })
     }
 
-    fn state_len(&self) -> usize {
+    pub fn execute_custom_step_native(
-        self.state_len
-    }
-    fn external_inputs_len(&self) -> usize {
-        self.external_inputs_len
-    }
-
-    fn step_native(
         &self,
         _i: usize,
         z_i: Vec<F>,
         external_inputs: Vec<F>,
     ) -> Result<Vec<F>, Error> {
+        if let Some(code) = &self.custom_step_native_code {
+            (code.func)(_i, z_i, external_inputs)
+        } else {
             #[cfg(test)]
             assert_eq!(z_i.len(), self.state_len());
             #[cfg(test)]
@@ -84,6 +88,42 @@ impl<F: PrimeField> FCircuit<F> for CircomFCircuit<F> {
             let z_i1 = witness[1..1 + self.state_len()].to_vec();
             Ok(z_i1)
         }
+    }
+}
+
+impl<F: PrimeField> FCircuit<F> for CircomFCircuit<F> {
+    /// (r1cs_path, wasm_path, state_len, external_inputs_len)
+    type Params = (PathBuf, PathBuf, usize, usize);
+
+    fn new(params: Self::Params) -> Result<Self, Error> {
+        let (r1cs_path, wasm_path, state_len, external_inputs_len) = params;
+        let circom_wrapper = CircomWrapper::new(r1cs_path, wasm_path);
+
+        let r1cs = circom_wrapper.extract_r1cs()?;
+        Ok(Self {
+            circom_wrapper,
+            state_len,
+            external_inputs_len,
+            r1cs,
+            custom_step_native_code: None,
+        })
+    }
+
+    fn state_len(&self) -> usize {
+        self.state_len
+    }
+    fn external_inputs_len(&self) -> usize {
+        self.external_inputs_len
+    }
+
+    fn step_native(
+        &self,
+        _i: usize,
+        z_i: Vec<F>,
+        external_inputs: Vec<F>,
+    ) -> Result<Vec<F>, Error> {
+        self.execute_custom_step_native(_i, z_i, external_inputs)
+    }
 
     fn generate_step_constraints(
         &self,
@@ -305,4 +345,34 @@ pub mod tests {
         // Disable check for now
         // assert!(z_i1_var.is_err())
     }
+
+    #[test]
+    fn test_custom_code() {
+        let r1cs_path = PathBuf::from("./src/frontend/circom/test_folder/cubic_circuit.r1cs");
+        let wasm_path =
+            PathBuf::from("./src/frontend/circom/test_folder/cubic_circuit_js/cubic_circuit.wasm");
+
+        let mut circom_fcircuit = CircomFCircuit::<Fr>::new((r1cs_path, wasm_path, 1, 0)).unwrap(); // state_len:1, external_inputs_len:0
+
+        circom_fcircuit.set_custom_step_native(Rc::new(|_i, z_i, _external| {
+            let z = z_i[0];
+            Ok(vec![z * z * z + z + Fr::from(5)])
+        }));
+
+        // Allocates z_i1 by using step_native function.
+        let z_i = vec![Fr::from(3_u32)];
+        let wrapper_circuit = crate::frontend::tests::WrapperCircuit {
+            FC: circom_fcircuit.clone(),
+            z_i: Some(z_i.clone()),
+            z_i1: Some(circom_fcircuit.step_native(0, z_i.clone(), vec![]).unwrap()),
+        };
+
+        let cs = ConstraintSystem::<Fr>::new_ref();
+
+        wrapper_circuit.generate_constraints(cs.clone()).unwrap();
+        assert!(
+            cs.is_satisfied().unwrap(),
+            "Constraint system is not satisfied"
+        );
+    }
 }