separate circuit into init/final

hero78119 · hero78119 · commit b84f74e4de3a · 2025-10-14T15:05:41.000+08:00
diff --git a/ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs b/ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs
@@ -16,6 +16,7 @@ use crate::{
 
 pub struct RegConfigs<E: ExtensionField> {
     pub reg_init_config: <RegTableInitCircuit<E> as TableCircuit<E>>::TableConfig,
+    pub reg_final_config: <RegTableFinalCircuit<E> as TableCircuit<E>>::TableConfig,
     pub reg_mem_bus: <RegTableInitCircuit<E> as TableCircuit<E>>::TableConfig,
 }
 
diff --git a/ceno_zkvm/src/tables/ram.rs b/ceno_zkvm/src/tables/ram.rs
@@ -8,7 +8,10 @@ use crate::{
 
 mod ram_circuit;
 mod ram_impl;
-use crate::tables::ram::ram_circuit::NonVolatileInitRamCircuit;
+use crate::tables::ram::ram_impl::{
+    DynVolatileRamTableConfig, DynVolatileRamTableFinalConfig, DynVolatileRamTableInitConfig,
+    NonVolatileFinalTableConfig, NonVolatileInitTableConfig,
+};
 pub use ram_circuit::{DynVolatileRamTable, MemFinalRecord, MemInitRecord, NonVolatileTable};
 
 #[derive(Clone)]
@@ -33,7 +36,10 @@ impl DynVolatileRamTable for HeapTable {
     }
 }
 
-pub type HeapCircuit<E> = DynVolatileRamCircuit<E, HeapTable>;
+pub type HeapInitCircuit<E> =
+    DynVolatileRamCircuit<E, HeapTable, DynVolatileRamTableInitConfig<HeapTable>>;
+pub type HeapFinalCircuit<E> =
+    DynVolatileRamCircuit<E, HeapTable, DynVolatileRamTableFinalConfig<HeapTable>>;
 
 #[derive(Clone)]
 pub struct StackTable;
@@ -67,7 +73,10 @@ impl DynVolatileRamTable for StackTable {
     }
 }
 
-pub type StackCircuit<E> = DynVolatileRamCircuit<E, StackTable>;
+pub type StackInitCircuit<E> =
+    DynVolatileRamCircuit<E, StackTable, DynVolatileRamTableInitConfig<StackTable>>;
+pub type StackFinalCircuit<E> =
+    DynVolatileRamCircuit<E, StackTable, DynVolatileRamTableFinalConfig<StackTable>>;
 
 #[derive(Clone)]
 pub struct HintsTable;
@@ -89,7 +98,8 @@ impl DynVolatileRamTable for HintsTable {
         "HintsTable"
     }
 }
-pub type HintsCircuit<E> = DynVolatileRamCircuit<E, HintsTable>;
+pub type HintsCircuit<E> =
+    DynVolatileRamCircuit<E, HintsTable, DynVolatileRamTableConfig<HintsTable>>;
 
 /// RegTable, fix size without offset
 #[derive(Clone)]
@@ -109,8 +119,10 @@ impl NonVolatileTable for RegTable {
     }
 }
 
-// pub type RegTableCircuit<E> = NonVolatileRamCircuit<E, RegTable>;
-pub type RegTableInitCircuit<E> = NonVolatileInitRamCircuit<E, RegTable>;
+pub type RegTableInitCircuit<E> =
+    NonVolatileRamCircuit<E, RegTable, NonVolatileInitTableConfig<RegTable>>;
+pub type RegTableFinalCircuit<E> =
+    NonVolatileRamCircuit<E, RegTable, NonVolatileFinalTableConfig<RegTable>>;
 
 #[derive(Clone)]
 pub struct StaticMemTable;
@@ -129,7 +141,10 @@ impl NonVolatileTable for StaticMemTable {
     }
 }
 
