From 1f50d52a5fca799440c0c1486f679d9329fb779c Mon Sep 17 00:00:00 2001
From: winderica <winderica@gmail.com>
Date: Tue, 2 Apr 2024 05:54:21 +0800
Subject: [PATCH] Add convenient method for variable allocation with inferred
 mode

This closes #141.
---
 src/alloc.rs | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/src/alloc.rs b/src/alloc.rs
index ccce47f..d961405 100644
--- a/src/alloc.rs
+++ b/src/alloc.rs
@@ -76,6 +76,53 @@ pub trait AllocVar<V: ?Sized, F: Field>: Sized {
     ) -> Result<Self, SynthesisError> {
         Self::new_variable(cs, f, AllocationMode::Witness)
     }
+
+    /// Allocates a new constant or private witness of type `Self` in the
+    /// `ConstraintSystem` `cs` with the allocation mode inferred from `cs`.
+    /// A constant is allocated if `cs` is `None`, and a private witness is
+    /// allocated otherwise.
+    ///
+    /// A common use case is the creation of non-deterministic advice (a.k.a.
+    /// hints) in the circuit, where this method can avoid boilerplate code
+    /// while allowing optimization on circuit size.
+    ///
+    /// For example, to compute `x_var / y_var` where `y_var` is a non-zero
+    /// variable, one can write:
+    /// ```
+    /// use ark_ff::PrimeField;
+    /// use ark_r1cs_std::{alloc::AllocVar, fields::{fp::FpVar, FieldVar}, R1CSVar};
+    /// use ark_relations::r1cs::SynthesisError;
+    ///
+    /// fn div<F: PrimeField>(x_var: &FpVar<F>, y_var: &FpVar<F>) -> Result<FpVar<F>, SynthesisError> {
+    ///   let cs = x_var.cs().or(y_var.cs());
+    ///   let z_var = FpVar::new_variable_with_inferred_mode(cs, || Ok(x_var.value()? / y_var.value()?))?;
+    ///   z_var.mul_equals(y_var, x_var)?;
+    ///   Ok(z_var)
+    /// }
+    /// ```
+    /// In this example, if either `x_var` or `y_var` is a witness variable,
+    /// then `z_var` is also a witness variable. On the other hand, `z_var`
+    /// is a constant if both `x_var` and `y_var` are constants (i.e., `cs`
+    /// is `None`), and future operations on `z_var` do not generate any
+    /// constraints.
+    ///
+    /// (Note that we use division as an example for simplicity. You may
+    /// call `x_var.mul_by_inverse(y_var)?` directly, which internally works
+    /// similarly to the above code.)
+    #[tracing::instrument(target = "r1cs", skip(cs, f))]
+    fn new_variable_with_inferred_mode<T: Borrow<V>>(
+        cs: impl Into<Namespace<F>>,
+        f: impl FnOnce() -> Result<T, SynthesisError>,
+    ) -> Result<Self, SynthesisError> {
+        let ns: Namespace<F> = cs.into();
+        let cs = ns.cs();
+        let mode = if cs.is_none() {
+            AllocationMode::Constant
+        } else {
+            AllocationMode::Witness
+        };
+        Self::new_variable(cs, f, mode)
+    }
 }
 
 /// This blanket implementation just allocates variables in `Self`