chore: remove all Stateful implementations (#1376)

* remove all stateful implementations

* remove all ByteStateful trait

* fix clippy, lint

* nit

* chore: update stark-backend commit

---------

Co-authored-by: Jonathan Wang <[email protected]>

Author: 2 people

Cargo.lock

@@ -104,8 +104,8 @@ lto = "thin"
 
 [workspace.dependencies]
 # Stark Backend
-openvm-stark-backend = { git = "https://github.com/openvm-org/stark-backend.git", rev = "5819e43af08f6625a246e5e389757970f7c2a80c", default-features = false }
-openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", rev = "5819e43af08f6625a246e5e389757970f7c2a80c", default-features = false }
+openvm-stark-backend = { git = "https://github.com/openvm-org/stark-backend.git", rev = "bc364134b8315c27bfd29c6e77ac79fe77090137", default-features = false }
+openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", rev = "bc364134b8315c27bfd29c6e77ac79fe77090137", default-features = false }
 
 # OpenVM
 openvm-sdk = { path = "crates/sdk", default-features = false }

Cargo.toml

@@ -1,13 +1,9 @@
 use std::{
     borrow::{Borrow, BorrowMut},
     mem::size_of,
-    sync::{
-        atomic::{AtomicU32, Ordering},
-        Arc,
-    },
+    sync::{atomic::AtomicU32, Arc},
 };
 
-use itertools::Itertools;
 use openvm_circuit_primitives_derive::AlignedBorrow;
 use openvm_stark_backend::{
     config::{StarkGenericConfig, Val},
@@ -17,7 +13,7 @@ use openvm_stark_backend::{
     p3_matrix::{dense::RowMajorMatrix, Matrix},
     prover::types::AirProofInput,
     rap::{get_air_name, BaseAirWithPublicValues, PartitionedBaseAir},
-    AirRef, Chip, ChipUsageGetter, Stateful,
+    AirRef, Chip, ChipUsageGetter,
 };
 
 mod bus;
@@ -260,20 +256,3 @@ impl<const NUM_BITS: usize> ChipUsageGetter for SharedBitwiseOperationLookupChip
         self.0.trace_width()
     }
 }
-
-impl<const NUM_BITS: usize> Stateful<Vec<u8>> for SharedBitwiseOperationLookupChip<NUM_BITS> {
-    fn load_state(&mut self, state: Vec<u8>) {
-        // AtomicU32 can be deserialized as u32
-        let (count_range, count_xor): (Vec<u32>, Vec<u32>) = bitcode::deserialize(&state).unwrap();
-        for (x, v) in self.0.count_range.iter().zip_eq(count_range) {
-            x.store(v, Ordering::Relaxed);
-        }
-        for (x, v) in self.0.count_xor.iter().zip_eq(count_xor) {
-            x.store(v, Ordering::Relaxed);
-        }
-    }
-
-    fn store_state(&self) -> Vec<u8> {
-        bitcode::serialize(&(&self.0.count_range, &self.0.count_xor)).unwrap()
-    }
-}

crates/circuits/primitives/src/bitwise_op_lookup/mod.rs

@@ -5,13 +5,9 @@
 
 use std::{
     mem::size_of,
-    sync::{
-        atomic::{AtomicU32, Ordering},
-        Arc,
-    },
+    sync::{atomic::AtomicU32, Arc},
 };
 
-use itertools::Itertools;
 use openvm_stark_backend::{
     config::{StarkGenericConfig, Val},
     interaction::InteractionBuilder,
@@ -20,7 +16,7 @@ use openvm_stark_backend::{
     p3_matrix::{dense::RowMajorMatrix, Matrix},
     prover::types::AirProofInput,
     rap::{get_air_name, BaseAirWithPublicValues, PartitionedBaseAir},
-    AirRef, Chip, ChipUsageGetter, Stateful,
+    AirRef, Chip, ChipUsageGetter,
 };
 
 mod bus;
@@ -239,16 +235,3 @@ impl<const N: usize> ChipUsageGetter for SharedRangeTupleCheckerChip<N> {
         self.0.trace_width()
     }
 }
