chore: sync stark-backend commit

openvm-org · Jan 24, 2025 · 2d3e8b6 · 2d3e8b6
1 parent 1f7f905
commit 2d3e8b6
Show file tree

Hide file tree

Showing 35 changed files with 207 additions and 182 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -107,8 +107,8 @@ lto = "thin"
 
 [workspace.dependencies]
 # Stark Backend
-openvm-stark-backend = { git = "https://github.com/openvm-org/stark-backend.git", rev = "cfa33b", default-features = false }
-openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", rev = "cfa33b", default-features = false }
+openvm-stark-backend = { git = "https://github.com/openvm-org/stark-backend.git", rev = "495692", default-features = false }
+openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", rev = "495692", default-features = false }
 
 # OpenVM
 openvm-sdk = { path = "crates/sdk", default-features = false }

diff --git a/crates/vm/src/arch/extensions.rs b/crates/vm/src/arch/extensions.rs
@@ -26,7 +26,7 @@ use openvm_stark_backend::{
     p3_matrix::Matrix,
     prover::types::{AirProofInput, CommittedTraceData, ProofInput},
     rap::AnyRap,
-    Chip, ChipUsageGetter, Stateful,
+    AirRef, Chip, ChipUsageGetter, Stateful,
 };
 use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
@@ -982,7 +982,7 @@ impl<F: PrimeField32, E, P> VmChipComplex<F, E, P> {
             .collect()
     }
 
-    pub(crate) fn airs<SC: StarkGenericConfig>(&self) -> Vec<Arc<dyn AnyRap<SC>>>
+    pub fn airs<SC: StarkGenericConfig>(&self) -> Vec<AirRef<SC>>
     where
         Domain<SC>: PolynomialSpace<Val = F>,
         E: Chip<SC>,

diff --git a/crates/vm/src/arch/integration_api.rs b/crates/vm/src/arch/integration_api.rs
@@ -300,7 +300,6 @@ where
     }
 
     fn generate_air_proof_input(self) -> AirProofInput<SC> {
-        let air = self.air();
         let num_records = self.records.len();
         let height = next_power_of_two_or_zero(num_records);
         let core_width = self.core.air().width();
@@ -325,7 +324,7 @@ where
         let mut trace = RowMajorMatrix::new(values, width);
         self.core.finalize(&mut trace, num_records);
 
-        AirProofInput::simple(air, trace, self.core.generate_public_values())
+        AirProofInput::simple(trace, self.core.generate_public_values())
     }
 }
 

diff --git a/crates/vm/src/arch/testing/execution/mod.rs b/crates/vm/src/arch/testing/execution/mod.rs
@@ -58,7 +58,6 @@ where
     }
 
     fn generate_air_proof_input(self) -> AirProofInput<SC> {
-        let air = self.air();
         let height = self.records.len().next_power_of_two();
         let width = self.trace_width();
         let mut values = Val::<SC>::zero_vec(height * width);
@@ -67,7 +66,7 @@ where
         for (row, record) in values.chunks_mut(width).zip(self.records) {
             *row.borrow_mut() = record;
         }
-        AirProofInput::simple_no_pis(air, RowMajorMatrix::new(values, width))
+        AirProofInput::simple_no_pis(RowMajorMatrix::new(values, width))
     }
 }
 impl<F: Field> ChipUsageGetter for ExecutionTester<F> {

diff --git a/crates/vm/src/arch/testing/memory/mod.rs b/crates/vm/src/arch/testing/memory/mod.rs
@@ -85,7 +85,6 @@ where
         let offline_memory = self.controller.borrow().offline_memory();
         let offline_memory = offline_memory.lock().unwrap();
 
-        let air = self.air();
         let height = self.records.len().next_power_of_two();
         let width = self.trace_width();
         let mut values = Val::<SC>::zero_vec(2 * height * width);
@@ -113,7 +112,7 @@ where
             row.timestamp = Val::<SC>::from_canonical_u32(record.timestamp);
             row.count = Val::<SC>::ONE;
         }
