abstract heaps to be able to write a fuzzing-friendly heap

src/heap.rs

@@ -1,5 +1,6 @@
-use std::ops::{Index, IndexMut};
-
+use crate::instruction_handlers::HeapInterface;
+use std::ops::{Index, IndexMut, Range};
+use u256::U256;
 use zkevm_opcode_defs::system_params::NEW_FRAME_MEMORY_STIPEND;
 
 #[derive(Copy, Clone, PartialEq, Debug)]
@@ -16,8 +17,44 @@ impl HeapId {
     }
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct Heap(Vec<u8>);
+
+impl HeapInterface for Heap {
+    fn read_u256(&self, start_address: u32) -> U256 {
+        self.read_u256_partially(start_address..start_address + 32)
+    }
+    fn read_u256_partially(&self, range: Range<u32>) -> U256 {
+        let end = (range.end as usize).min(self.0.len());
+
+        let mut bytes = [0; 32];
+        for (i, j) in (range.start as usize..end).enumerate() {
+            bytes[i] = self.0[j];
+        }
+        U256::from_big_endian(&bytes)
+    }
+    fn write_u256(&mut self, start_address: u32, value: U256) {
+        let end = (start_address + 32) as usize;
+        if end > self.0.len() {
+            self.0.resize(end, 0);
+        }
+
+        let mut bytes = [0; 32];
+        value.to_big_endian(&mut bytes);
+        self.0[start_address as usize..end].copy_from_slice(&bytes);
+    }
+    fn read_range_big_endian(&self, range: Range<u32>) -> Vec<u8> {
+        let end = (range.end as usize).min(self.0.len());
+        let mut result = vec![0; range.len()];
+        for (i, j) in (range.start as usize..end).enumerate() {
+            result[i] = self.0[j];
+        }
+        result
+    }
+}
+
 #[derive(Debug, Clone)]
-pub struct Heaps(Vec<Vec<u8>>);
+pub struct Heaps(Vec<Heap>);
 
 pub(crate) const CALLDATA_HEAP: HeapId = HeapId(1);
 pub const FIRST_HEAP: HeapId = HeapId(2);
