From 48f2915b64a524ec5c4ff3856c0e9467a056db04 Mon Sep 17 00:00:00 2001
From: arnaucube <root@arnaucube.com>
Date: Mon, 25 Mar 2024 22:31:47 +0100
Subject: [PATCH] Add parameter 'inputs_already_allocated' to CircomCircuit

To allow allocating the inputs in previous stages.
---
 src/circom/builder.rs |  1 +
 src/circom/circuit.rs | 30 +++++++++++++++++++-----------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/circom/builder.rs b/src/circom/builder.rs
index 63fab22..5608210 100644
--- a/src/circom/builder.rs
+++ b/src/circom/builder.rs
@@ -58,6 +58,7 @@ impl<F: PrimeField> CircomBuilder<F> {
         let mut circom = CircomCircuit {
             r1cs: self.cfg.r1cs.clone(),
             witness: None,
+            inputs_already_allocated: false,
         };
 
         // Disable the wire mapping
diff --git a/src/circom/circuit.rs b/src/circom/circuit.rs
index 7895712..17e5d9f 100644
--- a/src/circom/circuit.rs
+++ b/src/circom/circuit.rs
@@ -11,6 +11,7 @@ use color_eyre::Result;
 pub struct CircomCircuit<F: PrimeField> {
     pub r1cs: R1CS<F>,
     pub witness: Option<Vec<F>>,
+    pub inputs_already_allocated: bool,
 }
 
 impl<F: PrimeField> CircomCircuit<F> {
@@ -31,16 +32,18 @@ impl<F: PrimeField> ConstraintSynthesizer<F> for CircomCircuit<F> {
         let wire_mapping = &self.r1cs.wire_mapping;
 
         // Start from 1 because Arkworks implicitly allocates One for the first input
-        for i in 1..self.r1cs.num_inputs {
-            cs.new_input_variable(|| {
-                Ok(match witness {
-                    None => F::from(1u32),
-                    Some(w) => match wire_mapping {
-                        Some(m) => w[m[i]],
-                        None => w[i],
-                    },
-                })
-            })?;
+        if !self.inputs_already_allocated {
+            for i in 1..self.r1cs.num_inputs {
+                cs.new_input_variable(|| {
+                    Ok(match witness {
+                        None => F::from(1u32),
+                        Some(w) => match wire_mapping {
+                            Some(m) => w[m[i]],
+                            None => w[i],
+                        },
+                    })
+                })?;
+            }
         }
 
         for i in 0..self.r1cs.num_aux {
@@ -69,7 +72,12 @@ impl<F: PrimeField> ConstraintSynthesizer<F> for CircomCircuit<F> {
             )
         };
 
-        for constraint in &self.r1cs.constraints {
+        let skip = if self.inputs_already_allocated {
+            self.r1cs.num_inputs
+        } else {
+            0
+        };
+        for constraint in self.r1cs.constraints.iter().skip(skip) {
             cs.enforce_constraint(
                 make_lc(&constraint.0),
                 make_lc(&constraint.1),