From b60a02862bba29cb180951aae81990f3843c7f51 Mon Sep 17 00:00:00 2001
From: Sai Deng <deng.sai@gmail.com>
Date: Tue, 22 Oct 2024 15:28:05 -0700
Subject: [PATCH] more fixes

---
 zero/src/prover_state/mod.rs | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/zero/src/prover_state/mod.rs b/zero/src/prover_state/mod.rs
index 8cdfb45ce..376d4c90a 100644
--- a/zero/src/prover_state/mod.rs
+++ b/zero/src/prover_state/mod.rs
@@ -15,6 +15,7 @@ use std::borrow::Borrow;
 use std::{fmt::Display, sync::OnceLock};
 
 use clap::ValueEnum;
+use evm_arithmetization::testing_utils::TEST_STARK_CONFIG;
 use evm_arithmetization::{
     fixed_recursive_verifier::ProverOutputData, prover::prove, AllProof, AllRecursiveCircuits,
     AllStark, GenerationSegmentData, RecursiveCircuitsForTableSize, StarkConfig,
@@ -207,20 +208,20 @@ impl ProverStateManager {
         &self,
         input: TrimmedGenerationInputs,
         segment_data: &mut GenerationSegmentData,
+        config: &StarkConfig,
     ) -> anyhow::Result<ProofWithPublicValues> {
-        let config = StarkConfig::standard_fast_config();
         let all_stark = AllStark::default();
 
         let all_proof = prove(
             &all_stark,
-            &config,
+            config,
             input,
             segment_data,
             &mut TimingTree::default(),
             None,
         )?;
 
-        let table_circuits = self.load_table_circuits(&config, &all_proof)?;
+        let table_circuits = self.load_table_circuits(config, &all_proof)?;
 
         let proof_with_pvs =
             p_state()
@@ -236,10 +237,11 @@ impl ProverStateManager {
         &self,
         input: TrimmedGenerationInputs,
         segment_data: &mut GenerationSegmentData,
+        config: &StarkConfig,
     ) -> anyhow::Result<ProofWithPublicValues> {
         let p_out = p_state().state.prove_segment(
             &AllStark::default(),
-            &StarkConfig::standard_fast_config(),
+            config,
             input,
             segment_data,
             &mut TimingTree::default(),
@@ -269,15 +271,20 @@ impl ProverStateManager {
         input: (TrimmedGenerationInputs, GenerationSegmentData),
     ) -> anyhow::Result<ProofWithPublicValues> {
         let (generation_inputs, mut segment_data) = input;
+        let config = if self.circuit_config.use_test_config {
+            TEST_STARK_CONFIG
+        } else {
+            StarkConfig::standard_fast_config()
+        };
 
         match self.persistence {
             CircuitPersistence::None | CircuitPersistence::Disk(TableLoadStrategy::Monolithic) => {
                 info!("using monolithic circuit {:?}", self);
-                self.segment_proof_monolithic(generation_inputs, &mut segment_data)
+                self.segment_proof_monolithic(generation_inputs, &mut segment_data, &config)
             }
             CircuitPersistence::Disk(TableLoadStrategy::OnDemand) => {
                 info!("using on demand circuit {:?}", self);
-                self.segment_proof_on_demand(generation_inputs, &mut segment_data)
+                self.segment_proof_on_demand(generation_inputs, &mut segment_data, &config)
             }
         }
     }