extract initialization state

crates/specs/src/lib.rs

@@ -39,6 +39,42 @@ pub mod mtable;
 pub mod step;
 pub mod types;
 
+#[derive(Clone, Debug, Serialize)]
+pub struct InitializationState<T> {
+    pub eid: T,
+    // FIXME: in StateInitialization or ImageTableLayouter
+    pub fid: T,
+    pub iid: T,
+    pub frame_id: T,
+    pub sp: T,
+    // FIXME: in StateInitialization or ImageTableLayouter
+    pub initial_memory_pages: T,
+    // FIXME: in StateInitialization or ImageTableLayouter
+    pub maximal_memory_pages: T,
+
+    // TODO: open mtable
+    // pub rest_mops: Option<T>,
+    pub rest_jops: T,
+
+    pub is_very_first_step: bool,
+}
+
+impl Default for InitializationState<u32> {
+    fn default() -> Self {
+        Self {
+            eid: Default::default(),
+            fid: Default::default(),
+            iid: Default::default(),
+            frame_id: Default::default(),
+            sp: Default::default(),
+            initial_memory_pages: Default::default(),
+            maximal_memory_pages: Default::default(),
+            rest_jops: Default::default(),
+            is_very_first_step: Default::default(),
+        }
+    }
+}
+
 #[derive(Default, Serialize, Debug, Clone)]
 pub struct CompilationTable {
     pub itable: InstructionTable,
@@ -51,6 +87,7 @@ pub struct CompilationTable {
 
 #[derive(Default, Serialize, Clone)]
 pub struct ExecutionTable {
+    pub initialization_state: InitializationState<u32>,
     pub etable: EventTable,
     pub jtable: JumpTable,
 }

crates/zkwasm/src/circuits/cell.rs

@@ -178,25 +178,27 @@ define_cell!(AllocatedU16Cell, F::from(u16::MAX as u64));
 define_cell!(AllocatedUnlimitedCell, -F::one());
 
 impl<F: FieldExt> AllocatedU32Cell<F> {
-    pub(crate) fn assign(&self, ctx: &mut Context<'_, F>, value: u32) -> Result<(), Error> {
+    pub(crate) fn assign(
+        &self,
+        ctx: &mut Context<'_, F>,
+        value: u32,
+    ) -> Result<AssignedCell<F, F>, Error> {
         for i in 0..2 {
             self.u16_cells_le[i].assign(ctx, (((value >> (i * 16)) & 0xffffu32) as u64).into())?;
         }
-        self.u32_cell.assign(ctx, (value as u64).into())?;
-        Ok(())
+        self.u32_cell.assign(ctx, (value as u64).into())
     }
 
     pub(crate) fn assign_constant(
         &self,
         ctx: &mut Context<'_, F>,
         value: u32,
-    ) -> Result<(), Error> {
+    ) -> Result<AssignedCell<F, F>, Error> {
         for i in 0..2 {
             self.u16_cells_le[i].assign(ctx, (((value >> (i * 16)) & 0xffffu32) as u64).into())?;
         }
 
-        self.u32_cell.assign_constant(ctx, (value as u64).into())?;
-        Ok(())
+        self.u32_cell.assign_constant(ctx, (value as u64).into())
     }
 }
 

crates/zkwasm/src/circuits/etable/assign.rs

@@ -5,6 +5,7 @@ use log::debug;
 use specs::configure_table::ConfigureTable;
 use specs::itable::Opcode;
 use specs::itable::OpcodeClassPlain;
+use specs::InitializationState;
 use std::collections::BTreeMap;
 use std::rc::Rc;
 use wasmi::DEFAULT_VALUE_STACK_LIMIT;
@@ -111,15 +112,115 @@ impl<F: FieldExt> EventTableChip<F> {
         Ok((rest_mops_cell.cell(), rest_mops_jell.cell()))
     }
 
+    #[cfg(not(feature = "continuation"))]
+    fn assign_initialization_state(
+        &self,
+        ctx: &mut Context<'_, F>,
+        initialization_state: &InitializationState<u32>,
+    ) -> Result<InitializationState<Cell>, Error> {
+        macro_rules! assign_advice {
+            ($cell:ident, $value:expr) => {
+                self.config.common_config.$cell.assign(ctx, $value)?.cell()
+            };
+        }
+
+        macro_rules! assign_constant {
+            ($cell:ident, $value:expr) => {
+                self.config
+                    .common_config
+                    .$cell
+                    .assign_constant(ctx, $value)?
+                    .cell()
+            };
+        }
+
+        let eid = assign_constant!(eid_cell, initialization_state.eid);
+        let fid = assign_advice!(fid_cell, F::from(initialization_state.fid as u64));
+        let iid = assign_constant!(iid_cell, F::from(initialization_state.iid as u64));
+        let sp = assign_constant!(sp_cell, F::from(initialization_state.sp as u64));
+        let frame_id =
+            assign_constant!(frame_id_cell, F::from(initialization_state.frame_id as u64));
+
+        let initial_memory_pages = assign_advice!(
+            mpages_cell,
+            F::from(initialization_state.initial_memory_pages as u64)
+        );
+        let maximal_memory_pages = assign_advice!(
+            maximal_memory_pages_cell,
+            F::from(initialization_state.maximal_memory_pages as u64)
+        );
+
+        let rest_jops = assign_advice!(
+            rest_jops_cell,
+            F::from(initialization_state.rest_jops as u64)
+        );
+
+        Ok(InitializationState {
+            eid,
+            fid,
+            iid,
+            frame_id,
+            sp,
+            initial_memory_pages,
+            maximal_memory_pages,
+            rest_jops,
+            is_very_first_step: initialization_state.is_very_first_step,
+        })
+    }
+
+    #[cfg(feature = "continuation")]
+    fn assign_initialization_state(
+        &self,
+        ctx: &mut Context<'_, F>,
+        initialization_state: &InitializationState<u32>,
+    ) -> Result<InitializationState<Cell>, Error> {
+        macro_rules! assign_advice {
+            ($cell:ident, $value:expr) => {
+                self.config.common_config.$cell.assign(ctx, $value)?.cell()
+            };
+        }
+
+        let eid = assign_advice!(eid_cell, initialization_state.eid);
+        let fid = assign_advice!(fid_cell, F::from(initialization_state.fid as u64));
+        let iid = assign_advice!(iid_cell, F::from(initialization_state.iid as u64));
+        let sp = assign_advice!(sp_cell, F::from(initialization_state.sp as u64));
+        let frame_id = assign_advice!(frame_id_cell, F::from(initialization_state.frame_id as u64));
+
+        let initial_memory_pages = assign_advice!(
+            mpages_cell,
+            F::from(initialization_state.initial_memory_pages as u64)
+        );
+        let maximal_memory_pages = assign_advice!(
+            maximal_memory_pages_cell,
+            F::from(initialization_state.maximal_memory_pages as u64)
+        );
+
+        let rest_jops = assign_advice!(
+            rest_jops_cell,
+            F::from(initialization_state.rest_jops as u64)
+        );
+
+        Ok(InitializationState {
+            eid,
+            fid,
+            iid,
+            frame_id,
+            sp,
+            initial_memory_pages,
+            maximal_memory_pages,
+            rest_jops,
+            is_very_first_step: initialization_state.is_very_first_step,
+        })
+    }
+
     fn assign_entries(
         &self,
         ctx: &mut Context<'_, F>,
         op_configs: &BTreeMap<OpcodeClassPlain, Rc<Box<dyn EventTableOpcodeConfig<F>>>>,
         event_table: &EventTableWithMemoryInfo,
-        configure_table: &ConfigureTable,
-        fid_of_entry: u32,
+        initialization_state: &InitializationState<u32>,
         rest_ops: Vec<(u32, u32)>,
-    ) -> Result<(Cell, Cell, Cell), Error> {
+    ) -> Result<(), Error> {
         macro_rules! assign_advice {
             ($cell:ident, $value:expr) => {
                 self.config.common_config.$cell.assign(ctx, $value)?
@@ -153,43 +254,22 @@ impl<F: FieldExt> EventTableChip<F> {
             external_host_call_index_cell,
             F::from(external_host_call_call_index as u64)
         );
-        let initial_memory_pages_cell = assign_advice!(
-            mpages_cell,
-            F::from(configure_table.init_memory_pages as u64)
-        );
-        let maximal_memory_pages_cell = assign_advice!(
-            maximal_memory_pages_cell,
-            F::from(configure_table.maximal_memory_pages as u64)
-        );
-        assign_constant!(sp_cell, F::from(DEFAULT_VALUE_STACK_LIMIT as u64 - 1));
-        assign_constant!(frame_id_cell, F::zero());
-        assign_constant!(eid_cell, 1);
-        let fid_of_entry_cell = assign_advice!(fid_cell, F::from(fid_of_entry as u64));
-        assign_constant!(iid_cell, F::zero());
 
         /*
          * Skip subsequent advice assignment in the first pass to enhance performance.
          */
         {
             let assigned_cell = assign_advice!(enabled_cell, F::zero());
             if assigned_cell.value().is_none() {
-                return Ok((
-                    fid_of_entry_cell.cell(),
-                    initial_memory_pages_cell.cell(),
-                    maximal_memory_pages_cell.cell(),
-                ));
+                return Ok(());
             }
         }
 
         /*
          * The length of event_table equals 0: without_witness
          */
         if event_table.0.len() == 0 {
-            return Ok((
-                fid_of_entry_cell.cell(),
-                initial_memory_pages_cell.cell(),
-                maximal_memory_pages_cell.cell(),
-            ));
+            return Ok(());
         }
 
         let status = {
@@ -234,7 +314,7 @@ impl<F: FieldExt> EventTableChip<F> {
                 current: &status[index],
                 next: &status[index + 1],
                 current_external_host_call_index: external_host_call_call_index,
-                configure_table: *configure_table,
+                maximal_memory_pages: initialization_state.maximal_memory_pages,
                 host_public_inputs,
                 context_in_index,
                 context_out_index,
@@ -264,7 +344,7 @@ impl<F: FieldExt> EventTableChip<F> {
             );
             assign_advice!(
                 maximal_memory_pages_cell,
-                F::from(configure_table.maximal_memory_pages as u64)
+                F::from(initialization_state.maximal_memory_pages as u64)
             );
             assign_advice!(frame_id_cell, F::from(entry.eentry.last_jump_eid as u64));
             assign_advice!(eid_cell, entry.eentry.eid);
@@ -308,7 +388,7 @@ impl<F: FieldExt> EventTableChip<F> {
         );
         assign_advice!(
             maximal_memory_pages_cell,
-            F::from(configure_table.maximal_memory_pages as u64)
+            F::from(initialization_state.maximal_memory_pages as u64)
         );
         assign_advice!(input_index_cell, F::from(host_public_inputs as u64));
         assign_advice!(context_input_index_cell, F::from(context_in_index as u64));
@@ -318,20 +398,15 @@ impl<F: FieldExt> EventTableChip<F> {
             F::from(external_host_call_call_index as u64)
         );
 
-        Ok((
-            fid_of_entry_cell.cell(),
-            initial_memory_pages_cell.cell(),
-            maximal_memory_pages_cell.cell(),
-        ))
+        Ok(())
     }
 
     pub(in crate::circuits) fn assign(
         &self,
         ctx: &mut Context<'_, F>,
         event_table: &EventTableWithMemoryInfo,
-        configure_table: &ConfigureTable,
-        fid_of_entry: u32,
-    ) -> Result<EventTablePermutationCells, Error> {
+        initialization_state: &InitializationState<u32>,
+    ) -> Result<InitializationState<Cell>, Error> {
         debug!("size of execution table: {}", event_table.0.len());
         assert!(event_table.0.len() * EVENT_TABLE_ENTRY_ROWS as usize <= self.max_available_rows);
 
@@ -340,29 +415,26 @@ impl<F: FieldExt> EventTableChip<F> {
         self.init(ctx)?;
         ctx.reset();
 
-        let (rest_mops_cell, rest_jops_cell) = self.assign_rest_ops_first_step(
-            ctx,
-            rest_ops.first().map_or(0u32, |(rest_mops, _)| *rest_mops),
-            rest_ops.first().map_or(0u32, |(_, rest_jops)| *rest_jops),
-        )?;
+        // let (rest_mops_cell, rest_jops_cell) = self.assign_rest_ops_first_step(
+        //     ctx,
+        //     rest_ops.first().map_or(0u32, |(rest_mops, _)| *rest_mops),
+        //     rest_ops.first().map_or(0u32, |(_, rest_jops)| *rest_jops),
+        // )?;
+        // ctx.reset();
+
+        let initialization_state_cells =
+            self.assign_initialization_state(ctx, initialization_state)?;
         ctx.reset();
 
-        let (fid_of_entry, initial_memory_pages, maximal_memory_pages) = self.assign_entries(
+        self.assign_entries(
             ctx,
             &self.config.op_configs,
             event_table,
-            configure_table,
-            fid_of_entry,
+            initialization_state,
             rest_ops,
         )?;
         ctx.reset();
 
-        Ok(EventTablePermutationCells {
-            rest_mops: Some(rest_mops_cell),
-            rest_jops: Some(rest_jops_cell),
-            fid_of_entry,
-            initial_memory_pages,
-            maximal_memory_pages,
-        })
+        Ok(initialization_state_cells)
     }
 }

crates/zkwasm/src/circuits/etable/op_configure/op_memory_grow.rs

@@ -133,7 +133,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for MemoryGrowConfig<F> {
                     self.current_maximal_diff.assign(
                         ctx,
                         F::from(
-                            (step.configure_table.maximal_memory_pages
+                            (step.maximal_memory_pages
                                 - (step.current.allocated_memory_pages + *grow_size as u32))
                                 as u64,
                         ),

crates/zkwasm/src/circuits/image_table/assign.rs

@@ -37,11 +37,13 @@ impl<F: FieldExt> ImageTableChip<F> {
                     }};
                 }
 
+                // FIXME: should not constrain equal if continuation is enabled.
                 let entry_fid_cell = assign_one_line!(image_table.entry_fid);
                 ctx.region
                     .constrain_equal(permutation_cells.entry_fid, entry_fid_cell)?;
 
                 let initial_memory_pages_cell = assign_one_line!(image_table.initial_memory_pages);
+                // FIXME: should not constrain equal if continuation is enabled.
                 ctx.region.constrain_equal(
                     permutation_cells.initial_memory_pages,
                     initial_memory_pages_cell,

crates/zkwasm/src/circuits/jtable/assign.rs

@@ -203,12 +203,10 @@ impl<F: FieldExt> JumpTableChip<F> {
         &self,
         ctx: &mut Context<'_, F>,
         jtable: &JumpTable,
-        etable_rest_jops_cell: Option<Cell>,
+        etable_rest_jops_cell: Cell,
         static_entries: &Vec<StaticFrameEntry>,
     ) -> Result<Vec<(Cell, Cell)>, Error> {
-        if etable_rest_jops_cell.is_some() {
-            self.constraint_to_etable_jops(ctx, etable_rest_jops_cell.unwrap())?;
-        }
+        self.constraint_to_etable_jops(ctx, etable_rest_jops_cell)?;
 
         self.init(ctx)?;
         ctx.reset();

crates/zkwasm/src/circuits/mod.rs

@@ -20,6 +20,7 @@ pub mod config;
 pub mod image_table;
 pub mod jtable;
 pub mod rtable;
+pub mod state_table;
 pub mod test_circuit;
 pub mod utils;