r1cs: single-variable allocation API + simpler API for assignments (#256)

This adds efficient and clean API for allocating one variable at a time.

Here is how it simplifies the code in Spacesuit and ZkVM: stellar/slingshot#211

Plus, the assignments are provided simply as `Option<Scalar>` and `Option<(Scalar,Scalar)>` to make code look nicer. The overhead in verifier of checking for `None` and not doing `Option::map` or something in the gadget code should be negligible.

Closes #206

Author: oleganza

src/r1cs/constraint_system.rs

@@ -39,16 +39,29 @@ pub trait ConstraintSystem {
         right: LinearCombination,
     ) -> (Variable, Variable, Variable);
 
+    /// Allocate a single variable.
+    ///
+    /// This either allocates a new multiplier and returns its `left` variable,
+    /// or returns a `right` variable of a multiplier previously allocated by this method.
+    /// The output of a multiplier is assigned on a even call, when `right` is assigned.
+    ///
+    /// When CS is committed at the end of the first or second phase, the half-assigned multiplier
+    /// has the `right` assigned to zero and all its variables committed.
+    ///
+    /// Returns unconstrained `Variable` for use in further constraints.
+    fn allocate(&mut self, assignment: Option<Scalar>) -> Result<Variable, R1CSError>;
+
     /// Allocate variables `left`, `right`, and `out`
     /// with the implicit constraint that
     /// ```text
     /// left * right = out
     /// ```
     ///
     /// Returns `(left, right, out)` for use in further constraints.
-    fn allocate<F>(&mut self, assign_fn: F) -> Result<(Variable, Variable, Variable), R1CSError>
-    where
-        F: FnOnce() -> Result<(Scalar, Scalar, Scalar), R1CSError>;
+    fn allocate_multiplier(
+        &mut self,
+        input_assignments: Option<(Scalar, Scalar)>,
+    ) -> Result<(Variable, Variable, Variable), R1CSError>;
 
     /// Enforce the explicit constraint that
     /// ```text

src/r1cs/prover.rs

@@ -43,6 +43,9 @@ pub struct Prover<'a, 'b> {
     /// This list holds closures that will be called in the second phase of the protocol,
     /// when non-randomized variables are committed.
     deferred_constraints: Vec<Box<Fn(&mut RandomizingProver<'a, 'b>) -> Result<(), R1CSError>>>,
+
+    /// Index of a pending multiplier that's not fully assigned yet.
+    pending_multiplier: Option<usize>,
 }
 
 /// Prover in the randomizing phase.
@@ -111,11 +114,33 @@ impl<'a, 'b> ConstraintSystem for Prover<'a, 'b> {
         (l_var, r_var, o_var)
     }
 
-    fn allocate<F>(&mut self, assign_fn: F) -> Result<(Variable, Variable, Variable), R1CSError>
-    where
-        F: FnOnce() -> Result<(Scalar, Scalar, Scalar), R1CSError>,
-    {
-        let (l, r, o) = assign_fn()?;
+    fn allocate(&mut self, assignment: Option<Scalar>) -> Result<Variable, R1CSError> {
+        let scalar = assignment.ok_or(R1CSError::MissingAssignment)?;
+
+        match self.pending_multiplier {
+            None => {
+                let i = self.a_L.len();
+                self.pending_multiplier = Some(i);
+                self.a_L.push(scalar);
+                self.a_R.push(Scalar::zero());
+                self.a_O.push(Scalar::zero());
+                Ok(Variable::MultiplierLeft(i))
+            }
+            Some(i) => {
+                self.pending_multiplier = None;
+                self.a_R[i] = scalar;
+                self.a_O[i] = self.a_L[i] * self.a_R[i];
+                Ok(Variable::MultiplierRight(i))
+            }
+        }
+    }
+
+    fn allocate_multiplier(
+        &mut self,
+        input_assignments: Option<(Scalar, Scalar)>,
+    ) -> Result<(Variable, Variable, Variable), R1CSError> {
+        let (l, r) = input_assignments.ok_or(R1CSError::MissingAssignment)?;
+        let o = l * r;
 
         // Create variables for l,r,o ...
         let l_var = Variable::MultiplierLeft(self.a_L.len());
@@ -155,11 +180,15 @@ impl<'a, 'b> ConstraintSystem for RandomizingProver<'a, 'b> {
         self.prover.multiply(left, right)
     }
 
-    fn allocate<F>(&mut self, assign_fn: F) -> Result<(Variable, Variable, Variable), R1CSError>
-    where
-        F: FnOnce() -> Result<(Scalar, Scalar, Scalar), R1CSError>,
-    {
-        self.prover.allocate(assign_fn)
+    fn allocate(&mut self, assignment: Option<Scalar>) -> Result<Variable, R1CSError> {
+        self.prover.allocate(assignment)
+    }
+
+    fn allocate_multiplier(
+        &mut self,
+        input_assignments: Option<(Scalar, Scalar)>,
+    ) -> Result<(Variable, Variable, Variable), R1CSError> {
+        self.prover.allocate_multiplier(input_assignments)
     }
 
     fn constrain(&mut self, lc: LinearCombination) {
@@ -219,6 +248,7 @@ impl<'a, 'b> Prover<'a, 'b> {
             a_R: Vec::new(),
             a_O: Vec::new(),
             deferred_constraints: Vec::new(),
+            pending_multiplier: None,
         }
     }
 
@@ -320,6 +350,9 @@ impl<'a, 'b> Prover<'a, 'b> {
     /// Calls all remembered callbacks with an API that
     /// allows generating challenge scalars.
     fn create_randomized_constraints(mut self) -> Result<Self, R1CSError> {
+        // Clear the pending multiplier (if any) because it was committed into A_L/A_R/S.
+        self.pending_multiplier = None;
+
         // Note: the wrapper could've used &mut instead of ownership,
         // but specifying lifetimes for boxed closures is not going to be nice,
         // so we move the self into wrapper and then move it back out afterwards.

src/r1cs/verifier.rs

@@ -42,6 +42,9 @@ pub struct Verifier<'a, 'b> {
     /// After that, the option will flip to None and additional calls to `randomize_constraints`
     /// will invoke closures immediately.
     deferred_constraints: Vec<Box<Fn(&mut RandomizingVerifier<'a, 'b>) -> Result<(), R1CSError>>>,
+
+    /// Index of a pending multiplier that's not fully assigned yet.
+    pending_multiplier: Option<usize>,
 }
 
 /// Verifier in the randomizing phase.
@@ -80,10 +83,25 @@ impl<'a, 'b> ConstraintSystem for Verifier<'a, 'b> {
         (l_var, r_var, o_var)
     }
 
-    fn allocate<F>(&mut self, _: F) -> Result<(Variable, Variable, Variable), R1CSError>
-    where
-        F: FnOnce() -> Result<(Scalar, Scalar, Scalar), R1CSError>,
-    {
+    fn allocate(&mut self, _: Option<Scalar>) -> Result<Variable, R1CSError> {
+        match self.pending_multiplier {
+            None => {
+                let i = self.num_vars;
+                self.num_vars += 1;
+                self.pending_multiplier = Some(i);
+                Ok(Variable::MultiplierLeft(i))
+            }
+            Some(i) => {
+                self.pending_multiplier = None;
+                Ok(Variable::MultiplierRight(i))
+            }
+        }
+    }
+
+    fn allocate_multiplier(
+        &mut self,
+        _: Option<(Scalar, Scalar)>,
+    ) -> Result<(Variable, Variable, Variable), R1CSError> {
         let var = self.num_vars;
         self.num_vars += 1;
 
@@ -122,11 +140,15 @@ impl<'a, 'b> ConstraintSystem for RandomizingVerifier<'a, 'b> {
         self.verifier.multiply(left, right)
     }
 
-    fn allocate<F>(&mut self, assign_fn: F) -> Result<(Variable, Variable, Variable), R1CSError>
-    where
-        F: FnOnce() -> Result<(Scalar, Scalar, Scalar), R1CSError>,
-    {
-        self.verifier.allocate(assign_fn)
+    fn allocate(&mut self, assignment: Option<Scalar>) -> Result<Variable, R1CSError> {
+        self.verifier.allocate(assignment)
+    }
+
+    fn allocate_multiplier(
+        &mut self,
+        input_assignments: Option<(Scalar, Scalar)>,
+    ) -> Result<(Variable, Variable, Variable), R1CSError> {
+        self.verifier.allocate_multiplier(input_assignments)
     }
 
     fn constrain(&mut self, lc: LinearCombination) {
@@ -194,6 +216,7 @@ impl<'a, 'b> Verifier<'a, 'b> {
             V: Vec::new(),
             constraints: Vec::new(),
             deferred_constraints: Vec::new(),
+            pending_multiplier: None,
         }
     }
 
@@ -279,6 +302,9 @@ impl<'a, 'b> Verifier<'a, 'b> {
     /// Calls all remembered callbacks with an API that
     /// allows generating challenge scalars.
     fn create_randomized_constraints(mut self) -> Result<Self, R1CSError> {
+        // Clear the pending multiplier (if any) because it was committed into A_L/A_R/S.
+        self.pending_multiplier = None;
+
         // Note: the wrapper could've used &mut instead of ownership,
         // but specifying lifetimes for boxed closures is not going to be nice,
         // so we move the self into wrapper and then move it back out afterwards.

tests/r1cs.rs

@@ -358,3 +358,92 @@ fn example_gadget_serialization_test() {
     // (3 + 4) * (6 + 1) != (40 + 10)
     assert!(example_gadget_roundtrip_serialization_helper(3, 4, 6, 1, 40, 10).is_err());
 }
+
+// Range Proof gadget
+
+/// Enforces that the quantity of v is in the range [0, 2^n).
+pub fn range_proof<CS: ConstraintSystem>(
+    cs: &mut CS,
+    mut v: LinearCombination,
+    v_assignment: Option<u64>,
+    n: usize,
+) -> Result<(), R1CSError> {
+    let mut exp_2 = Scalar::one();
+    for i in 0..n {
+        // Create low-level variables and add them to constraints
+        let (a, b, o) = cs.allocate_multiplier(v_assignment.map(|q| {
+            let bit: u64 = (q >> i) & 1;
+            ((1 - bit).into(), bit.into())
+        }))?;
+
+        // Enforce a * b = 0, so one of (a,b) is zero
+        cs.constrain(o.into());
+
+        // Enforce that a = 1 - b, so they both are 1 or 0.
+        cs.constrain(a + (b - 1u64));
+
+        // Add `-b_i*2^i` to the linear combination
+        // in order to form the following constraint by the end of the loop:
+        // v = Sum(b_i * 2^i, i = 0..n-1)
+        v = v - b * exp_2;
+
+        exp_2 = exp_2 + exp_2;
+    }
+
+    // Enforce that v = Sum(b_i * 2^i, i = 0..n-1)
+    cs.constrain(v);
+
+    Ok(())
+}
+
+#[test]
+fn range_proof_gadget() {
+    use rand::rngs::OsRng;
+    use rand::Rng;
+
+    let mut rng = OsRng::new().unwrap();
+    let m = 3; // number of values to test per `n`
+
+    for n in [2, 10, 32, 63].iter() {
+        let (min, max) = (0u64, ((1u128 << n) - 1) as u64);
+        let values: Vec<u64> = (0..m).map(|_| rng.gen_range(min, max)).collect();
+        for v in values {
+            assert!(range_proof_helper(v.into(), *n).is_ok());
+        }
+        assert!(range_proof_helper((max + 1).into(), *n).is_err());
+    }
+}
+
+fn range_proof_helper(v_val: u64, n: usize) -> Result<(), R1CSError> {
+    // Common
+    let pc_gens = PedersenGens::default();
+    let bp_gens = BulletproofGens::new(128, 1);
+
+    // Prover's scope
+    let (proof, commitment) = {
+        // Prover makes a `ConstraintSystem` instance representing a range proof gadget
+        let mut prover_transcript = Transcript::new(b"RangeProofTest");
+        let mut rng = rand::thread_rng();
+
+        let mut prover = Prover::new(&bp_gens, &pc_gens, &mut prover_transcript);
+
+        let (com, var) = prover.commit(v_val.into(), Scalar::random(&mut rng));
+        assert!(range_proof(&mut prover, var.into(), Some(v_val), n).is_ok());
+
+        let proof = prover.prove()?;
+
+        (proof, com)
+    };
+
+    // Verifier makes a `ConstraintSystem` instance representing a merge gadget
+    let mut verifier_transcript = Transcript::new(b"RangeProofTest");
+    let mut verifier = Verifier::new(&bp_gens, &pc_gens, &mut verifier_transcript);
+
+    let var = verifier.commit(commitment);
+
+    // Verifier adds constraints to the constraint system
+    assert!(range_proof(&mut verifier, var.into(), None, n).is_ok());
+
+    // Verifier verifies proof
+    Ok(verifier.verify(&proof)?)
+}