-pub type StaticMemCircuit<E> = NonVolatileRamCircuit<E, StaticMemTable>;
+pub type StaticMemInitCircuit<E> =
+    NonVolatileRamCircuit<E, StaticMemTable, NonVolatileInitTableConfig<StaticMemTable>>;
+pub type StaticMemFinalCircuit<E> =
+    NonVolatileRamCircuit<E, StaticMemTable, NonVolatileFinalTableConfig<StaticMemTable>>;
 
 #[derive(Clone)]
 pub struct PubIOTable;
diff --git a/ceno_zkvm/src/tables/ram/ram_circuit.rs b/ceno_zkvm/src/tables/ram/ram_circuit.rs
@@ -12,7 +12,8 @@ use gkr_iop::error::CircuitBuilderError;
 use witness::{InstancePaddingStrategy, RowMajorMatrix};
 
 use super::ram_impl::{
-    DynVolatileRamTableConfig, NonVolatileInitTableConfig, NonVolatileTableConfig, PubIOTableConfig,
+    DynVolatileRamTableConfig, NonVolatileInitTableConfig, NonVolatileTableConfig,
+    NonVolatileTableConfigTrait, PubIOTableConfig,
 };
 
 #[derive(Clone, Debug)]
@@ -62,12 +63,15 @@ pub trait NonVolatileTable {
 /// - with fixed initial content,
 /// - with witnessed final content that the program wrote, if WRITABLE,
 /// - or final content equal to initial content, if not WRITABLE.
-pub struct NonVolatileRamCircuit<E, R>(PhantomData<(E, R)>);
+pub struct NonVolatileRamCircuit<E, R, C>(PhantomData<(E, R, C)>);
 
-impl<E: ExtensionField, NVRAM: NonVolatileTable + Send + Sync + Clone> TableCircuit<E>
-    for NonVolatileRamCircuit<E, NVRAM>
+impl<
+    E: ExtensionField,
+    NVRAM: NonVolatileTable + Send + Sync + Clone,
+    C: NonVolatileTableConfigTrait<NVRAM>,
+> TableCircuit<E> for NonVolatileRamCircuit<E, NVRAM, C>
 {
-    type TableConfig = NonVolatileTableConfig<NVRAM>;
+    type TableConfig = C::Config;
     type FixedInput = [MemInitRecord];
     type WitnessInput = [MemFinalRecord];
 
@@ -79,10 +83,7 @@ impl<E: ExtensionField, NVRAM: NonVolatileTable + Send + Sync + Clone> TableCirc
         cb: &mut CircuitBuilder<E>,
         params: &ProgramParams,
     ) -> Result<Self::TableConfig, ZKVMError> {
-        Ok(cb.namespace(
-            || Self::name(),
-            |cb| Self::TableConfig::construct_circuit(cb, params),
-        )?)
+        Ok(cb.namespace(|| Self::name(), |cb| C::construct_circuit(cb, params))?)
     }
 
     fn generate_fixed_traces(
@@ -91,7 +92,7 @@ impl<E: ExtensionField, NVRAM: NonVolatileTable + Send + Sync + Clone> TableCirc
         init_v: &Self::FixedInput,
     ) -> RowMajorMatrix<E::BaseField> {
         // assume returned table is well-formed include padding
-        config.gen_init_state(num_fixed, init_v)
+        C::gen_init_state(config, num_fixed, init_v)
     }
 
     fn assign_instances(
@@ -102,57 +103,13 @@ impl<E: ExtensionField, NVRAM: NonVolatileTable + Send + Sync + Clone> TableCirc
         final_v: &Self::WitnessInput,
     ) -> Result<RMMCollections<E::BaseField>, ZKVMError> {
         // assume returned table is well-formed include padding
-        Ok(config.assign_instances(num_witin, num_structural_witin, final_v)?)
-    }
-}
-
-/// NonVolatileRamCircuit initializes and finalizes memory
-/// - at fixed addresses,
-/// - with fixed initial content,
-/// - with witnessed final content that the program wrote, if WRITABLE,
-/// - or final content equal to initial content, if not WRITABLE.
-pub struct NonVolatileInitRamCircuit<E, R>(PhantomData<(E, R)>);
-
-impl<E: ExtensionField, NVRAM: NonVolatileTable + Send + Sync + Clone> TableCircuit<E>
-    for NonVolatileInitRamCircuit<E, NVRAM>
-{
-    type TableConfig = NonVolatileInitTableConfig<NVRAM>;
-    type FixedInput = [MemInitRecord];
-    type WitnessInput = [MemFinalRecord];
-
-    fn name() -> String {
-        format!("RAM_{:?}_{}", NVRAM::RAM_TYPE, NVRAM::name())
-    }
-
-    fn construct_circuit(
-        cb: &mut CircuitBuilder<E>,
-        params: &ProgramParams,
-    ) -> Result<Self::TableConfig, ZKVMError> {
-        Ok(cb.namespace(
-            || Self::name(),
-            |cb| Self::TableConfig::construct_circuit(cb, params),
+        Ok(C::assign_instances(
+            config,
+            num_witin,
+            num_structural_witin,
+            final_v,
         )?)
     }
-
-    fn generate_fixed_traces(
-        config: &Self::TableConfig,
-        num_fixed: usize,
-        init_v: &Self::FixedInput,
-    ) -> RowMajorMatrix<E::BaseField> {
-        // assume returned table is well-formed include padding
-        config.gen_init_state(num_fixed, init_v)
-    }
-
-    fn assign_instances(
-        config: &Self::TableConfig,
-        num_witin: usize,
-        num_structural_witin: usize,
-        _multiplicity: &[HashMap<u64, usize>],
-        final_v: &Self::WitnessInput,
-    ) -> Result<RMMCollections<E::BaseField>, ZKVMError> {
-        // assume returned table is well-formed include padding
-        Ok(config.assign_instances(num_witin, num_structural_witin, final_v)?)
-    }
 }
 
 /// PubIORamCircuit initializes and finalizes memory
@@ -241,13 +198,13 @@ pub trait DynVolatileRamTable {
 }
 
 pub trait DynVolatileRamTableConfigTrait<DVRAM>: Sized + Send + Sync {
-    type Output: Sized + Send + Sync;
+    type Config: Sized + Send + Sync;
     fn construct_circuit<E: ExtensionField>(
         cb: &mut CircuitBuilder<E>,
         params: &ProgramParams,
-    ) -> Result<Self::Output, CircuitBuilderError>;
+    ) -> Result<Self::Config, CircuitBuilderError>;
     fn assign_instances<F: SmallField>(
-        &self,
+        config: &Self::Config,
         num_witin: usize,
         num_structural_witin: usize,
         final_mem: &[MemFinalRecord],
@@ -270,7 +227,7 @@ impl<
     C: DynVolatileRamTableConfigTrait<DVRAM>,
 > TableCircuit<E> for DynVolatileRamCircuit<E, DVRAM, C>
 {
-    type TableConfig = C::Output;
+    type TableConfig = C::Config;
     type FixedInput = ();
     type WitnessInput = [MemFinalRecord];
 
@@ -301,6 +258,13 @@ impl<
         final_v: &Self::WitnessInput,
     ) -> Result<RMMCollections<E::BaseField>, ZKVMError> {
         // assume returned table is well-formed include padding
-        Ok(config.assign_instances(num_witin, num_structural_witin, final_v)?)
+        Ok(
+            <C as DynVolatileRamTableConfigTrait<DVRAM>>::assign_instances(
+                config,
+                num_witin,
+                num_structural_witin,
+                final_v,
+            )?,
+        )
     }
 }
diff --git a/ceno_zkvm/src/tables/ram/ram_impl.rs b/ceno_zkvm/src/tables/ram/ram_impl.rs

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@ use crate::{`
`16`	`16`
`17`	`17`	`pub struct RegConfigs<E: ExtensionField> {`
`18`	`18`	`pub reg_init_config: <RegTableInitCircuit<E> as TableCircuit<E>>::TableConfig,`
	`19`	`+ pub reg_final_config: <RegTableFinalCircuit<E> as TableCircuit<E>>::TableConfig,`
`19`	`20`	`pub reg_mem_bus: <RegTableInitCircuit<E> as TableCircuit<E>>::TableConfig,`
`20`	`21`	`}`
`21`	`22`