From f6f7e5519138f78a2f7ff37dcefade6eb175e445 Mon Sep 17 00:00:00 2001
From: Nicholas Ward <npward@berkeley.edu>
Date: Wed, 16 Feb 2022 11:31:26 -0800
Subject: [PATCH] windowed mul fixes......

---
 plonky2/src/gadgets/arithmetic.rs         | 10 ++++++++++
 plonky2/src/gadgets/curve_windowed_mul.rs | 21 ++++++++++++++++++++-
 plonky2/src/gadgets/nonnative.rs          | 12 ++++++++++++
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/plonky2/src/gadgets/arithmetic.rs b/plonky2/src/gadgets/arithmetic.rs
index 70cd24ad..b7df3726 100644
--- a/plonky2/src/gadgets/arithmetic.rs
+++ b/plonky2/src/gadgets/arithmetic.rs
@@ -326,6 +326,16 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         BoolTarget::new_unsafe(res)
     }
 
+    pub fn and(&mut self, b1: BoolTarget, b2: BoolTarget) -> BoolTarget {
+        BoolTarget::new_unsafe(self.mul(b1.target, b2.target))
+    }
+
+    pub fn _if(&mut self, b: BoolTarget, x: Target, y: Target) -> Target {
+        let not_b = self.not(b);
+        let maybe_x = self.mul(b.target, x);
+        self.mul_add(not_b.target, y, maybe_x)
+    }
+
     pub fn is_equal(&mut self, x: Target, y: Target) -> BoolTarget {
         let zero = self.zero();
 
diff --git a/plonky2/src/gadgets/curve_windowed_mul.rs b/plonky2/src/gadgets/curve_windowed_mul.rs
index 002b6ec4..b8b0f804 100644
--- a/plonky2/src/gadgets/curve_windowed_mul.rs
+++ b/plonky2/src/gadgets/curve_windowed_mul.rs
@@ -9,7 +9,7 @@ use crate::gadgets::biguint::BigUintTarget;
 use crate::gadgets::curve::AffinePointTarget;
 use crate::gadgets::nonnative::NonNativeTarget;
 use crate::hash::hash_types::RichField;
-use crate::iop::target::Target;
+use crate::iop::target::{BoolTarget, Target};
 use crate::plonk::circuit_builder::CircuitBuilder;
 
 const WINDOW_SIZE: usize = 4;
@@ -69,12 +69,25 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         AffinePointTarget { x, y }
     }
 
+    pub fn if_affine_point<C: Curve>(
+        &mut self,
+        b: BoolTarget,
+        p1: &AffinePointTarget<C>,
+        p2: &AffinePointTarget<C>,
+    ) -> AffinePointTarget<C> {
+        let new_x = self.if_nonnative(b, &p1.x, &p2.x);
+        let new_y = self.if_nonnative(b, &p1.y, &p2.y);
+        AffinePointTarget { x: new_x, y: new_y }
+    }
+
     pub fn curve_scalar_mul_windowed<C: Curve>(
         &mut self,
         p: &AffinePointTarget<C>,
         n: &NonNativeTarget<C::ScalarField>,
     ) -> AffinePointTarget<C> {
         let mut result = self.constant_affine_point(C::GENERATOR_AFFINE);
+        let mut to_subtract = self.constant_affine_point(C::GENERATOR_AFFINE);
+        let mut to_subtract_grows = self._true();
 
         let precomputation = self.precompute_window(p);
         let zero = self.zero();
@@ -83,10 +96,15 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let m = C::ScalarField::BITS / WINDOW_SIZE;
         for i in (0..m).rev() {
             result = self.curve_repeated_double(&result, WINDOW_SIZE);
+
+            let to_subtract_increased = self.curve_repeated_double(&to_subtract, WINDOW_SIZE);
+            to_subtract = self.if_affine_point(to_subtract_grows, &to_subtract_increased, &to_subtract);
+
             let window = windows[i];
 
             let to_add = self.random_access_curve_points(window, precomputation.clone());
             let is_zero = self.is_equal(window, zero);
+            to_subtract_grows = self.and(to_subtract_grows, is_zero);
             let should_add = self.not(is_zero);
             result = self.curve_conditional_add(&result, &to_add, should_add);
         }
@@ -173,6 +191,7 @@ mod tests {
 
         builder.connect_affine_point(&neg_five_g_expected, &neg_five_g_actual);
 
+        println!("NUM GATES: {}", builder.num_gates());
         let data = builder.build::<C>();
         let proof = data.prove(pw).unwrap();
 
diff --git a/plonky2/src/gadgets/nonnative.rs b/plonky2/src/gadgets/nonnative.rs
index 3f8d29e8..44969ac9 100644
--- a/plonky2/src/gadgets/nonnative.rs
+++ b/plonky2/src/gadgets/nonnative.rs
@@ -106,6 +106,18 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         }
     }
 
+    pub fn if_nonnative<FF: Field>(
+        &mut self,
+        b: BoolTarget,
+        x: &NonNativeTarget<FF>,
+        y: &NonNativeTarget<FF>,
+    ) -> NonNativeTarget<FF> {
+        let not_b = self.not(b);
+        let maybe_x = self.mul_nonnative_by_bool(x, b);
+        let maybe_y = self.mul_nonnative_by_bool(y, not_b);
+        self.add_nonnative(&maybe_x, &maybe_y)
+    }
+
     pub fn add_many_nonnative<FF: PrimeField>(
         &mut self,
         to_add: &[NonNativeTarget<FF>],