-        AirProofInput::simple_no_pis(air, RowMajorMatrix::new(values, width))
+        AirProofInput::simple_no_pis(RowMajorMatrix::new(values, width))
     }
 }
 

diff --git a/crates/vm/src/arch/testing/mod.rs b/crates/vm/src/arch/testing/mod.rs
@@ -1,5 +1,6 @@
 use std::{
     cell::RefCell,
+    iter::zip,
     rc::Rc,
     sync::{Arc, Mutex},
 };
@@ -15,7 +16,7 @@ use openvm_stark_backend::{
     p3_matrix::dense::{DenseMatrix, RowMajorMatrix},
     prover::types::AirProofInput,
     verifier::VerificationError,
-    Chip,
+    AirRef, Chip,
 };
 use openvm_stark_sdk::{
     config::{
@@ -295,7 +296,7 @@ impl<F: PrimeField32> Default for VmChipTestBuilder<F> {
 
 pub struct VmChipTester<SC: StarkGenericConfig> {
     pub memory: Option<MemoryTester<Val<SC>>>,
-    pub air_proof_inputs: Vec<AirProofInput<SC>>,
+    pub air_proof_inputs: Vec<(AirRef<SC>, AirProofInput<SC>)>,
 }
 
 impl<SC: StarkGenericConfig> Default for VmChipTester<SC> {
@@ -313,12 +314,10 @@ where
 {
     pub fn load<C: Chip<SC>>(mut self, chip: C) -> Self {
         if chip.current_trace_height() > 0 {
+            let air = chip.air();
             let air_proof_input = chip.generate_air_proof_input();
-            tracing::debug!(
-                "Generated air proof input for {}",
-                air_proof_input.air.name()
-            );
-            self.air_proof_inputs.push(air_proof_input);
+            tracing::debug!("Generated air proof input for {}", air.name());
+            self.air_proof_inputs.push((air, air_proof_input));
         }
 
         self
@@ -330,22 +329,24 @@ where
             let range_checker = memory_controller.borrow().range_checker.clone();
             self = self.load(memory_tester); // dummy memory interactions
             {
+                let airs = memory_controller.borrow().airs();
                 let air_proof_inputs = Rc::try_unwrap(memory_controller)
                     .unwrap_or_else(|_| panic!("Memory controller was not dropped"))
                     .into_inner()
                     .generate_air_proof_inputs();
                 self.air_proof_inputs.extend(
-                    air_proof_inputs
-                        .into_iter()
-                        .filter(|api| api.main_trace_height() > 0),
+                    zip(airs, air_proof_inputs).filter(|(_, input)| input.main_trace_height() > 0),
                 );
             }
             self = self.load(range_checker); // this must be last because other trace generation mutates its state
         }
         self
     }
 
-    pub fn load_air_proof_input(mut self, air_proof_input: AirProofInput<SC>) -> Self {
+    pub fn load_air_proof_input(
+        mut self,
+        air_proof_input: (AirRef<SC>, AirProofInput<SC>),
+    ) -> Self {
         self.air_proof_inputs.push(air_proof_input);
         self
     }
@@ -355,20 +356,22 @@ where
         chip: C,
         trace: RowMajorMatrix<Val<SC>>,
     ) -> Self {
+        let air = chip.air();
         let mut air_proof_input = chip.generate_air_proof_input();
         air_proof_input.raw.common_main = Some(trace);
-        self.air_proof_inputs.push(air_proof_input);
+        self.air_proof_inputs.push((air, air_proof_input));
         self
     }
 
     pub fn load_and_prank_trace<C: Chip<SC>, P>(mut self, chip: C, modify_trace: P) -> Self
     where
         P: Fn(&mut DenseMatrix<Val<SC>>),
     {
+        let air = chip.air();
         let mut air_proof_input = chip.generate_air_proof_input();
         let trace = air_proof_input.raw.common_main.as_mut().unwrap();
         modify_trace(trace);
-        self.air_proof_inputs.push(air_proof_input);
+        self.air_proof_inputs.push((air, air_proof_input));
         self
     }
 
@@ -379,7 +382,8 @@ where
         engine_provider: P,
     ) -> Result<VerificationData<SC>, VerificationError> {
         assert!(self.memory.is_none(), "Memory must be finalized");
-        engine_provider().run_test_impl(self.air_proof_inputs.clone())
+        let (airs, air_proof_inputs) = self.air_proof_inputs.iter().cloned().unzip();
+        engine_provider().run_test_impl(airs, air_proof_inputs)
     }
 }
 

diff --git a/crates/vm/src/arch/testing/program/mod.rs b/crates/vm/src/arch/testing/program/mod.rs
@@ -59,7 +59,6 @@ impl<SC: StarkGenericConfig> Chip<SC> for ProgramTester<Val<SC>> {
     }
 
     fn generate_air_proof_input(self) -> AirProofInput<SC> {
-        let air = self.air();
         let height = self.records.len().next_power_of_two();
         let width = self.trace_width();
         let mut values = Val::<SC>::zero_vec(height * width);
@@ -69,7 +68,7 @@ impl<SC: StarkGenericConfig> Chip<SC> for ProgramTester<Val<SC>> {
             *(row[..width - 1]).borrow_mut() = record;
             row[width - 1] = Val::<SC>::ONE;
         }
-        AirProofInput::simple_no_pis(air, RowMajorMatrix::new(values, width))
+        AirProofInput::simple_no_pis(RowMajorMatrix::new(values, width))
     }
 }
 

diff --git a/crates/vm/src/arch/vm.rs b/crates/vm/src/arch/vm.rs
@@ -7,7 +7,8 @@ use openvm_stark_backend::{
     keygen::types::{MultiStarkProvingKey, MultiStarkVerifyingKey},
     p3_commit::PolynomialSpace,
     p3_field::PrimeField32,
-    prover::types::{CommittedTraceData, Proof, ProofInput},
+    proof::Proof,
+    prover::types::{CommittedTraceData, ProofInput},
     utils::metrics_span,
     verifier::VerificationError,
     Chip,

diff --git a/crates/vm/src/system/connector/mod.rs b/crates/vm/src/system/connector/mod.rs
@@ -198,7 +198,7 @@ where
             exit_code: final_state.exit_code,
             is_terminate: final_state.is_terminate,
         };
-        AirProofInput::simple(Arc::new(self.air), trace, public_values)
+        AirProofInput::simple(trace, public_values)
     }
 }
 