-
-impl<const N: usize> Stateful<Vec<u8>> for SharedRangeTupleCheckerChip<N> {
-    fn load_state(&mut self, state: Vec<u8>) {
-        let vals: Vec<u32> = bitcode::deserialize(&state).unwrap();
-        for (x, v) in self.0.count.iter().zip_eq(vals) {
-            x.store(v, Ordering::Relaxed);
-        }
-    }
-
-    fn store_state(&self) -> Vec<u8> {
-        bitcode::serialize(&self.0.count).unwrap()
-    }
-}

crates/circuits/primitives/src/range_tuple/mod.rs

@@ -6,10 +6,7 @@
 use core::mem::size_of;
 use std::{
     borrow::{Borrow, BorrowMut},
-    sync::{
-        atomic::{AtomicU32, Ordering},
-        Arc,
-    },
+    sync::{atomic::AtomicU32, Arc},
 };
 
 use openvm_circuit_primitives_derive::AlignedBorrow;
@@ -21,7 +18,7 @@ use openvm_stark_backend::{
     p3_matrix::{dense::RowMajorMatrix, Matrix},
     prover::types::AirProofInput,
     rap::{get_air_name, BaseAirWithPublicValues, PartitionedBaseAir},
-    AirRef, Chip, ChipUsageGetter, Stateful,
+    AirRef, Chip, ChipUsageGetter,
 };
 use tracing::instrument;
 
@@ -271,19 +268,6 @@ impl ChipUsageGetter for SharedVariableRangeCheckerChip {
     }
 }
 
