implement heap for zk_evm

afl-fuzz/src/main.rs

@@ -1,5 +1,5 @@
 use differential_fuzzing::VmAndWorld;
-use vm2::single_instruction_test::{vm2_to_zk_evm, NoTracer, UniversalVmState};
+use vm2::single_instruction_test::{add_heap_to_zk_evm, vm2_to_zk_evm, NoTracer, UniversalVmState};
 
 fn main() {
     afl::fuzz!(|data: &[u8]| {
@@ -18,7 +18,9 @@ fn main() {
                 let pc = vm.run_single_instruction(&mut world).unwrap();
                 assert!(vm.is_in_valid_state());
 
+                add_heap_to_zk_evm(&mut zk_evm, &vm);
                 let _ = zk_evm.cycle(&mut NoTracer);
+
                 assert_eq!(
                     UniversalVmState::from(zk_evm),
                     vm2_to_zk_evm(&vm, world.clone(), pc).into()

afl-fuzz/src/show_testcase.rs

@@ -2,6 +2,7 @@ use differential_fuzzing::VmAndWorld;
 use pretty_assertions::assert_eq;
 use std::env;
 use std::fs;
+use vm2::single_instruction_test::add_heap_to_zk_evm;
 use vm2::single_instruction_test::vm2_to_zk_evm;
 use vm2::single_instruction_test::NoTracer;
 use vm2::single_instruction_test::UniversalVmState;
@@ -37,6 +38,7 @@ fn main() {
 
     assert!(vm.is_in_valid_state());
 
+    add_heap_to_zk_evm(&mut zk_evm, &vm);
     let _ = zk_evm.cycle(&mut NoTracer);
     assert_eq!(
         UniversalVmState::from(zk_evm),

src/single_instruction_test/callframe.rs

@@ -1,5 +1,7 @@
 use super::{heap::FIRST_AUX_HEAP, stack::StackPool};
-use crate::{callframe::Callframe, decommit::is_kernel, predication::Flags, Program, WorldDiff};
+use crate::{
+    callframe::Callframe, decommit::is_kernel, predication::Flags, HeapId, Program, WorldDiff,
+};
 use arbitrary::Arbitrary;
 use u256::H160;
 
@@ -12,6 +14,10 @@ impl<'a> Arbitrary<'a> for Flags {
 impl<'a> Arbitrary<'a> for Callframe {
     fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
         let address = u.arbitrary()?;
+
+        // zk_evm requires a base page, which makes heap ids of 0 and 1 invalid
+        let base_page = u.arbitrary::<u16>()? as u32;
+
         let mut me = Self {
             address,
             code_address: u.arbitrary()?,
@@ -26,8 +32,8 @@ impl<'a> Arbitrary<'a> for Callframe {
             stipend: u.arbitrary()?,
             near_calls: vec![],
             program: u.arbitrary()?,
-            heap: u.arbitrary()?,
-            aux_heap: u.arbitrary()?,
+            heap: HeapId::from_u32_unchecked(base_page + 2),
+            aux_heap: HeapId::from_u32_unchecked(base_page + 3),
             heap_size: u.arbitrary()?,
             aux_heap_size: u.arbitrary()?,
             calldata_heap: u.arbitrary()?,

src/single_instruction_test/into_zk_evm.rs

@@ -6,7 +6,7 @@ use crate::{
 };
 use u256::U256;
 use zk_evm::{
-    abstractions::{DecommittmentProcessor, Memory, PrecompilesProcessor, Storage},
+    abstractions::{DecommittmentProcessor, Memory, MemoryType, PrecompilesProcessor, Storage},
     aux_structures::PubdataCost,
     block_properties::BlockProperties,
     reference_impls::event_sink::InMemoryEventSink,
@@ -16,11 +16,7 @@ use zk_evm::{
 };
 use zk_evm_abstractions::vm::EventSink;
 
-pub fn vm2_to_zk_evm(
-    vm: &VirtualMachine,
-    world: MockWorld,
-    program_counter: *const Instruction,
-) -> VmState<
+type ZkEvmState = VmState<
     MockWorldWrapper,
     MockMemory,
     InMemoryEventSink,
@@ -29,7 +25,13 @@ pub fn vm2_to_zk_evm(
     NoOracle,
     8,
     EncodingModeProduction,
-> {
+>;
+
+pub fn vm2_to_zk_evm(
+    vm: &VirtualMachine,
+    world: MockWorld,
+    program_counter: *const Instruction,
+) -> ZkEvmState {
     let mut event_sink = InMemoryEventSink::new();
     event_sink.start_frame(zk_evm::aux_structures::Timestamp(0));
 
@@ -44,6 +46,7 @@ pub fn vm2_to_zk_evm(
         memory: MockMemory {
             code_page: vm.state.current_frame.program.code_page().clone(),
             stack: *vm.state.current_frame.stack.clone(),
+            heap_read: None,
         },
         event_sink,
         precompiles_processor: NoOracle,
@@ -52,10 +55,38 @@ pub fn vm2_to_zk_evm(
     }
 }
 
+pub fn add_heap_to_zk_evm(zk_evm: &mut ZkEvmState, vm_after_execution: &VirtualMachine) {
+    zk_evm.memory.heap_read = vm_after_execution
+        .state
+        .heaps
+        .read
+        .read_that_happened()
+        .map(|(heapid, heap)| {
+            let (start_index, mut value) = heap.read.read_that_happened().unwrap();
+            value.reverse();
+
+            ExpectedHeapRead {
+                heap: heapid.to_u32(),
+                start_index,
+                value,
+            }
+        });
+}
+
 #[derive(Debug)]
 pub struct MockMemory {
     code_page: Arc<[U256]>,
     stack: Stack,
+    heap_read: Option<ExpectedHeapRead>,
+}
+
+// One arbitrary heap value is not enough for zk_evm
+// because it reads two U256s to read one U256.
+#[derive(Debug, Copy, Clone)]
+pub struct ExpectedHeapRead {
+    heap: u32,
+    start_index: u32,
+    value: [u8; 32],
 }
 
 impl Memory for MockMemory {
@@ -65,7 +96,7 @@ impl Memory for MockMemory {
         mut query: zk_evm::aux_structures::MemoryQuery,
     ) -> zk_evm::aux_structures::MemoryQuery {
         match query.location.memory_type {
-            zk_evm::abstractions::MemoryType::Stack => {
+            MemoryType::Stack => {
                 let slot = query.location.index.0 as u16;
                 if query.rw_flag {
                     self.stack.set(slot, query.value);
@@ -80,6 +111,28 @@ impl Memory for MockMemory {
                 }
                 query
             }
+            MemoryType::Heap => {
+                if query.rw_flag {
+                    dbg!(query);
+                    todo!()
+                } else {
+                    let heap = self.heap_read.unwrap();
+                    assert_eq!(query.location.page.0, heap.heap);
+
+                    let start = query.location.index.0 * 32;
+                    let mut read = [0; 32];
+                    for i in 0..32 {
+                        if start + i >= heap.start_index {
+                            let j = start + i - heap.start_index;
+                            if j < 32 {
+                                read[i as usize] = heap.value[j as usize];
+                            }
+                        }
+                    }
+                    query.value = U256::from_big_endian(&read);
+                }
+                query
+            }
             _ => todo!(),
         }
     }

src/single_instruction_test/mod.rs

@@ -18,6 +18,6 @@ mod validation;
 mod vm;
 mod world;
 
-pub use into_zk_evm::{vm2_to_zk_evm, NoTracer};
+pub use into_zk_evm::{add_heap_to_zk_evm, vm2_to_zk_evm, NoTracer};
 pub use universal_state::UniversalVmState;
 pub use world::MockWorld;

src/single_instruction_test/print_mock_info.rs

@@ -14,14 +14,14 @@ impl VirtualMachine {
 
 impl State {
     pub fn print_mock_info(&self) {
-        if let Some(heap) = self.heaps.read.read_that_happened() {
-            println!("Heap: {:?}", heap);
-        }
-        if let Some((address, value)) = self.heaps.read.value_read.read.read_that_happened() {
-            println!("  {value:?} read from {address:?}");
-        }
-        if let Some((address, value)) = self.heaps.read.value_read.write {
-            println!("  {value:?} written to {address:?}");
+        if let Some((heapid, heap)) = self.heaps.read.read_that_happened() {
+            println!("Heap: {:?}", heapid);
+            if let Some((address, value)) = heap.read.read_that_happened() {
+                println!("  {value:?} read from {address:?}");
+            }
+            if let Some((address, value)) = heap.write {
+                println!("  {value:?} written to {address:?}");
+            }
         }
 
         println!("Current frame:");

src/single_instruction_test/state_to_zk_evm.rs

@@ -66,7 +66,7 @@ fn vm2_frame_to_zk_evm_frames(frame: Callframe) -> impl Iterator<Item = CallStac
         this_address: frame.address,
         msg_sender: frame.caller,
         code_address: frame.code_address,
-        base_memory_page: MemoryPage(frame.heap.to_u32()),
+        base_memory_page: MemoryPage(frame.heap.to_u32() - 2),
         code_page: MemoryPage(0), // TODO
         sp: frame.sp,
         pc: 0,