From 1ebeab2c3a1cf83e3d8d1075dca4e3c65666585f Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Fri, 11 Jun 2021 10:27:03 +0200
Subject: [PATCH] Implement Frobenius optimization discussed in #61 comments to
 avoid calling the Frobenius for every wires.

---
 src/field/extension_field/mod.rs    |  8 ++++++++
 src/field/extension_field/target.rs |  9 +++++++++
 src/field/field.rs                  | 16 ++++++++++++++++
 src/fri/verifier.rs                 |  9 ++++-----
 4 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/src/field/extension_field/mod.rs b/src/field/extension_field/mod.rs
index d706a341..525d8ecc 100644
--- a/src/field/extension_field/mod.rs
+++ b/src/field/extension_field/mod.rs
@@ -28,6 +28,14 @@ pub trait OEF<const D: usize>: FieldExtension<D> {
 
         Self::from_basefield_array(res)
     }
+
+    /// Repeated Frobenius automorphisms: x -> x^(p^k).
+    // TODO: Implement this. Is basically the same as `frobenius` above, but using
+    // `z = W^floor(j*p^k/D)`. I'm not sure there is a closed form for these so
+    // might require to hardcode them.
+    fn repeated_frobenius(&self, k: usize) -> Self {
+        todo!()
+    }
 }
 
 impl<F: Field> OEF<1> for F {
diff --git a/src/field/extension_field/target.rs b/src/field/extension_field/target.rs
index 7efd0f3b..aa5bce7f 100644
--- a/src/field/extension_field/target.rs
+++ b/src/field/extension_field/target.rs
@@ -31,6 +31,15 @@ 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>>(
+        &self,
+        k: usize,
+        builder: &mut CircuitBuilder<F, D>,
+    ) -> Self {
+        todo!()
+    }
+
     pub fn from_range(gate: usize, range: Range<usize>) -> Self {
         debug_assert_eq!(range.end - range.start, D);
         Target::wires_from_range(gate, range).try_into().unwrap()
diff --git a/src/field/field.rs b/src/field/field.rs
index 3c6f0e47..44a20a58 100644
--- a/src/field/field.rs
+++ b/src/field/field.rs
@@ -7,6 +7,7 @@ use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssi
 use num::Integer;
 use rand::Rng;
 
+use crate::field::extension_field::OEF;
 use crate::util::bits_u64;
 
 /// A finite field with prime order less than 2^64.
@@ -283,3 +284,18 @@ impl<F: Field> Iterator for Powers<F> {
         Some(result)
     }
 }
+
+impl<F: Field> Powers<F> {
+    /// Apply the Frobenius automorphism `k` times.
+    // TODO: Use `OEF::repeated_frobenius` when it is implemented.
+    pub fn repeated_frobenius<const D: usize>(self, k: usize) -> Self
+    where
+        F: OEF<D>,
+    {
+        let Self { base, current } = self;
+        Self {
+            base: (0..k).fold(base, |acc, _| acc.frobenius()),
+            current: (0..k).fold(current, |acc, _| acc.frobenius()),
+        }
+    }
+}
diff --git a/src/fri/verifier.rs b/src/fri/verifier.rs
index 6275ada6..27fdf53e 100644
--- a/src/fri/verifier.rs
+++ b/src/fri/verifier.rs
@@ -199,11 +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 wire_evals_frob = os.wires.iter().map(|e| e.frobenius()).collect::<Vec<_>>();
-    let wires_interpol = interpolant(&[
-        (zeta, reduce_with_iter(&os.wires, alpha_powers.clone())),
-        (zeta_frob, reduce_with_iter(&wire_evals_frob, alpha_powers)),
-    ]);
+    let ev_zeta = 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 numerator = ev - wires_interpol.eval(subgroup_x);
     let denominator = (subgroup_x - zeta) * (subgroup_x - zeta_frob);
     sum += numerator / denominator;