implement code page for zk_evm

matter-labs · May 30, 2024 · f0306ab · f0306ab
1 parent ab8b633
commit f0306ab
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/src/single_instruction_test/into_zk_evm.rs b/src/single_instruction_test/into_zk_evm.rs
@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use super::{
     mock_array::MockRead, stack::Stack, state_to_zk_evm::vm2_state_to_zk_evm_state, MockWorld,
 };
@@ -35,12 +37,7 @@ pub fn vm2_to_zk_evm(
         },
         storage: MockWorldWrapper(world),
         memory: MockMemory {
-            instruction_to_read: MockRead::new(U256([
-                0,
-                0,
-                0,
-                vm.state.current_frame.program.raw_first_instruction,
-            ])),
+            code_page: vm.state.current_frame.program.code_page().clone(),
             stack: *vm.state.current_frame.stack.clone(),
         },
         event_sink: InMemoryEventSink::new(),
@@ -52,7 +49,7 @@ pub fn vm2_to_zk_evm(
 
 #[derive(Debug)]
 pub struct MockMemory {
-    instruction_to_read: MockRead<MemoryLocation, U256>,
+    code_page: Arc<[U256]>,
     stack: Stack,
 }
 
@@ -95,7 +92,12 @@ impl Memory for MockMemory {
         _: u32,
         mut query: zk_evm::aux_structures::MemoryQuery,
     ) -> zk_evm::aux_structures::MemoryQuery {
-        query.value = *self.instruction_to_read.get(query.location);
+        // Code page read, instruction reads don't happen because the code word cache has been set up
+        query.value = self
+            .code_page
+            .get(query.location.index.0 as usize)
+            .cloned()
+            .unwrap_or_default();
         query
     }
 }

diff --git a/src/single_instruction_test/state_to_zk_evm.rs b/src/single_instruction_test/state_to_zk_evm.rs
@@ -1,3 +1,4 @@
+use u256::U256;
 use zk_evm::{
     aux_structures::{MemoryPage, PubdataCost},
     vm_state::{execution_stack::CallStackEntry, Callstack, PrimitiveValue, VmLocalState},
@@ -10,8 +11,7 @@ pub(crate) fn vm2_state_to_zk_evm_state(
     state: &crate::State,
 ) -> VmLocalState<8, EncodingModeProduction> {
     VmLocalState {
-        // To ensure that this field is not read, we make previous_super_pc != super_pc
-        previous_code_word: 0.into(),
+        previous_code_word: U256([0, 0, 0, state.current_frame.raw_first_instruction()]),
         previous_code_memory_page: MemoryPage(0),
         registers: state
             .registers
@@ -33,7 +33,7 @@ pub(crate) fn vm2_state_to_zk_evm_state(
         absolute_execution_step: 0,
         tx_number_in_block: state.transaction_number,
         pending_exception: false,
-        previous_super_pc: 13, // Current pc is zero so anything else is fine
+        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(),