continuation optimization

Cargo.toml

@@ -36,6 +36,7 @@ std = [
     "parity-wasm/std",
     "validation/std",
 ]
+continuation = ["specs/continuation"]
 # Enables OS supported virtual memory.
 #
 # Note

src/lib.rs

@@ -279,7 +279,7 @@ pub use self::{
     host::{Externals, NopExternals, RuntimeArgs},
     imports::{ImportResolver, ImportsBuilder, ModuleImportResolver},
     memory::{MemoryInstance, MemoryRef, LINEAR_MEMORY_PAGE_SIZE},
-    module::{ExternVal, ModuleInstance, ModuleRef, NotStartedModuleRef},
+    module::{ExternVal, ModuleInstance, ModuleRef, NotStartedModuleRef, ENTRY},
     runner::{StackRecycler, DEFAULT_CALL_STACK_LIMIT, DEFAULT_VALUE_STACK_LIMIT},
     table::{TableInstance, TableRef},
     types::{GlobalDescriptor, MemoryDescriptor, Signature, TableDescriptor},

src/module.rs

@@ -30,7 +30,7 @@ use core::{
     fmt,
 };
 use parity_wasm::elements::{External, InitExpr, Instruction, Internal, ResizableLimits, Type};
-use specs::configure_table::ConfigureTable;
+use specs::{configure_table::ConfigureTable, jtable::StaticFrameEntry};
 use validation::{DEFAULT_MEMORY_INDEX, DEFAULT_TABLE_INDEX};
 
 /// Reference to a [`ModuleInstance`].
@@ -56,6 +56,8 @@ impl ::core::ops::Deref for ModuleRef {
     }
 }
 
+pub const ENTRY: &str = "zkmain";
+
 /// An external value is the runtime representation of an entity
 /// that can be imported or exported.
 #[derive(PartialEq)]
@@ -607,7 +609,6 @@ impl ModuleInstance {
         tracer: Option<Rc<RefCell<Tracer>>>,
     ) -> Result<NotStartedModuleRef<'m>, Error> {
         let module = loaded_module.module();
-
         let mut extern_vals = Vec::new();
         for import_entry in module.import_section().map(|s| s.entries()).unwrap_or(&[]) {
             let module_name = import_entry.module();
@@ -646,6 +647,58 @@ impl ModuleInstance {
 
         let module_ref = Self::with_externvals(loaded_module, extern_vals.iter(), tracer.clone());
 
+        let instance = module_ref
+            .as_ref()
+            .expect("failed to instantiate wasm module");
+        // set tracer's initial fid_of_entry
+        let tracer = tracer.expect("failed to initialize tracer");
+        let fid_of_entry = {
+            let idx_of_entry = instance.lookup_function_by_name(tracer.clone(), ENTRY);
+            tracer
+                .clone()
+                .borrow_mut()
+                .static_jtable_entries
+                .push(StaticFrameEntry {
+                    enable: true,
+                    frame_id: 0,
+                    next_frame_id: 0,
+                    callee_fid: idx_of_entry,
+                    fid: 0,
+                    iid: 0,
+                });
+
+            tracer
+                .clone()
+                .borrow_mut()
+                .static_jtable_entries
+                .push(if instance.has_start() {
+                    StaticFrameEntry {
+                        enable: true,
+                        frame_id: 0,
+                        next_frame_id: 0,
+                        callee_fid: 0, // the fid of start function is always 0
+                        fid: idx_of_entry,
+                        iid: 0,
+                    }
+                } else {
+                    StaticFrameEntry {
+                        enable: false,
+                        frame_id: 0,
+                        next_frame_id: 0,
+                        callee_fid: 0,
+                        fid: 0,
+                        iid: 0,
+                    }
+                });
+
+            if instance.has_start() {
+                0
+            } else {
+                idx_of_entry
+            }
+        };
+        tracer.clone().borrow_mut().set_fid_of_entry(fid_of_entry);
+
         module_ref
     }
 

src/prepare/tests.rs

@@ -282,7 +282,7 @@ fn if_without_else() {
             }),
             isa::Instruction::I32Const(2),
             isa::Instruction::Return(isa::DropKeep {
-                drop: 1,                 // 1 param
+                drop: 1,                                 // 1 param
                 keep: isa::Keep::Single(ValueType::I32), // 1 result
             }),
             isa::Instruction::I32Const(3),

src/runner.rs

