From 7ac4ae9f94aa16b8a5db42af2bd0d181d8bd4fc8 Mon Sep 17 00:00:00 2001
From: Philippe Laferriere <plafer@protonmail.com>
Date: Sat, 15 Feb 2025 10:24:25 -0500
Subject: [PATCH] fix: fix FRI folded layers domain

---
 fri/src/prover/mod.rs   | 20 ++++++++++++++------
 fri/src/verifier/mod.rs |  4 +++-
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/fri/src/prover/mod.rs b/fri/src/prover/mod.rs
index bf6d8da41..71bb2f001 100644
--- a/fri/src/prover/mod.rs
+++ b/fri/src/prover/mod.rs
@@ -184,14 +184,17 @@ where
 
         // reduce the degree by folding_factor at each iteration until the remaining polynomial
         // has small enough degree
+        let mut domain_offset = self.options.domain_offset();
         for _ in 0..self.options.num_fri_layers(evaluations.len()) {
             match self.folding_factor() {
-                2 => self.build_layer::<2>(channel, &mut evaluations),
-                4 => self.build_layer::<4>(channel, &mut evaluations),
-                8 => self.build_layer::<8>(channel, &mut evaluations),
-                16 => self.build_layer::<16>(channel, &mut evaluations),
+                2 => self.build_layer::<2>(channel, &mut evaluations, domain_offset),
+                4 => self.build_layer::<4>(channel, &mut evaluations, domain_offset),
+                8 => self.build_layer::<8>(channel, &mut evaluations, domain_offset),
+                16 => self.build_layer::<16>(channel, &mut evaluations, domain_offset),
                 _ => unimplemented!("folding factor {} is not supported", self.folding_factor()),
             }
+
+            domain_offset = domain_offset.exp_vartime((self.folding_factor() as u32).into());
         }
 
         self.set_remainder(channel, &mut evaluations);
@@ -199,7 +202,12 @@ where
 
     /// Builds a single FRI layer by first committing to the `evaluations`, then drawing a random
     /// alpha from the channel and use it to perform degree-respecting projection.
-    fn build_layer<const N: usize>(&mut self, channel: &mut C, evaluations: &mut Vec<E>) {
+    fn build_layer<const N: usize>(
+        &mut self,
+        channel: &mut C,
+        evaluations: &mut Vec<E>,
+        domain_offset: E::BaseField,
+    ) {
         // commit to the evaluations at the current layer; we do this by first transposing the
         // evaluations into a matrix of N columns, then hashing each row into a digest, and finally
         // commiting to vector of these digests; we do this so that we could de-commit to N values
@@ -213,7 +221,7 @@ where
         // draw a pseudo-random coefficient from the channel, and use it in degree-respecting
         // projection to reduce the degree of evaluations by N
         let alpha = channel.draw_fri_alpha();
-        *evaluations = apply_drp(&transposed_evaluations, self.domain_offset(), alpha);
+        *evaluations = apply_drp(&transposed_evaluations, domain_offset, alpha);
         self.layers.push(FriLayer {
             commitment: evaluation_vector_commitment,
             evaluations: flatten_vector_elements(transposed_evaluations),
diff --git a/fri/src/verifier/mod.rs b/fri/src/verifier/mod.rs
index 59395b2b7..0e2eedc2b 100644
--- a/fri/src/verifier/mod.rs
+++ b/fri/src/verifier/mod.rs
@@ -250,6 +250,7 @@ where
         let mut max_degree_plus_1 = self.max_poly_degree + 1;
         let mut positions = positions.to_vec();
         let mut evaluations = evaluations.to_vec();
+        let mut domain_offset = self.options.domain_offset();
 
         for depth in 0..self.options.num_fri_layers(self.domain_size) {
             // determine which evaluations were queried in the folded layer
@@ -275,7 +276,7 @@ where
             // build a set of x coordinates for each row polynomial
             #[rustfmt::skip]
             let xs = folded_positions.iter().map(|&i| {
-                let xe = domain_generator.exp_vartime((i as u64).into()) * self.options.domain_offset();
+                let xe = domain_generator.exp_vartime((i as u64).into()) * domain_offset;
                 folding_roots.iter()
                     .map(|&r| E::from(xe * r))
                     .collect::<Vec<_>>().try_into().unwrap()
@@ -302,6 +303,7 @@ where
             max_degree_plus_1 /= N;
             domain_size /= N;
             mem::swap(&mut positions, &mut folded_positions);
+            domain_offset = domain_offset.exp_vartime((N as u32).into());
         }
 
         // 2 ----- verify the remainder polynomial of the FRI proof -------------------------------