Make synthesis limit tracking robust against function calls

circuit/environment/src/canary_circuit.rs

@@ -296,8 +296,9 @@ impl Environment for CanaryCircuit {
         panic!("{}", &error)
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn inject_r1cs(r1cs: R1CS<Self::BaseField>) {
+    /// Injects the ejected R1CS circuit and restores the preserved environment state.
+    fn inject_r1cs(ejected: EjectedR1cs<Self::BaseField>) {
+        let EjectedR1cs { r1cs, variable_limit, constraint_limit, non_zero_limit, in_witness } = ejected;
         CANARY_CIRCUIT.with(|circuit| {
             // Ensure the circuit is empty before injecting.
             assert_eq!(0, circuit.borrow().num_constants());
@@ -313,19 +314,31 @@ impl Environment for CanaryCircuit {
             assert_eq!(0, r1cs.num_private());
             assert_eq!(1, r1cs.num_variables());
             assert_eq!(0, r1cs.num_constraints());
-        })
+        });
+        // Restore the preserved limits.
+        Self::set_variable_limit(variable_limit);
+        Self::set_constraint_limit(constraint_limit);
+        Self::set_non_zero_limit(non_zero_limit);
+        // Restore the preserved witness mode.
+        IN_WITNESS.with(|in_witness_cell| in_witness_cell.replace(in_witness));
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn eject_r1cs_and_reset() -> R1CS<Self::BaseField> {
+    /// Returns the R1CS circuit and preserved environment state, resetting the circuit.
+    fn eject_r1cs_and_reset() -> EjectedR1cs<Self::BaseField> {
         CANARY_CIRCUIT.with(|circuit| {
-            // Reset the witness mode.
-            IN_WITNESS.with(|in_witness| in_witness.replace(false));
-            // Reset the variable limit.
+            // Preserve the witness mode.
+            let in_witness = IN_WITNESS.with(|in_witness| {
+                let value = in_witness.get();
+                in_witness.replace(false);
+                value
+            });
+            // Preserve the limits.
+            let variable_limit = Self::get_variable_limit();
+            let constraint_limit = Self::get_constraint_limit();
+            let non_zero_limit = Self::get_non_zero_limit();
+            // Reset the limits for the temporary environment.
             Self::set_variable_limit(None);
-            // Reset the constraint limit.
             Self::set_constraint_limit(None);
-            // Reset the density limit.
             Self::set_non_zero_limit(None);
             // Eject the R1CS instance.
             let r1cs = circuit.replace(R1CS::<<Self as Environment>::BaseField>::new());
@@ -335,8 +348,8 @@ impl Environment for CanaryCircuit {
             assert_eq!(0, circuit.borrow().num_private());
             assert_eq!(1, circuit.borrow().num_variables());
             assert_eq!(0, circuit.borrow().num_constraints());
-            // Return the R1CS instance.
-            r1cs
+            // Return the ejected environment state.
+            EjectedR1cs::new(r1cs, variable_limit, constraint_limit, non_zero_limit, in_witness)
         })
     }
 

circuit/environment/src/circuit.rs

@@ -320,8 +320,9 @@ impl Environment for Circuit {
         panic!("{}", &error)
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn inject_r1cs(r1cs: R1CS<Self::BaseField>) {
+    /// Injects the ejected R1CS circuit and restores the preserved environment state.
+    fn inject_r1cs(ejected: EjectedR1cs<Self::BaseField>) {
+        let EjectedR1cs { r1cs, variable_limit, constraint_limit, non_zero_limit, in_witness } = ejected;
         CIRCUIT.with(|circuit| {
             // Ensure the circuit is empty before injecting.
             assert_eq!(0, circuit.borrow().num_constants());
@@ -337,19 +338,31 @@ impl Environment for Circuit {
             assert_eq!(0, r1cs.num_private());
             assert_eq!(1, r1cs.num_variables());
             assert_eq!(0, r1cs.num_constraints());
-        })
+        });
+        // Restore the preserved limits.
+        Self::set_variable_limit(variable_limit);
+        Self::set_constraint_limit(constraint_limit);
+        Self::set_non_zero_limit(non_zero_limit);
+        // Restore the preserved witness mode.
+        IN_WITNESS.with(|in_witness_cell| in_witness_cell.replace(in_witness));
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn eject_r1cs_and_reset() -> R1CS<Self::BaseField> {
+    /// Returns the R1CS circuit and preserved environment state, resetting the circuit.
+    fn eject_r1cs_and_reset() -> EjectedR1cs<Self::BaseField> {
         CIRCUIT.with(|circuit| {
-            // Reset the witness mode.
-            IN_WITNESS.with(|in_witness| in_witness.replace(false));
-            // Reset the variable limit.
+            // Preserve the witness mode.
+            let in_witness = IN_WITNESS.with(|in_witness| {
+                let value = in_witness.get();
+                in_witness.replace(false);
+                value
+            });
+            // Preserve the limits.
+            let variable_limit = Self::get_variable_limit();
+            let constraint_limit = Self::get_constraint_limit();
+            let non_zero_limit = Self::get_non_zero_limit();
+            // Reset the limits for the temporary environment.
             Self::set_variable_limit(None);
-            // Reset the constraint limit.
             Self::set_constraint_limit(None);
-            // Reset the density limit.
             Self::set_non_zero_limit(None);
             // Eject the R1CS instance.
             let r1cs = circuit.replace(R1CS::<<Self as Environment>::BaseField>::new());
@@ -359,8 +372,8 @@ impl Environment for Circuit {
             assert_eq!(0, circuit.borrow().num_private());
             assert_eq!(1, circuit.borrow().num_variables());
             assert_eq!(0, circuit.borrow().num_constraints());
-            // Return the R1CS instance.
-            r1cs
+            // Return the ejected environment state.
+            EjectedR1cs::new(r1cs, variable_limit, constraint_limit, non_zero_limit, in_witness)
         })
     }
 
@@ -464,4 +477,46 @@ mod tests {
             assert_eq!(0, Circuit::num_constraints_in_scope());
         })
     }
+
+    #[test]
+    fn test_eject_inject_preserves_limits() {
+        use std::panic::{self, AssertUnwindSafe};
+
+        /// Compute 2^EXPONENT - 1, in a purposefully constraint-inefficient manner for testing.
+        fn add_constraints<E: Environment>(exponent: u64) {
+            let one = snarkvm_console_types::Field::<<E as Environment>::Network>::one();
+            let two = one + one;
+
+            let mut candidate = Field::<E>::new(Mode::Public, one);
+            let mut accumulator = Field::<E>::new(Mode::Private, two);
+            for _ in 0..exponent {
+                candidate += &accumulator;
+                accumulator *= Field::<E>::new(Mode::Private, two);
+            }
+        }
+
+        Circuit::reset();
+        Circuit::set_constraint_limit(Some(20));
+
+        add_constraints::<Circuit>(10);
+        assert_eq!(10, Circuit::num_constraints());
+
+        let ejected = Circuit::eject_r1cs_and_reset();
+        assert_eq!(None, Circuit::get_constraint_limit());
+        assert_eq!(Some(20), ejected.constraint_limit());
+
+        // Without limits, the temporary environment can synthesize beyond the preserved limit.
+        add_constraints::<Circuit>(25);
+        assert_eq!(25, Circuit::num_constraints());
+
+        // Nested execution resets the circuit before returning control to the caller.
+        Circuit::reset();
+
+        Circuit::inject_r1cs(ejected);
+        assert_eq!(Some(20), Circuit::get_constraint_limit());
+        assert_eq!(10, Circuit::num_constraints());
+
+        let result = panic::catch_unwind(AssertUnwindSafe(|| add_constraints::<Circuit>(12)));
+        assert!(result.is_err());
+    }
 }

circuit/environment/src/environment.rs

@@ -13,7 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::{Assignment, ConstraintUnsatisfied, Inject, LinearCombination, Mode, R1CS, Variable, witness_mode};
+use crate::{Assignment, ConstraintUnsatisfied, EjectedR1cs, Inject, LinearCombination, Mode, Variable, witness_mode};
 use snarkvm_curves::AffineCurve;
 use snarkvm_fields::traits::*;
 
@@ -185,11 +185,11 @@ pub trait Environment: 'static + Copy + Clone + fmt::Debug + fmt::Display + Eq +
         <Self::Network as console::Environment>::halt(message)
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn inject_r1cs(r1cs: R1CS<Self::BaseField>);
+    /// Injects the ejected R1CS circuit and restores the preserved environment state.
+    fn inject_r1cs(ejected: EjectedR1cs<Self::BaseField>);
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn eject_r1cs_and_reset() -> R1CS<Self::BaseField>;
+    /// Returns the R1CS circuit and preserved environment state, resetting the circuit.
+    fn eject_r1cs_and_reset() -> EjectedR1cs<Self::BaseField>;
 
     /// Returns the R1CS assignment of the circuit, resetting the circuit.
     fn eject_assignment_and_reset() -> Assignment<<Self::Network as console::Environment>::Field>;

circuit/environment/src/helpers/r1cs.rs

@@ -266,3 +266,50 @@ impl<F: PrimeField> Display for R1CS<F> {
         write!(f, "{output}")
     }
 }
+
+/// Captures the circuit environment state preserved across a temporary R1CS ejection.
+#[derive(Debug)]
+pub struct EjectedR1cs<F: PrimeField> {
+    pub(crate) r1cs: R1CS<F>,
+    pub(crate) variable_limit: Option<u64>,
+    pub(crate) constraint_limit: Option<u64>,
+    pub(crate) non_zero_limit: Option<(u64, u64, u64)>,
+    pub(crate) in_witness: bool,
+}
+
+impl<F: PrimeField> EjectedR1cs<F> {
+    /// Returns the ejected R1CS without restoring the environment state.
+    pub fn into_r1cs(self) -> R1CS<F> {
+        self.r1cs
+    }
+
+    /// Returns the preserved variable limit.
+    pub const fn variable_limit(&self) -> Option<u64> {
+        self.variable_limit
+    }
+
+    /// Returns the preserved constraint limit.
+    pub const fn constraint_limit(&self) -> Option<u64> {
+        self.constraint_limit
+    }
+
+    /// Returns the preserved non-zero limit.
+    pub const fn non_zero_limit(&self) -> Option<(u64, u64, u64)> {
+        self.non_zero_limit
+    }
+
+    /// Returns the preserved witness mode.
+    pub const fn in_witness(&self) -> bool {
+        self.in_witness
+    }
+
+    pub(crate) fn new(
+        r1cs: R1CS<F>,
+        variable_limit: Option<u64>,
+        constraint_limit: Option<u64>,
+        non_zero_limit: Option<(u64, u64, u64)>,
+        in_witness: bool,
+    ) -> Self {
+        Self { r1cs, variable_limit, constraint_limit, non_zero_limit, in_witness }
+    }
+}

circuit/environment/src/testnet_circuit.rs

@@ -296,8 +296,9 @@ impl Environment for TestnetCircuit {
         panic!("{}", &error)
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn inject_r1cs(r1cs: R1CS<Self::BaseField>) {
+    /// Injects the ejected R1CS circuit and restores the preserved environment state.
+    fn inject_r1cs(ejected: EjectedR1cs<Self::BaseField>) {
+        let EjectedR1cs { r1cs, variable_limit, constraint_limit, non_zero_limit, in_witness } = ejected;
         TESTNET_CIRCUIT.with(|circuit| {
             // Ensure the circuit is empty before injecting.
             assert_eq!(0, circuit.borrow().num_constants());
@@ -313,19 +314,31 @@ impl Environment for TestnetCircuit {
             assert_eq!(0, r1cs.num_private());
             assert_eq!(1, r1cs.num_variables());
             assert_eq!(0, r1cs.num_constraints());
-        })
+        });
+        // Restore the preserved limits.
+        Self::set_variable_limit(variable_limit);
+        Self::set_constraint_limit(constraint_limit);
+        Self::set_non_zero_limit(non_zero_limit);
+        // Restore the preserved witness mode.
+        IN_WITNESS.with(|in_witness_cell| in_witness_cell.replace(in_witness));
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn eject_r1cs_and_reset() -> R1CS<Self::BaseField> {
+    /// Returns the R1CS circuit and preserved environment state, resetting the circuit.
+    fn eject_r1cs_and_reset() -> EjectedR1cs<Self::BaseField> {
         TESTNET_CIRCUIT.with(|circuit| {
-            // Reset the witness mode.
-            IN_WITNESS.with(|in_witness| in_witness.replace(false));
-            // Reset the variable limit.
+            // Preserve the witness mode.
+            let in_witness = IN_WITNESS.with(|in_witness| {
+                let value = in_witness.get();
+                in_witness.replace(false);
+                value
+            });
+            // Preserve the limits.
+            let variable_limit = Self::get_variable_limit();
+            let constraint_limit = Self::get_constraint_limit();
+            let non_zero_limit = Self::get_non_zero_limit();
+            // Reset the limits for the temporary environment.
             Self::set_variable_limit(None);
-            // Reset the constraint limit.
             Self::set_constraint_limit(None);
-            // Reset the density limit.
             Self::set_non_zero_limit(None);
             // Eject the R1CS instance.
             let r1cs = circuit.replace(R1CS::<<Self as Environment>::BaseField>::new());
@@ -335,8 +348,8 @@ impl Environment for TestnetCircuit {
             assert_eq!(0, circuit.borrow().num_private());
             assert_eq!(1, circuit.borrow().num_variables());
             assert_eq!(0, circuit.borrow().num_constraints());
-            // Return the R1CS instance.
-            r1cs
+            // Return the ejected environment state.
+            EjectedR1cs::new(r1cs, variable_limit, constraint_limit, non_zero_limit, in_witness)
         })
     }
 

circuit/network/src/canary_v0.rs

@@ -44,7 +44,7 @@ use snarkvm_circuit_types::{
     Field,
     Group,
     Scalar,
-    environment::{Assignment, CanaryCircuit, R1CS, prelude::*},
+    environment::{Assignment, CanaryCircuit, EjectedR1cs, prelude::*},
 };
 
 use core::fmt;
@@ -538,13 +538,13 @@ impl Environment for AleoCanaryV0 {
         E::halt(message)
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn inject_r1cs(r1cs: R1CS<Self::BaseField>) {
-        E::inject_r1cs(r1cs)
+    /// Injects the ejected R1CS circuit and restores the preserved environment state.
+    fn inject_r1cs(ejected: EjectedR1cs<Self::BaseField>) {
+        E::inject_r1cs(ejected)
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn eject_r1cs_and_reset() -> R1CS<Self::BaseField> {
+    /// Returns the R1CS circuit and preserved environment state, resetting the circuit.
+    fn eject_r1cs_and_reset() -> EjectedR1cs<Self::BaseField> {
         E::eject_r1cs_and_reset()
     }
 

circuit/network/src/testnet_v0.rs

@@ -44,7 +44,7 @@ use snarkvm_circuit_types::{
     Field,
     Group,
     Scalar,
-    environment::{Assignment, R1CS, TestnetCircuit, prelude::*},
+    environment::{Assignment, EjectedR1cs, TestnetCircuit, prelude::*},
 };
 
 use core::fmt;
@@ -538,13 +538,13 @@ impl Environment for AleoTestnetV0 {
         E::halt(message)
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn inject_r1cs(r1cs: R1CS<Self::BaseField>) {
-        E::inject_r1cs(r1cs)
+    /// Injects the ejected R1CS circuit and restores the preserved environment state.
+    fn inject_r1cs(ejected: EjectedR1cs<Self::BaseField>) {
+        E::inject_r1cs(ejected)
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn eject_r1cs_and_reset() -> R1CS<Self::BaseField> {
+    /// Returns the R1CS circuit and preserved environment state, resetting the circuit.
+    fn eject_r1cs_and_reset() -> EjectedR1cs<Self::BaseField> {
         E::eject_r1cs_and_reset()
     }
 

circuit/network/src/v0.rs

@@ -44,7 +44,7 @@ use snarkvm_circuit_types::{
     Field,
     Group,
     Scalar,
-    environment::{Assignment, Circuit, R1CS, prelude::*},
+    environment::{Assignment, Circuit, EjectedR1cs, prelude::*},
 };
 
 use core::fmt;
@@ -538,13 +538,13 @@ impl Environment for AleoV0 {
         E::halt(message)
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn inject_r1cs(r1cs: R1CS<Self::BaseField>) {
-        E::inject_r1cs(r1cs)
+    /// Injects the ejected R1CS circuit and restores the preserved environment state.
+    fn inject_r1cs(ejected: EjectedR1cs<Self::BaseField>) {
+        E::inject_r1cs(ejected)
     }
 
-    /// Returns the R1CS circuit, resetting the circuit.
-    fn eject_r1cs_and_reset() -> R1CS<Self::BaseField> {
+    /// Returns the R1CS circuit and preserved environment state, resetting the circuit.
+    fn eject_r1cs_and_reset() -> EjectedR1cs<Self::BaseField> {
         E::eject_r1cs_and_reset()
     }