@@ -2013,7 +2013,7 @@ impl Interpreter {
             };
 
             macro_rules! trace_post {
-                () => {{
+                ($is_return: ident) => {{
                     if let Some(tracer) = self.get_tracer_if_active() {
                         let post_status =
                             self.run_instruction_post(pre_status, function_context, &instruction);
@@ -2039,28 +2039,39 @@ impl Interpreter {
                             last_jump_eid,
                             post_status,
                         );
+
+                        // invoke callback to dump continuation slice tables
+                        if tracer.dump_enabled() {
+                            // println!("capacity: {}, tracer eid: {}, {}", tracer.slice_capability(), tracer.eid(), tracer.get_prev_eid());
+                            let is_last_slice = self.call_stack.is_empty() && $is_return;
+                            // println!("is last slice: {}", is_last_slice);
+                            assert!(tracer.eid() > tracer.get_prev_eid(), "eid: {}, prev_edi: {}", tracer.eid(), tracer.get_prev_eid());
+                            if (tracer.eid() - tracer.get_prev_eid() > tracer.slice_capability()) || is_last_slice {
+                                tracer.dump_witness(is_last_slice);
+                            }
+                        }
                     }
                 }};
             }
 
             match self.run_instruction(function_context, &instruction)? {
                 InstructionOutcome::RunNextInstruction => {
-                    trace_post!();
+                    trace_post!(false);
                 }
                 InstructionOutcome::Branch(target) => {
-                    trace_post!();
+                    trace_post!(false);
                     iter = instructions.iterate_from(target.dst_pc);
                     self.value_stack.drop_keep(target.drop_keep);
                 }
                 InstructionOutcome::ExecuteCall(func_ref) => {
                     // We don't record updated pc, the value should be recorded in the next trace log.
-                    trace_post!();
+                    trace_post!(false);
 
                     function_context.position = iter.position();
                     return Ok(RunResult::NestedCall(func_ref));
                 }
                 InstructionOutcome::Return(drop_keep) => {
-                    trace_post!();
+                    trace_post!(true);
 
                     if let Some(tracer) = self.tracer.clone() {
                         if tracer
@@ -2996,7 +3007,7 @@ impl Interpreter {
 }
 
 /// Function execution context.
-struct FunctionContext {
+pub(crate) struct FunctionContext {
     /// Is context initialized.
     pub is_initialized: bool,
     /// Internal function reference.
@@ -3027,7 +3038,7 @@ impl FunctionContext {
         self.is_initialized
     }
 
-    pub fn initialize(
+    fn initialize(
         &mut self,
         _locals: &[Local],
         _value_stack: &mut ValueStack,

src/tracer/etable.rs

@@ -131,12 +131,15 @@ pub(crate) trait ETable {
         allocated_memory_pages: u32,
         last_jump_eid: u32,
         step_info: StepInfo,
-    ) -> EventTableEntry;
+    );
 }
 
 impl ETable for EventTable {
     fn get_latest_eid(&self) -> u32 {
-        self.entries().last().unwrap().eid
+        match self.entries().last() {
+            Some(e) => e.eid,
+            None => 0,
+        }
     }
 
     fn get_last_entry_mut(&mut self) -> Option<&mut EventTableEntry> {
@@ -150,24 +153,22 @@ impl ETable for EventTable {
         allocated_memory_pages: u32,
         last_jump_eid: u32,
         step_info: StepInfo,
-    ) -> EventTableEntry {
+    ) {
         let sp = (DEFAULT_VALUE_STACK_LIMIT as u32)
             .checked_sub(sp)
             .unwrap()
             .checked_sub(1)
             .unwrap();
 
         let eentry = EventTableEntry {
-            eid: (self.entries().len() + 1).try_into().unwrap(),
+            eid: (self.get_latest_eid() + 1).try_into().unwrap(),
             sp,
             allocated_memory_pages,
             last_jump_eid,
             inst,
             step_info,
         };
 
-        self.entries_mut().push(eentry.clone());
-
-        eentry
+        self.entries_mut().push(eentry);
     }
 }

src/tracer/imtable.rs

@@ -3,7 +3,7 @@ use specs::{
     mtable::{LocationType, VarType},
 };
 
-#[derive(Debug, Default)]
+#[derive(Debug, Default, Clone)]
 pub struct IMTable(Vec<InitMemoryTableEntry>);
 
 impl IMTable {