@@ -27,22 +64,28 @@ impl Heaps {
     pub(crate) fn new(calldata: Vec<u8>) -> Self {
         // The first heap can never be used because heap zero
         // means the current heap in precompile calls
-        Self(vec![vec![], calldata, vec![], vec![]])
+        Self(vec![
+            Heap(vec![]),
+            Heap(calldata),
+            Heap(vec![]),
+            Heap(vec![]),
+        ])
     }
 
     pub(crate) fn allocate(&mut self) -> HeapId {
         let id = HeapId(self.0.len() as u32);
-        self.0.push(vec![0; NEW_FRAME_MEMORY_STIPEND as usize]);
+        self.0
+            .push(Heap(vec![0; NEW_FRAME_MEMORY_STIPEND as usize]));
         id
     }
 
     pub(crate) fn deallocate(&mut self, heap: HeapId) {
-        self.0[heap.0 as usize] = vec![];
+        self.0[heap.0 as usize].0 = vec![];
     }
 }
 
 impl Index<HeapId> for Heaps {
-    type Output = Vec<u8>;
+    type Output = Heap;
 
     fn index(&self, index: HeapId) -> &Self::Output {
         &self.0[index.0 as usize]
@@ -58,10 +101,28 @@ impl IndexMut<HeapId> for Heaps {
 impl PartialEq for Heaps {
     fn eq(&self, other: &Self) -> bool {
         for i in 0..self.0.len().max(other.0.len()) {
-            if self.0.get(i).unwrap_or(&vec![]) != other.0.get(i).unwrap_or(&vec![]) {
+            if self.0.get(i).unwrap_or(&Heap(vec![])) != other.0.get(i).unwrap_or(&Heap(vec![])) {
                 return false;
             }
         }
         true
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn heap_write_resizes() {
+        let mut heap = Heap(vec![]);
+        heap.write_u256(5, 1.into());
+        assert_eq!(heap.read_u256(5), 1.into());
+    }
+
+    #[test]
+    fn heap_read_out_of_bounds() {
+        let heap = Heap(vec![]);
+        assert_eq!(heap.read_u256(5), 0.into());
+    }
+}

src/instruction_handlers/heap_access.rs

@@ -9,16 +9,24 @@ use crate::{
     state::State,
     ExecutionEnd, Instruction, VirtualMachine, World,
 };
+use std::ops::Range;
 use u256::U256;
 
+pub trait HeapInterface {
+    fn read_u256(&self, start_address: u32) -> U256;
+    fn read_u256_partially(&self, range: Range<u32>) -> U256;
+    fn write_u256(&mut self, start_address: u32, value: U256);
+    fn read_range_big_endian(&self, range: Range<u32>) -> Vec<u8>;
+}
+
 pub trait HeapFromState {
-    fn get_heap(state: &mut State) -> &mut Vec<u8>;
+    fn get_heap(state: &mut State) -> &mut impl HeapInterface;
     fn get_heap_size(state: &mut State) -> &mut u32;
 }
 
 pub struct Heap;
 impl HeapFromState for Heap {
-    fn get_heap(state: &mut State) -> &mut Vec<u8> {
+    fn get_heap(state: &mut State) -> &mut impl HeapInterface {
         &mut state.heaps[state.current_frame.heap]
     }
     fn get_heap_size(state: &mut State) -> &mut u32 {
@@ -28,7 +36,7 @@ impl HeapFromState for Heap {
 
 pub struct AuxHeap;
 impl HeapFromState for AuxHeap {
-    fn get_heap(state: &mut State) -> &mut Vec<u8> {
+    fn get_heap(state: &mut State) -> &mut impl HeapInterface {
         &mut state.heaps[state.current_frame.aux_heap]
     }
     fn get_heap_size(state: &mut State) -> &mut u32 {
@@ -63,8 +71,7 @@ fn load<H: HeapFromState, In: Source, const INCREMENT: bool>(
             return Ok(&PANIC);
         };
 
-        let heap = H::get_heap(&mut vm.state);
-        let value = U256::from_big_endian(&heap[address as usize..new_bound as usize]);
+        let value = H::get_heap(&mut vm.state).read_u256(address);
         Register1::set(args, &mut vm.state, value);
 
         if INCREMENT {
@@ -100,8 +107,7 @@ fn store<H: HeapFromState, In: Source, const INCREMENT: bool, const HOOKING_ENAB
             return Ok(&PANIC);
         }
 
-        let heap = H::get_heap(&mut vm.state);
-        value.to_big_endian(&mut heap[address as usize..new_bound as usize]);
+        H::get_heap(&mut vm.state).write_u256(address, value);
 
         if INCREMENT {
             Register1::set(args, &mut vm.state, pointer + 32)
@@ -130,11 +136,6 @@ pub fn grow_heap<H: HeapFromState>(state: &mut State, new_bound: u32) -> Result<
         state.use_gas(to_pay)?;
     }
 
-    let heap = H::get_heap(state);
-    if (heap.len() as u32) < new_bound {
-        heap.resize(new_bound as usize, 0);
-    }
-
     Ok(())
 }
 
@@ -150,26 +151,18 @@ fn load_pointer<const INCREMENT: bool>(
         }
         let pointer = FatPointer::from(input);
 
-        let heap = &vm.state.heaps[pointer.memory_page];
-
         // Usually, we just read zeroes instead of out-of-bounds bytes
         // but if offset + 32 is not representable, we panic, even if we could've read some bytes.
         // This is not a bug, this is how it must work to be backwards compatible.
         if pointer.offset > LAST_ADDRESS {
             return Ok(&PANIC);
         };
 
-        let mut buffer = [0; 32];
-        if pointer.offset < pointer.length {
-            let start = pointer.start + pointer.offset;
-            let end = start.saturating_add(32).min(pointer.start + pointer.length);
-
-            for (i, addr) in (start..end).enumerate() {
-                buffer[i] = heap[addr as usize];
-            }
-        }
+        let start = pointer.start + pointer.offset.min(pointer.length);
+        let end = start.saturating_add(32).min(pointer.start + pointer.length);
 
-        Register1::set(args, &mut vm.state, U256::from_big_endian(&buffer));
+        let value = vm.state.heaps[pointer.memory_page].read_u256_partially(start..end);
+        Register1::set(args, &mut vm.state, value);
 
         if INCREMENT {
             // This addition does not overflow because we checked that the offset is small enough above.

src/instruction_handlers/mod.rs

@@ -1,6 +1,6 @@
 pub use binop::{Add, And, Div, Mul, Or, RotateLeft, RotateRight, ShiftLeft, ShiftRight, Sub, Xor};
 pub use far_call::CallingMode;
-pub use heap_access::{AuxHeap, Heap};
+pub use heap_access::{AuxHeap, Heap, HeapInterface};
 pub use pointer::{PtrAdd, PtrPack, PtrShrink, PtrSub};
 pub(crate) use ret::{free_panic, PANIC};
 

src/instruction_handlers/precompiles.rs

@@ -1,11 +1,10 @@
-use super::{common::instruction_boilerplate_with_panic, PANIC};
+use super::{common::instruction_boilerplate_with_panic, HeapInterface, PANIC};
 use crate::{
     addressing_modes::{Arguments, Destination, Register1, Register2, Source},
     heap::{HeapId, Heaps},
     instruction::InstructionResult,
     Instruction, VirtualMachine, World,
 };
-use u256::U256;
 use zk_evm_abstractions::{
     aux::Timestamp,
     precompiles::{
@@ -97,21 +96,11 @@ impl Memory for Heaps {
     ) -> zk_evm_abstractions::queries::MemoryQuery {
         let page = HeapId::from_u32_unchecked(query.location.page.0);
 
-        let start = query.location.index.0 as usize * 32;
-        let range = start..start + 32;
+        let start = query.location.index.0 * 32;
         if query.rw_flag {
-            if range.end > self[page].len() {
-                self[page].resize(range.end, 0);
-            }
-            query.value.to_big_endian(&mut self[page][range]);
+            self[page].write_u256(start, query.value);
         } else {
-            let mut buffer = [0; 32];
-            for (i, page_index) in range.enumerate() {
-                if let Some(byte) = self[page].get(page_index) {
-                    buffer[i] = *byte;
-                }
-            }
-            query.value = U256::from_big_endian(&buffer);
+            query.value = self[page].read_u256(start);
             query.value_is_pointer = false;
         }
         query

src/instruction_handlers/ret.rs

@@ -1,4 +1,4 @@
-use super::far_call::get_far_call_calldata;
+use super::{far_call::get_far_call_calldata, HeapInterface};
 use crate::{
     addressing_modes::{Arguments, Immediate1, Register1, Source, INVALID_INSTRUCTION_COST},
     callframe::FrameRemnant,
@@ -94,8 +94,10 @@ fn ret<const RETURN_TYPE: u8, const TO_LABEL: bool>(
             // these would break were the initial frame to be rolled back.
 
             return if let Some(return_value) = return_value_or_panic {
-                let output = vm.state.heaps[return_value.memory_page][return_value.start as usize
-                    ..(return_value.start + return_value.length) as usize]
+                let output = vm.state.heaps[return_value.memory_page]
+                    .read_range_big_endian(
+                        return_value.start..return_value.start + return_value.length,
+                    )
                     .to_vec();
                 if return_type == ReturnType::Revert {
                     Err(ExecutionEnd::Reverted(output))

src/single_instruction_test/heap.rs

@@ -1,23 +1,55 @@
 use super::mock_array::MockRead;
+use crate::instruction_handlers::HeapInterface;
 use arbitrary::Arbitrary;
 use std::ops::{Index, IndexMut};
-
-//#[derive(Debug, Clone)]
-type Heap = Vec<u8>;
+use u256::U256;
 
 #[derive(Debug, Clone)]
-pub struct Heaps {
-    read: MockRead<HeapId, Heap>,
+pub struct Heap {
+    read: MockRead<u32, [u8; 32]>,
+    write: Option<(u32, U256)>,
+}
+
+impl HeapInterface for Heap {
+    fn read_u256(&self, start_address: u32) -> U256 {
+        assert!(self.write.is_none());
+        U256::from_little_endian(self.read.get(start_address))
+    }
+
+    fn read_u256_partially(&self, range: std::ops::Range<u32>) -> U256 {
+        assert!(self.write.is_none());
+        let mut result = *self.read.get(range.start);
+        for byte in &mut result[0..range.len()] {
+            *byte = 0;
+        }
+        U256::from_little_endian(&result)
+    }
+
+    fn write_u256(&mut self, start_address: u32, value: U256) {
+        assert!(self.write.is_none());
+        self.write = Some((start_address, value));
+    }
+
+    fn read_range_big_endian(&self, _: std::ops::Range<u32>) -> Vec<u8> {
+        // This is wrong, but this method is only used to get the final return value.
+        vec![]
+    }
 }
 
-impl<'a> Arbitrary<'a> for Heaps {
+impl<'a> Arbitrary<'a> for Heap {
     fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
         Ok(Self {
-            read: MockRead::new(vec![u.arbitrary()?; 1]),
+            read: u.arbitrary()?,
+            write: None,
         })
     }
 }
 
+#[derive(Debug, Clone, Arbitrary)]
+pub struct Heaps {
+    read: MockRead<HeapId, Heap>,
+}
+
 pub(crate) const CALLDATA_HEAP: HeapId = HeapId(1);
 pub const FIRST_HEAP: HeapId = HeapId(2);
 pub(crate) const FIRST_AUX_HEAP: HeapId = HeapId(3);