From 260bcb29c7097bfbdaa9c89faa58f68d5d936d4e Mon Sep 17 00:00:00 2001
From: ilyalesokhin-starkware <ilya@starkware.co>
Date: Thu, 28 Nov 2024 10:04:07 +0200
Subject: [PATCH] FRI: Enforce one column per size. (#905)

---
 crates/prover/src/core/fri.rs | 9 ++++++---
 crates/prover/src/lib.rs      | 1 +
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/crates/prover/src/core/fri.rs b/crates/prover/src/core/fri.rs
index 03dac47da..0f0a89b9b 100644
--- a/crates/prover/src/core/fri.rs
+++ b/crates/prover/src/core/fri.rs
@@ -136,7 +136,7 @@ pub struct FriProver<'a, B: FriOps + MerkleOps<MC::H>, MC: MerkleChannel> {
 impl<'a, B: FriOps + MerkleOps<MC::H>, MC: MerkleChannel> FriProver<'a, B, MC> {
     /// Commits to multiple circle polynomials.
     ///
-    /// `columns` must be provided in descending order by size.
+    /// `columns` must be provided in descending order by size with at most one column per size.
     ///
     /// This is a batched commitment that handles multiple mixed-degree polynomials, each
     /// evaluated over domains of varying sizes. Instead of combining these evaluations into
@@ -147,7 +147,7 @@ impl<'a, B: FriOps + MerkleOps<MC::H>, MC: MerkleChannel> FriProver<'a, B, MC> {
     /// # Panics
     ///
     /// Panics if:
-    /// * `columns` is empty or not sorted in ascending order by domain size.
+    /// * `columns` is empty or not sorted in descending order by domain size.
     /// * An evaluation is not from a sufficiently low degree circle polynomial.
     /// * An evaluation's domain is smaller than the last layer.
     /// * An evaluation's domain is not a canonic circle domain.
@@ -159,8 +159,11 @@ impl<'a, B: FriOps + MerkleOps<MC::H>, MC: MerkleChannel> FriProver<'a, B, MC> {
         twiddles: &TwiddleTree<B>,
     ) -> Self {
         assert!(!columns.is_empty(), "no columns");
-        assert!(columns.is_sorted_by_key(|e| Reverse(e.len())), "not sorted");
         assert!(columns.iter().all(|e| e.domain.is_canonic()), "not canonic");
+        assert!(
+            columns.array_windows().all(|[a, b]| a.len() > b.len()),
+            "column sizes not decreasing"
+        );
 
         let first_layer = Self::commit_first_layer(channel, columns);
         let (inner_layers, last_layer_evaluation) =
diff --git a/crates/prover/src/lib.rs b/crates/prover/src/lib.rs
index 49adff0f3..34a5c6701 100644
--- a/crates/prover/src/lib.rs
+++ b/crates/prover/src/lib.rs
@@ -3,6 +3,7 @@
     array_chunks,
     array_methods,
     array_try_from_fn,
+    array_windows,
     assert_matches,
     exact_size_is_empty,
     generic_const_exprs,