-impl Stateful<Vec<u8>> for SharedVariableRangeCheckerChip {
-    fn load_state(&mut self, state: Vec<u8>) {
-        let count_vals: Vec<u32> = bitcode::deserialize(&state).unwrap();
-        for (x, val) in self.0.count.iter().zip(count_vals) {
-            x.store(val, Ordering::Relaxed);
-        }
-    }
-
-    fn store_state(&self) -> Vec<u8> {
-        bitcode::serialize(&self.0.count).unwrap()
-    }
-}
-
 impl AsRef<VariableRangeCheckerChip> for SharedVariableRangeCheckerChip {
     fn as_ref(&self) -> &VariableRangeCheckerChip {
         &self.0

crates/circuits/primitives/src/var_range/mod.rs

@@ -11,7 +11,7 @@ use openvm_circuit::{
     arch::{
         SystemConfig, SystemExecutor, SystemPeriphery, VmChipComplex, VmConfig, VmInventoryError,
     },
-    circuit_derive::{BytesStateful, Chip, ChipUsageGetter},
+    circuit_derive::{Chip, ChipUsageGetter},
     derive::{AnyEnum, InstructionExecutor},
 };
 use openvm_ecc_circuit::{
@@ -63,7 +63,7 @@ pub struct SdkVmConfig {
     pub castf: Option<CastFExtension>,
 }
 
-#[derive(ChipUsageGetter, Chip, InstructionExecutor, From, AnyEnum, BytesStateful)]
+#[derive(ChipUsageGetter, Chip, InstructionExecutor, From, AnyEnum)]
 pub enum SdkVmConfigExecutor<F: PrimeField32> {
     #[any_enum]
     System(SystemExecutor<F>),
@@ -93,7 +93,7 @@ pub enum SdkVmConfigExecutor<F: PrimeField32> {
     CastF(CastFExtensionExecutor<F>),
 }
 
-#[derive(From, ChipUsageGetter, Chip, AnyEnum, BytesStateful)]
+#[derive(From, ChipUsageGetter, Chip, AnyEnum)]
 pub enum SdkVmConfigPeriphery<F: PrimeField32> {
     #[any_enum]
     System(SystemPeriphery<F>),

crates/sdk/src/config/global.rs

@@ -344,14 +344,14 @@ pub fn vm_generic_config_derive(input: proc_macro::TokenStream) -> proc_macro::T
             let periphery_type = Ident::new(&format!("{}Periphery", name), name.span());
 
             TokenStream::from(quote! {
-                #[derive(ChipUsageGetter, Chip, InstructionExecutor, From, AnyEnum, ::openvm_circuit_primitives_derive::BytesStateful)]
+                #[derive(ChipUsageGetter, Chip, InstructionExecutor, From, AnyEnum)]
                 pub enum #executor_type<F: PrimeField32> {
                     #[any_enum]
                     #source_name_upper(#source_executor_type<F>),
                     #(#executor_enum_fields)*
                 }
 
-                #[derive(ChipUsageGetter, Chip, From, AnyEnum, ::openvm_circuit_primitives_derive::BytesStateful)]
+                #[derive(ChipUsageGetter, Chip, From, AnyEnum)]
                 pub enum #periphery_type<F: PrimeField32> {
                     #[any_enum]
                     #source_name_upper(#source_periphery_type<F>),

crates/vm/derive/src/lib.rs

@@ -4,7 +4,7 @@ use derive_new::new;
 use openvm_circuit::system::memory::MemoryTraceHeights;
 use openvm_instructions::program::DEFAULT_MAX_NUM_PUBLIC_VALUES;
 use openvm_poseidon2_air::Poseidon2Config;
-use openvm_stark_backend::{p3_field::PrimeField32, ChipUsageGetter, Stateful};
+use openvm_stark_backend::{p3_field::PrimeField32, ChipUsageGetter};
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 use super::{
@@ -25,8 +25,8 @@ pub fn vm_poseidon2_config<F: PrimeField32>() -> Poseidon2Config<F> {
 }
 
 pub trait VmConfig<F: PrimeField32>: Clone + Serialize + DeserializeOwned {
-    type Executor: InstructionExecutor<F> + AnyEnum + ChipUsageGetter + Stateful<Vec<u8>>;
-    type Periphery: AnyEnum + ChipUsageGetter + Stateful<Vec<u8>>;
+    type Executor: InstructionExecutor<F> + AnyEnum + ChipUsageGetter;
+    type Periphery: AnyEnum + ChipUsageGetter;
 
     /// Must contain system config
     fn system(&self) -> &SystemConfig;

crates/vm/src/arch/config.rs

@@ -7,15 +7,14 @@ use std::{
 
 use derive_more::derive::From;
 use getset::Getters;
-use itertools::Itertools;
 #[cfg(feature = "bench-metrics")]
 use metrics::counter;
 use openvm_circuit_derive::{AnyEnum, InstructionExecutor};
 use openvm_circuit_primitives::{
     utils::next_power_of_two_or_zero,
     var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus},
 };
-use openvm_circuit_primitives_derive::{BytesStateful, Chip, ChipUsageGetter};
+use openvm_circuit_primitives_derive::{Chip, ChipUsageGetter};
 use openvm_instructions::{
     program::Program, LocalOpcode, PhantomDiscriminant, PublishOpcode, SystemOpcode, VmOpcode,
 };
@@ -25,7 +24,7 @@ use openvm_stark_backend::{
     p3_field::{FieldAlgebra, PrimeField32},
     p3_matrix::Matrix,
     prover::types::{AirProofInput, CommittedTraceData, ProofInput},
-    AirRef, Chip, ChipUsageGetter, Stateful,
+    AirRef, Chip, ChipUsageGetter,
 };
 use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
@@ -39,7 +38,6 @@ use crate::metrics::VmMetrics;
 use crate::system::{
     connector::VmConnectorChip,
     memory::{
-        interface::MemoryInterface,
         merkle::{DirectCompressionBus, MemoryMerkleBus},
         offline_checker::{MemoryBridge, MemoryBus},
         online::MemoryLogEntry,
@@ -406,26 +404,6 @@ impl<E, P> VmInventory<E, P> {
     }
 }
 
-impl<E: Stateful<Vec<u8>>, P: Stateful<Vec<u8>>> Stateful<VmInventoryState> for VmInventory<E, P> {
-    fn load_state(&mut self, state: VmInventoryState) {
-        for (e, s) in self.executors.iter_mut().zip_eq(state.executors) {
-            e.load_state(s)
-        }
-        for (p, s) in self.periphery.iter_mut().zip_eq(state.periphery) {
-            p.load_state(s)
-        }
-    }
-
-    fn store_state(&self) -> VmInventoryState {
-        let executors = self.executors.iter().map(|e| e.store_state()).collect();
-        let periphery = self.periphery.iter().map(|p| p.store_state()).collect();
-        VmInventoryState {
-            executors,
-            periphery,
-        }
-    }
-}
-
 impl VmInventoryTraceHeights {
     /// Round all trace heights to the next power of two. This will round trace heights of 0 to 1.
     pub fn round_to_next_power_of_two(&mut self) {
@@ -548,37 +526,13 @@ impl<F: PrimeField32> SystemBase<F> {
     }
 }
 
-impl<F: PrimeField32> Stateful<SystemBaseState<F>> for SystemBase<F> {
-    fn load_state(&mut self, state: SystemBaseState<F>) {
-        self.range_checker_chip.load_state(state.range_checker_chip);
-        if let Some(initial_memory) = state.initial_memory {
-            self.memory_controller.set_initial_memory(initial_memory);
-        }
-        self.memory_controller.set_memory_logs(state.memory_logs);
-        self.connector_chip.load_state(state.connector_chip);
-        self.program_chip.load_state(state.program_chip);
-    }
-    fn store_state(&self) -> SystemBaseState<F> {
-        SystemBaseState {
-            range_checker_chip: self.range_checker_chip.store_state(),
-            initial_memory: match &self.memory_controller.interface_chip {
-                MemoryInterface::Volatile { .. } => None,
-                MemoryInterface::Persistent { initial_memory, .. } => Some(initial_memory.clone()),
-            },
-            memory_logs: self.memory_controller.get_memory_logs(),
-            connector_chip: self.connector_chip.store_state(),
-            program_chip: self.program_chip.store_state(),
-        }
-    }
-}
-
-#[derive(ChipUsageGetter, Chip, AnyEnum, From, InstructionExecutor, BytesStateful)]
+#[derive(ChipUsageGetter, Chip, AnyEnum, From, InstructionExecutor)]
 pub enum SystemExecutor<F: PrimeField32> {
     PublicValues(PublicValuesChip<F>),
     Phantom(RefCell<PhantomChip<F>>),
 }
 
-#[derive(ChipUsageGetter, Chip, AnyEnum, From, BytesStateful)]
+#[derive(ChipUsageGetter, Chip, AnyEnum, From)]
 pub enum SystemPeriphery<F: PrimeField32> {
     /// Poseidon2 chip with direct compression interactions
     Poseidon2(Poseidon2PeripheryChip<F>),
@@ -1122,22 +1076,6 @@ impl<F: PrimeField32, E, P> VmChipComplex<F, E, P> {
     }
 }
 
-impl<F: PrimeField32, E: Stateful<Vec<u8>>, P: Stateful<Vec<u8>>> Stateful<VmChipComplexState<F>>
-    for VmChipComplex<F, E, P>
-{
-    fn load_state(&mut self, state: VmChipComplexState<F>) {
-        self.base.load_state(state.base);
-        self.inventory.load_state(state.inventory);
-    }
-
-    fn store_state(&self) -> VmChipComplexState<F> {
-        VmChipComplexState {
-            base: self.base.store_state(),
-            inventory: self.inventory.store_state(),
-        }
-    }
-}
-
 struct VmProofInputBuilder<SC: StarkGenericConfig> {
     curr_air_id: usize,
     proof_input_per_air: Vec<(usize, AirProofInput<SC>)>,

crates/vm/src/arch/extensions.rs

@@ -17,7 +17,7 @@ use openvm_stark_backend::{
     p3_maybe_rayon::prelude::*,
     prover::types::AirProofInput,
     rap::{get_air_name, BaseAirWithPublicValues, PartitionedBaseAir},
-    AirRef, Chip, ChipUsageGetter, Stateful,
+    AirRef, Chip, ChipUsageGetter,
 };
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
@@ -225,18 +225,6 @@ where
     }
 }
 
-impl<F, A: VmAdapterChip<F>, C: VmCoreChip<F, A::Interface>> Stateful<Vec<u8>>
-    for VmChipWrapper<F, A, C>
-{
-    fn load_state(&mut self, state: Vec<u8>) {
-        self.records = bitcode::deserialize(&state).unwrap();
-    }
-
-    fn store_state(&self) -> Vec<u8> {
-        bitcode::serialize(&self.records).unwrap()
-    }
-}
-
 impl<F, A, M> InstructionExecutor<F> for VmChipWrapper<F, A, M>
 where
     F: PrimeField32,