[feat] Implement Stateful for VmChipComplex

crates/circuits/mod-builder/src/core_chip.rs

@@ -1,13 +1,11 @@
-use std::sync::Arc;
-
 use itertools::Itertools;
 use num_bigint_dig::BigUint;
 use openvm_circuit::arch::{
     AdapterAirContext, AdapterRuntimeContext, DynAdapterInterface, DynArray, MinimalInstruction,
     Result, VmAdapterInterface, VmCoreAir, VmCoreChip,
 };
 use openvm_circuit_primitives::{
-    var_range::VariableRangeCheckerChip, SubAir, TraceSubRowGenerator,
+    var_range::SharedVariableRangeCheckerChip, SubAir, TraceSubRowGenerator,
 };
 use openvm_instructions::instruction::Instruction;
 use openvm_stark_backend::{
@@ -170,7 +168,7 @@ pub struct FieldExpressionRecord {
 
 pub struct FieldExpressionCoreChip {
     pub air: FieldExpressionCoreAir,
-    pub range_checker: Arc<VariableRangeCheckerChip>,
+    pub range_checker: SharedVariableRangeCheckerChip,
 
     pub name: String,
 
@@ -184,7 +182,7 @@ impl FieldExpressionCoreChip {
         offset: usize,
         local_opcode_idx: Vec<usize>,
         opcode_flag_idx: Vec<usize>,
-        range_checker: Arc<VariableRangeCheckerChip>,
+        range_checker: SharedVariableRangeCheckerChip,
         name: &str,
         should_finalize: bool,
     ) -> Self {
@@ -277,7 +275,7 @@ where
 
     fn generate_trace_row(&self, row_slice: &mut [F], record: Self::Record) {
         self.air.expr.generate_subrow(
-            (&self.range_checker, record.inputs, record.flags),
+            (self.range_checker.as_ref(), record.inputs, record.flags),
             row_slice,
         );
     }

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

@@ -6,7 +6,10 @@
 use core::mem::size_of;
 use std::{
     borrow::{Borrow, BorrowMut},
-    sync::{atomic::AtomicU32, Arc},
+    sync::{
+        atomic::{AtomicU32, Ordering},
+        Arc,
+    },
 };
 
 use openvm_circuit_primitives_derive::AlignedBorrow;
@@ -18,7 +21,7 @@ use openvm_stark_backend::{
     p3_matrix::{dense::RowMajorMatrix, Matrix},
     prover::types::AirProofInput,
     rap::{get_air_name, AnyRap, BaseAirWithPublicValues, PartitionedBaseAir},
-    Chip, ChipUsageGetter,
+    Chip, ChipUsageGetter, Stateful,
 };
 use tracing::instrument;
 
@@ -99,6 +102,9 @@ pub struct VariableRangeCheckerChip {
     count: Vec<AtomicU32>,
 }
 
+#[derive(Clone)]
+pub struct SharedVariableRangeCheckerChip(Arc<VariableRangeCheckerChip>);
+
 impl VariableRangeCheckerChip {
     pub fn new(bus: VariableRangeCheckerBus) -> Self {
         let num_rows = (1 << (bus.range_max_bits + 1)) as usize;
@@ -179,6 +185,36 @@ impl VariableRangeCheckerChip {
     }
 }
 
+impl SharedVariableRangeCheckerChip {
+    pub fn new(bus: VariableRangeCheckerBus) -> Self {
+        Self(Arc::new(VariableRangeCheckerChip::new(bus)))
+    }
+
+    pub fn bus(&self) -> VariableRangeCheckerBus {
+        self.0.bus()
+    }
+
+    pub fn range_max_bits(&self) -> usize {
+        self.0.range_max_bits()
+    }
+
+    pub fn air_width(&self) -> usize {
+        self.0.air_width()
+    }
+
+    pub fn add_count(&self, value: u32, max_bits: usize) {
+        self.0.add_count(value, max_bits)
+    }
+
+    pub fn clear(&self) {
+        self.0.clear()
+    }
+
+    pub fn generate_trace<F: Field>(&self) -> RowMajorMatrix<F> {
+        self.0.generate_trace()
+    }
+}
+
 impl<SC: StarkGenericConfig> Chip<SC> for VariableRangeCheckerChip
 where
     Val<SC>: PrimeField32,
@@ -193,6 +229,19 @@ where
     }
 }
 
+impl<SC: StarkGenericConfig> Chip<SC> for SharedVariableRangeCheckerChip
+where
+    Val<SC>: PrimeField32,
+{
+    fn air(&self) -> Arc<dyn AnyRap<SC>> {
+        self.0.air()
+    }
+
+    fn generate_air_proof_input(self) -> AirProofInput<SC> {
+        self.0.generate_air_proof_input()
+    }
+}
+
 impl ChipUsageGetter for VariableRangeCheckerChip {
     fn air_name(&self) -> String {
         get_air_name(&self.air)
@@ -207,3 +256,36 @@ impl ChipUsageGetter for VariableRangeCheckerChip {
         NUM_VARIABLE_RANGE_COLS
     }
 }
+
+impl ChipUsageGetter for SharedVariableRangeCheckerChip {
+    fn air_name(&self) -> String {
+        self.0.air_name()
+    }
+
+    fn current_trace_height(&self) -> usize {
+        self.0.current_trace_height()
+    }
+
+    fn trace_width(&self) -> usize {
+        self.0.trace_width()
+    }
+}
+
+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/sdk/src/config/global.rs

@@ -12,7 +12,7 @@ use openvm_circuit::{
         SystemConfig, SystemExecutor, SystemPeriphery, VmChipComplex, VmConfig, VmInventoryError,
     },
     circuit_derive::{Chip, ChipUsageGetter},
-    derive::{AnyEnum, InstructionExecutor},
+    derive::{AnyEnum, InstructionExecutor, Stateful},
 };
 use openvm_ecc_circuit::{
     WeierstrassExtension, WeierstrassExtensionExecutor, WeierstrassExtensionPeriphery,
@@ -59,7 +59,7 @@ pub struct SdkVmConfig {
     pub ecc: Option<WeierstrassExtension>,
 }
 
-#[derive(ChipUsageGetter, Chip, InstructionExecutor, From, AnyEnum)]
+#[derive(ChipUsageGetter, Chip, InstructionExecutor, From, AnyEnum, Stateful)]
 pub enum SdkVmConfigExecutor<F: PrimeField32> {
     #[any_enum]
     System(SystemExecutor<F>),
@@ -87,7 +87,7 @@ pub enum SdkVmConfigExecutor<F: PrimeField32> {
     Ecc(WeierstrassExtensionExecutor<F>),
 }
 
-#[derive(From, ChipUsageGetter, Chip, AnyEnum)]
+#[derive(From, ChipUsageGetter, Chip, AnyEnum, Stateful)]
 pub enum SdkVmConfigPeriphery<F: PrimeField32> {
     #[any_enum]
     System(SystemPeriphery<F>),

crates/vm/derive/src/lib.rs

@@ -403,7 +403,7 @@ pub fn vm_generic_config_derive(input: proc_macro::TokenStream) -> proc_macro::T
                     #(#executor_enum_fields)*
                 }
 
-                #[derive(ChipUsageGetter, Chip, From, AnyEnum)]
+                #[derive(ChipUsageGetter, Chip, From, AnyEnum, ::openvm_circuit_derive::Stateful)]
                 pub enum #periphery_type<F: PrimeField32> {
                     #[any_enum]
                     #system_name_upper(SystemPeriphery<F>),

crates/vm/src/arch/config.rs

@@ -2,7 +2,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};
+use openvm_stark_backend::{p3_field::PrimeField32, ChipUsageGetter, Stateful};
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 // TODO[jpw]: re-exporting hardcoded bus constants for tests. Import paths should be
@@ -30,8 +30,8 @@ pub fn vm_poseidon2_config<F: PrimeField32>() -> Poseidon2Config<F> {
 }
 
 pub trait VmConfig<F: PrimeField32>: Clone + Serialize + DeserializeOwned {
-    type Executor: InstructionExecutor<F> + AnyEnum + ChipUsageGetter;
-    type Periphery: AnyEnum + ChipUsageGetter;
+    type Executor: InstructionExecutor<F> + AnyEnum + ChipUsageGetter + Stateful<Vec<u8>>;
+    type Periphery: AnyEnum + ChipUsageGetter + Stateful<Vec<u8>>;
 
     /// Must contain system config
     fn system(&self) -> &SystemConfig;