diff --git a/crates/vm/src/system/memory/adapter/mod.rs b/crates/vm/src/system/memory/adapter/mod.rs
@@ -285,9 +285,8 @@ where
     }
 
     fn generate_air_proof_input(self) -> AirProofInput<SC> {
-        let air = self.air.clone();
         let trace = self.generate_trace();
-        AirProofInput::simple(Arc::new(air), trace, vec![])
+        AirProofInput::simple_no_pis(trace)
     }
 }
 

diff --git a/crates/vm/src/system/memory/merkle/tests/mod.rs b/crates/vm/src/system/memory/merkle/tests/mod.rs
@@ -96,6 +96,7 @@ fn test<const CHUNK: usize>(
         chip.final_state.as_ref().unwrap().final_root,
         final_tree_check.hash()
     );
+    let chip_air = chip.air();
     let chip_api = chip.generate_air_proof_input();
 
     let dummy_interaction_air = DummyInteractionAir::new(4 + CHUNK, true, merkle_bus.0);
@@ -162,14 +163,20 @@ fn test<const CHUNK: usize>(
         dummy_interaction_trace_rows,
         dummy_interaction_air.field_width() + 1,
     );
-    let dummy_interaction_api =
-        AirProofInput::simple_no_pis(Arc::new(dummy_interaction_air), dummy_interaction_trace);
-
-    BabyBearPoseidon2Engine::run_test_fast(vec![
-        chip_api,
-        dummy_interaction_api,
-        hash_test_chip.generate_air_proof_input(),
-    ])
+    let dummy_interaction_api = AirProofInput::simple_no_pis(dummy_interaction_trace);
+
+    BabyBearPoseidon2Engine::run_test_fast(
+        vec![
+            chip_air,
+            Arc::new(dummy_interaction_air),
+            Arc::new(hash_test_chip.air()),
+        ],
+        vec![
+            chip_api,
+            dummy_interaction_api,
+            hash_test_chip.generate_air_proof_input(),
+        ],
+    )
     .expect("Verification failed");
 }
 
@@ -283,10 +290,13 @@ fn expand_test_no_accesses() {
 
     let partition = memory_to_partition(&memory);
     chip.finalize(&tree, &partition, &mut hash_test_chip);
-    BabyBearPoseidon2Engine::run_test_fast(vec![
-        chip.generate_air_proof_input(),
-        hash_test_chip.generate_air_proof_input(),
-    ])
+    BabyBearPoseidon2Engine::run_test_fast(
+        vec![chip.air(), Arc::new(hash_test_chip.air())],
+        vec![
+            chip.generate_air_proof_input(),
+            hash_test_chip.generate_air_proof_input(),
+        ],
+    )
     .expect("This should occur");
 }
 
@@ -320,6 +330,7 @@ fn expand_test_negative() {
 
     let partition = memory_to_partition(&memory);
     chip.finalize(&tree, &partition, &mut hash_test_chip);
+    let air = chip.air();
     let mut chip_api = chip.generate_air_proof_input();
     {
         let trace = chip_api.raw.common_main.as_mut().unwrap();
@@ -332,9 +343,10 @@ fn expand_test_negative() {
         }
     }
 
-    BabyBearPoseidon2Engine::run_test_fast(vec![
-        chip_api,
-        hash_test_chip.generate_air_proof_input(),
-    ])
+    let hash_air = Arc::new(hash_test_chip.air());
+    BabyBearPoseidon2Engine::run_test_fast(
+        vec![air, hash_air],
+        vec![chip_api, hash_test_chip.generate_air_proof_input()],
+    )
     .expect("This should occur");
 }
diff --git a/crates/vm/src/system/memory/merkle/tests/util.rs b/crates/vm/src/system/memory/merkle/tests/util.rs
@@ -1,4 +1,4 @@
-use std::{array::from_fn, sync::Arc};
+use std::array::from_fn;
 
 use openvm_stark_backend::{
     config::{Domain, StarkGenericConfig},
@@ -51,7 +51,7 @@ impl<const CHUNK: usize, F: Field> HashTestChip<CHUNK, F> {
     where
         Domain<SC>: PolynomialSpace<Val = F>,
     {
-        AirProofInput::simple_no_pis(Arc::new(self.air()), self.trace())
+        AirProofInput::simple_no_pis(self.trace())
     }
 }
 

diff --git a/crates/vm/src/system/memory/merkle/trace.rs b/crates/vm/src/system/memory/merkle/trace.rs
@@ -66,7 +66,6 @@ where
     }
 
     fn generate_air_proof_input(self) -> AirProofInput<SC> {
-        let air = Arc::new(self.air);
         assert!(
             self.final_state.is_some(),
             "Merkle chip must finalize before trace generation"
@@ -99,7 +98,7 @@ where
 
         let trace = RowMajorMatrix::new(trace, width);
         let pvs = init_root.into_iter().chain(final_root).collect();
-        AirProofInput::simple(air, trace, pvs)
+        AirProofInput::simple(trace, pvs)
     }
 }
 impl<const CHUNK: usize, F: PrimeField32> ChipUsageGetter for MemoryMerkleChip<CHUNK, F> {

diff --git a/crates/vm/src/system/memory/persistent.rs b/crates/vm/src/system/memory/persistent.rs
@@ -242,7 +242,6 @@ where
     }
 
     fn generate_air_proof_input(self) -> AirProofInput<SC> {
-        let air = Arc::new(self.air);
         let trace = {
             let width = PersistentBoundaryCols::<Val<SC>, CHUNK>::width();
             // Boundary AIR should always present in order to fix the AIR ID of merkle AIR.
@@ -286,7 +285,7 @@ where
                 });
             RowMajorMatrix::new(rows, width)
         };
-        AirProofInput::simple_no_pis(air, trace)
+        AirProofInput::simple_no_pis(trace)
     }
 }