From 4b1f368e890df5165ee77fbcd7a50accb26f6197 Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Fri, 11 Jun 2021 14:16:40 +0200
Subject: [PATCH] Use Frobenius optimization in the circuit

---
 src/field/extension_field/target.rs |  3 +--
 src/fri/recursive_verifier.rs       | 26 ++++++++++++--------------
 src/fri/verifier.rs                 |  6 +++---
 src/gadgets/arithmetic.rs           | 12 ++++++++++++
 4 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/src/field/extension_field/target.rs b/src/field/extension_field/target.rs
index 2123079d..2e743098 100644
--- a/src/field/extension_field/target.rs
+++ b/src/field/extension_field/target.rs
@@ -32,8 +32,7 @@ impl<const D: usize> ExtensionTarget<D> {
         res.try_into().unwrap()
     }
 
-    // TODO: Implement this. See comment in `OEF::repeated_frobenius`.
-    fn repeated_frobenius<F: Extendable<D>>(
+    pub fn repeated_frobenius<F: Extendable<D>>(
         &self,
         k: usize,
         builder: &mut CircuitBuilder<F, D>,
diff --git a/src/fri/recursive_verifier.rs b/src/fri/recursive_verifier.rs
index 4f8c3864..3a4332cf 100644
--- a/src/fri/recursive_verifier.rs
+++ b/src/fri/recursive_verifier.rs
@@ -225,22 +225,20 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             ev = self.mul_add_extension(a, e, ev);
         }
         let zeta_frob = zeta.frobenius(self);
-        let wire_evals_frob = os
+        let wire_eval = os.wires.iter().fold(self.zero_extension(), |acc, &w| {
+            let a = alpha_powers.next(self);
+            self.mul_add_extension(a, w, acc)
+        });
+        let mut alpha_powers_frob = alpha_powers.repeated_frobenius(D - 1, self);
+        let wire_eval_frob = os
             .wires
             .iter()
-            .map(|e| e.frobenius(self))
-            .collect::<Vec<_>>();
-        let mut ev_zeta = self.zero_extension();
-        for &t in &os.wires {
-            let a = alpha_powers.next(self);
-            ev_zeta = self.mul_add_extension(a, t, ev_zeta);
-        }
-        let mut ev_zeta_frob = self.zero_extension();
-        for &t in &wire_evals_frob {
-            let a = alpha_powers.next(self);
-            ev_zeta_right = self.mul_add_extension(a, t, ev_zeta);
-        }
-        let wires_interpol = self.interpolate2([(zeta, ev_zeta), (zeta_frob, ev_zeta_frob)]);
+            .fold(self.zero_extension(), |acc, &w| {
+                let a = alpha_powers_frob.next(self);
+                self.mul_add_extension(a, w, acc)
+            })
+            .frobenius(self);
+        let wires_interpol = self.interpolate2([(zeta, wire_eval), (zeta_frob, wire_eval_frob)]);
         let interpol_val = wires_interpol.eval(self, subgroup_x);
         let numerator = self.sub_extension(ev, interpol_val);
         let vanish_frob = self.sub_extension(subgroup_x, zeta_frob);
diff --git a/src/fri/verifier.rs b/src/fri/verifier.rs
index a3f3864a..104deaa3 100644
--- a/src/fri/verifier.rs
+++ b/src/fri/verifier.rs
@@ -199,10 +199,10 @@ fn fri_combine_initial<F: Field + Extendable<D>, const D: usize>(
         .map(|(&e, a)| a * e.into())
         .sum();
     let zeta_frob = zeta.frobenius();
-    let ev_zeta = reduce_with_iter(&os.wires, alpha_powers.clone());
+    let wire_eval = reduce_with_iter(&os.wires, alpha_powers.clone());
     let mut alpha_powers_frob = alpha_powers.repeated_frobenius(D - 1);
-    let ev_zeta_frob = reduce_with_iter(&os.wires, alpha_powers_frob).frobenius();
-    let wires_interpol = interpolant(&[(zeta, ev_zeta), (zeta_frob, ev_zeta_frob)]);
+    let wire_eval_frob = reduce_with_iter(&os.wires, alpha_powers_frob).frobenius();
+    let wires_interpol = interpolant(&[(zeta, wire_eval), (zeta_frob, wire_eval_frob)]);
     let numerator = ev - wires_interpol.eval(subgroup_x);
     let denominator = (subgroup_x - zeta) * (subgroup_x - zeta_frob);
     sum += numerator / denominator;
diff --git a/src/gadgets/arithmetic.rs b/src/gadgets/arithmetic.rs
index 9c35d251..c6c3b2b7 100644
--- a/src/gadgets/arithmetic.rs
+++ b/src/gadgets/arithmetic.rs
@@ -338,6 +338,18 @@ impl<const D: usize> PowersTarget<D> {
         self.current = builder.mul_extension(self.base, self.current);
         result
     }
+
+    pub fn repeated_frobenius<F: Extendable<D>>(
+        self,
+        k: usize,
+        builder: &mut CircuitBuilder<F, D>,
+    ) -> Self {
+        let Self { base, current } = self;
+        Self {
+            base: base.repeated_frobenius(k, builder),
+            current: current.repeated_frobenius(k, builder),
+        }
+    }
 }
 
 impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {