From a3c23aa7a218a7c6c02e42f0f4eb4fd7c25f2bc9 Mon Sep 17 00:00:00 2001
From: hratoanina <hratoanina@polygon.technology>
Date: Fri, 19 Jul 2024 00:02:45 -0400
Subject: [PATCH] Fix fri instances

---
 evm_arithmetization/src/prover.rs            | 62 ++++++++++++++++++--
 evm_arithmetization/tests/simple_transfer.rs |  3 +-
 2 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/evm_arithmetization/src/prover.rs b/evm_arithmetization/src/prover.rs
index 3f25b44eb..c4eaaaf6f 100644
--- a/evm_arithmetization/src/prover.rs
+++ b/evm_arithmetization/src/prover.rs
@@ -15,7 +15,7 @@ use plonky2::field::polynomial::{PolynomialCoeffs, PolynomialValues};
 use plonky2::field::types::Field;
 use plonky2::fri::oracle::PolynomialBatch;
 use plonky2::fri::reduction_strategies::FriReductionStrategy;
-use plonky2::fri::structure::{FriInstanceInfo, FriOpeningBatch};
+use plonky2::fri::structure::{FriInstanceInfo, FriOpeningBatch, FriOracleInfo};
 use plonky2::fri::FriConfig;
 use plonky2::hash::hash_types::RichField;
 use plonky2::hash::merkle_tree::MerkleCap;
@@ -87,7 +87,7 @@ impl GenerationSegmentData {
     }
 }
 
-pub(crate) fn zkevm_fast_config() -> StarkConfig {
+pub fn zkevm_fast_config() -> StarkConfig {
     StarkConfig {
         security_bits: 100,
         num_challenges: 2,
@@ -515,8 +515,10 @@ where
         &quotient_commitment,
     ];
 
+    let mut degree_bits_squashed = Table::all_degree_logs().to_vec();
+    degree_bits_squashed.dedup();
     let opening_proof = BatchFriOracle::prove_openings(
-        &Table::all_degree_logs(),
+        &degree_bits_squashed,
         &all_fri_instances,
         &initial_merkle_trees,
         &mut challenger,
@@ -1147,10 +1149,60 @@ where
         config,
     ));
 
-    Table::all_sorted()
+    let res_sorted: Vec<_> = Table::all_sorted()
         .iter()
         .map(|&table| res[*table].clone())
-        .collect()
+        .collect();
+
+    let mut squashed_res = Vec::new();
+    let mut i = 0;
+    let mut current_instance = FriInstanceInfo {
+        oracles: vec![
+            FriOracleInfo {
+                num_polys: 0,
+                blinding: false,
+            },
+            FriOracleInfo {
+                num_polys: 0,
+                blinding: false,
+            },
+            FriOracleInfo {
+                num_polys: 0,
+                blinding: false,
+            },
+        ],
+        batches: vec![],
+    };
+
+    while i < NUM_TABLES {
+        let instance = &res_sorted[i];
+        for (k, oracle) in instance.oracles.iter().enumerate() {
+            current_instance.oracles[k].num_polys += oracle.num_polys;
+        }
+        current_instance.batches.extend(instance.batches.clone());
+
+        if i == NUM_TABLES - 1 || Table::all_degree_logs()[i + 1] < Table::all_degree_logs()[i] {
+            squashed_res.push(current_instance.clone());
+            current_instance.oracles = vec![
+                FriOracleInfo {
+                    num_polys: 0,
+                    blinding: false,
+                },
+                FriOracleInfo {
+                    num_polys: 0,
+                    blinding: false,
+                },
+                FriOracleInfo {
+                    num_polys: 0,
+                    blinding: false,
+                },
+            ];
+            current_instance.batches = vec![];
+        }
+        i += 1;
+    }
+
+    squashed_res
 }
 
 fn all_openings_single_stark<F, C, S, const D: usize>(
diff --git a/evm_arithmetization/tests/simple_transfer.rs b/evm_arithmetization/tests/simple_transfer.rs
index 7e3e00b8e..c68ed0976 100644
--- a/evm_arithmetization/tests/simple_transfer.rs
+++ b/evm_arithmetization/tests/simple_transfer.rs
@@ -8,6 +8,7 @@ use evm_arithmetization::generation::mpt::{AccountRlp, LegacyReceiptRlp};
 use evm_arithmetization::generation::{GenerationInputs, TrieInputs};
 use evm_arithmetization::proof::{BlockHashes, BlockMetadata, TrieRoots};
 use evm_arithmetization::prover::testing::{prove_all_segments, prove_all_segments_batch};
+use evm_arithmetization::prover::zkevm_fast_config;
 use evm_arithmetization::verifier::testing::verify_all_proofs;
 use evm_arithmetization::{AllStark, Node, StarkConfig};
 use hex_literal::hex;
@@ -31,7 +32,7 @@ fn test_simple_transfer() -> anyhow::Result<()> {
     init_logger();
 
     let all_stark = AllStark::<F, D>::default();
-    let config = StarkConfig::standard_fast_config();
+    let config = zkevm_fast_config();
 
     let beneficiary = hex!("deadbeefdeadbeefdeadbeefdeadbeefdeadbeef");
     let sender = hex!("2c7536e3605d9c16a7a3d7b1898e529396a65c23");