translate near calls to zk_evm

matter-labs · Jun 4, 2024 · ad6f534 · ad6f534
1 parent d55475b
commit ad6f534
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 33 deletions.
diff --git a/src/callframe.rs b/src/callframe.rs
@@ -50,10 +50,10 @@ pub struct Callframe {
 
 #[derive(Clone, PartialEq, Debug)]
 pub(crate) struct NearCallFrame {
-    call_instruction: u16,
-    exception_handler: u16,
-    previous_frame_sp: u16,
-    previous_frame_gas: u32,
+    pub(crate) call_instruction: u16,
+    pub(crate) exception_handler: u16,
+    pub(crate) previous_frame_sp: u16,
+    pub(crate) previous_frame_gas: u32,
     world_before_this_frame: Snapshot,
 }
 

diff --git a/src/single_instruction_test/state_to_zk_evm.rs b/src/single_instruction_test/state_to_zk_evm.rs
@@ -1,15 +1,27 @@
+use crate::callframe::{Callframe, NearCallFrame};
+use std::iter;
 use u256::U256;
 use zk_evm::{
     aux_structures::{MemoryPage, PubdataCost},
     vm_state::{execution_stack::CallStackEntry, Callstack, PrimitiveValue, VmLocalState},
 };
 use zkevm_opcode_defs::decoding::EncodingModeProduction;
 
-use crate::callframe::Callframe;
-
 pub(crate) fn vm2_state_to_zk_evm_state(
     state: &crate::State,
 ) -> VmLocalState<8, EncodingModeProduction> {
+    // zk_evm requires an unused bottom frame
+    let mut callframes: Vec<_> = iter::once(CallStackEntry::empty_context())
+        .chain(
+            state
+                .previous_frames
+                .iter()
+                .map(|x| x.1.clone())
+                .chain(iter::once(state.current_frame.clone()))
+                .flat_map(vm2_frame_to_zk_evm_frames),
+        )
+        .collect();
+
     VmLocalState {
         previous_code_word: U256([0, 0, 0, state.current_frame.raw_first_instruction()]),
         previous_code_memory_page: MemoryPage(0),
@@ -36,38 +48,57 @@ pub(crate) fn vm2_state_to_zk_evm_state(
         previous_super_pc: 0, // Same as current pc so the instruction is read from previous_code_word
         context_u128_register: state.context_u128,
         callstack: Callstack {
-            current: (&state.current_frame).into(),
+            current: callframes.pop().unwrap(),
             // zk_evm requires an unused bottom frame
-            inner: std::iter::once(CallStackEntry::empty_context())
-                .chain(state.previous_frames.iter().map(|(_, frame)| frame.into()))
-                .collect(),
+            inner: callframes,
         },
         pubdata_revert_counter: PubdataCost(0),
     }
 }
 
-impl From<&Callframe> for CallStackEntry {
-    fn from(frame: &Callframe) -> Self {
-        CallStackEntry {
-            this_address: frame.address,
-            msg_sender: frame.caller,
-            code_address: frame.code_address,
-            base_memory_page: MemoryPage(frame.heap.to_u32()),
-            code_page: MemoryPage(0), // TODO
-            sp: frame.sp,
-            pc: 0,
-            exception_handler_location: frame.exception_handler,
-            ergs_remaining: frame.gas,
-            this_shard_id: 0,
-            caller_shard_id: 0,
-            code_shard_id: 0,
-            is_static: frame.is_static,
-            is_local_frame: false, // TODO this is for making near calls
-            context_u128_value: frame.context_u128,
-            heap_bound: frame.heap_size,
-            aux_heap_bound: frame.aux_heap_size,
-            total_pubdata_spent: PubdataCost(0),
-            stipend: frame.stipend,
-        }
+fn vm2_frame_to_zk_evm_frames(frame: Callframe) -> impl Iterator<Item = CallStackEntry> {
+    let far_frame = CallStackEntry {
+        this_address: frame.address,
+        msg_sender: frame.caller,
+        code_address: frame.code_address,
+        base_memory_page: MemoryPage(frame.heap.to_u32()),
+        code_page: MemoryPage(0), // TODO
+        sp: frame.sp,
+        pc: 0,
+        exception_handler_location: frame.exception_handler,
+        ergs_remaining: frame.gas,
+        this_shard_id: 0,
+        caller_shard_id: 0,
+        code_shard_id: 0,
+        is_static: frame.is_static,
+        is_local_frame: false,
+        context_u128_value: frame.context_u128,
+        heap_bound: frame.heap_size,
+        aux_heap_bound: frame.aux_heap_size,
+        total_pubdata_spent: PubdataCost(0),
+        stipend: frame.stipend,
+    };
+
+    let mut result = vec![far_frame];
+    for NearCallFrame {
+        call_instruction,
+        exception_handler,
+        previous_frame_sp,
+        previous_frame_gas,
+        ..
+    } in frame.near_calls
+    {
+        let last = result.last_mut().unwrap();
+        last.pc = call_instruction;
+        last.sp = previous_frame_sp;
+        last.ergs_remaining = previous_frame_gas;
+
+        result.push(CallStackEntry {
+            is_local_frame: true,
+            exception_handler_location: exception_handler,
+            ..far_frame
+        });
     }
+
+    result.into_iter()
 }