Zk server compatibility

src/execution.rs

@@ -26,6 +26,7 @@ pub struct Stack {
 #[derive(Debug, Clone, PartialEq)]
 pub struct Heap {
     heap: Vec<u8>,
+    size: u32,
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -448,34 +449,47 @@ impl Stack {
 
 impl Default for Heap {
     fn default() -> Self {
-        Self::new(vec![])
+        Self::new(vec![], 0)
     }
 }
 
 impl Heap {
-    pub fn new(values: Vec<u8>) -> Self {
-        Self { heap: values }
+    pub fn new(values: Vec<u8>, size: u32) -> Self {
+        Self { heap: values, size }
     }
-    // Returns how many ergs the expand costs
-    pub fn expand_memory(&mut self, address: u32) -> u32 {
+    fn expand_memory_inner(&mut self, address: u32) {
         if address >= self.heap.len() as u32 {
-            let old_size = self.heap.len() as u32;
             self.heap.resize(address as usize, 0);
-            return MEMORY_GROWTH_ERGS_PER_BYTE * (address - old_size);
+        }
+    }
+
+    // Returns how many ergs the expand costs
+    pub fn expand_memory(&mut self, address: u32) -> u32 {
+        if address >= self.size {
+            let old_size = self.size;
+            self.size = address;
+            return MEMORY_GROWTH_ERGS_PER_BYTE * (self.size - old_size);
         }
         0
     }
 
     pub fn store(&mut self, address: u32, value: U256) {
         let start = address as usize;
         let end = start + 32;
+        self.expand_memory_inner(end as u32);
         value.to_big_endian(&mut self.heap[start..end]);
     }
 
     pub fn read(&self, address: u32) -> U256 {
         let start = address as usize;
-        let end = start + 32;
-        U256::from_big_endian(&self.heap[start..end])
+        let mut end = start + 32;
+        if start >= self.heap.len() {
+            return U256::zero();
+        }
+        end = end.min(self.heap.len());
+        let mut result = self.heap[start..end].to_vec();
+        result.resize(32, 0);
+        U256::from_big_endian(&result)
     }
 
     pub fn expanded_read(&mut self, address: u32) -> (U256, u32) {
@@ -485,13 +499,19 @@ impl Heap {
     }
 
     pub fn read_byte(&self, address: u32) -> u8 {
+        if address >= self.heap.len() as u32 {
+            return 0;
+        }
         self.heap[address as usize]
     }
 
     pub fn read_from_pointer(&self, pointer: &FatPointer) -> U256 {
         let mut result = U256::zero();
         for i in 0..32 {
             let addr = pointer.start + pointer.offset + (31 - i);
+            if addr >= self.heap.len() as u32 {
+                continue;
+            }
             result |= U256::from(self.heap[addr as usize]) << (i * 8);
         }
         result
@@ -504,10 +524,10 @@ impl Heap {
     }
 
     pub fn len(&self) -> usize {
-        self.heap.len()
+        self.size as usize
     }
 
     pub fn is_empty(&self) -> bool {
-        self.heap.is_empty()
+        self.size == 0
     }
 }

src/heaps.rs

@@ -11,9 +11,10 @@ impl Heaps {
     pub fn new(calldata: Vec<u8>) -> Self {
         // The first heap can never be used because heap zero
         // means the current heap in precompile calls
+        let size = calldata.len() as u32;
         let heaps = vec![
             Heap::default(),
-            Heap::new(calldata),
+            Heap::new(calldata, size),
             Heap::default(),
             Heap::default(),
         ];
@@ -23,15 +24,13 @@ impl Heaps {
 
     pub fn allocate(&mut self) -> u32 {
         let id = self.heaps.len() as u32;
-        self.heaps
-            .push(Heap::new(vec![0; NEW_FRAME_MEMORY_STIPEND as usize]));
+        self.heaps.push(Heap::new(vec![], NEW_FRAME_MEMORY_STIPEND));
         id
     }
 
     pub fn allocate_copy(&mut self) -> u32 {
-        let id = self.heaps.len() as u32;
-        self.heaps
-            .push(Heap::new(vec![0; NEW_FRAME_MEMORY_STIPEND as usize]));
+        let id: u32 = self.heaps.len() as u32;
+        self.heaps.push(Heap::new(vec![], NEW_FRAME_MEMORY_STIPEND));
         id
     }
 

src/op_handlers/far_call.rs

@@ -15,11 +15,13 @@ use crate::{
     rollbacks::Rollbackable,
     state::VMState,
     statistics::{VmStatistics, STORAGE_READ_STORAGE_APPLICATION_CYCLES},
-    store::{StorageError, StorageKey},
+    store::{Storage, StorageError, StorageKey},
     utils::{address_into_u256, is_kernel},
     value::{FatPointer, TaggedValue},
     Opcode,
 };
+
+use super::ret::panic_from_far_call;
 #[allow(dead_code)]
 struct FarCallParams {
     forward_memory: FatPointer,
@@ -134,14 +136,15 @@ fn decommit_code_hash(
     default_aa_code_hash: [u8; 32],
     evm_interpreter_code_hash: [u8; 32],
     is_constructor_call: bool,
+    storage: &mut dyn Storage,
 ) -> Result<(U256, bool, u32), EraVmError> {
     let mut is_evm = false;
     let deployer_system_contract_address =
         Address::from_low_u64_be(DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW as u64);
     let storage_key = StorageKey::new(deployer_system_contract_address, address_into_u256(address));
 
     // reading when decommiting doesn't refund
-    let code_info = state.storage_read_with_no_refund(storage_key);
+    let code_info = state.storage_read_with_no_refund(storage_key, storage);
     let mut code_info_bytes = [0; 32];
     code_info.to_big_endian(&mut code_info_bytes);
 
@@ -216,6 +219,7 @@ pub fn far_call(
     far_call: &FarCallOpcode,
     state: &mut VMState,
     statistics: &mut VmStatistics,
+    storage: &mut dyn Storage,
 ) -> Result<(), EraVmError> {
     let (src0, src1) = address_operands_read(vm, opcode)?;
     let contract_address = address_from_u256(&src1.value);
@@ -227,45 +231,57 @@ pub fn far_call(
     abi.is_constructor_call = abi.is_constructor_call && vm.current_context()?.is_kernel();
     abi.is_system_call = abi.is_system_call && is_kernel(&contract_address);
 
+    let FarCallParams {
+        ergs_passed,
+        forward_memory,
+        ..
+    } = far_call_params_from_register(src0, vm)?;
+
+    let mut mandated_gas = if abi.is_system_call && src1.value == ADDRESS_MSG_VALUE.into() {
+        MSG_VALUE_SIMULATOR_ADDITIVE_COST
+    } else {
+        0
+    };
+
     let (code_key, is_evm, decommit_cost) = decommit_code_hash(
         state,
         contract_address,
         vm.default_aa_code_hash,
         vm.evm_interpreter_code_hash,
         abi.is_constructor_call,
+        storage,
     )?;
 
-    // Unlike all other gas costs, this one is not paid if low on gas.
-    if decommit_cost <= vm.gas_left()? {
-        vm.decrease_gas(decommit_cost)?;
+    if vm.decrease_gas(mandated_gas).is_err() {
+        vm.set_gas_left(0)?;
+        mandated_gas = 0;
+        panic_from_far_call(vm, opcode)?;
     } else {
-        return Err(EraVmError::DecommitFailed);
+        // Pay for decommit
+        // Unlike all other gas costs, this one is not paid if low on gas.
+        if decommit_cost <= vm.gas_left()? {
+            vm.decrease_gas(decommit_cost)?;
+        } else {
+            return Err(EraVmError::DecommitFailed);
+        }
     }
 
-    let FarCallParams {
-        ergs_passed,
-        forward_memory,
-        ..
-    } = far_call_params_from_register(src0, vm)?;
-
-    let mandated_gas = if abi.is_system_call && src1.value == ADDRESS_MSG_VALUE.into() {
-        MSG_VALUE_SIMULATOR_ADDITIVE_COST
-    } else {
-        0
-    };
+    let gas_left = vm.gas_left()?;
+    let maximum_gas =
+        gas_left / FAR_CALL_GAS_SCALAR_MODIFIER_DIVISOR * FAR_CALL_GAS_SCALAR_MODIFIER_DIVIDEND;
+    let ergs_passed = ergs_passed.min(maximum_gas);
+    vm.decrease_gas(ergs_passed)?;
 
     // mandated gas can surprass the 63/64 limit
     let ergs_passed = ergs_passed + mandated_gas;
 
-    vm.decrease_gas(ergs_passed)?;
-
     let stipend = if is_evm { EVM_SIMULATOR_STIPEND } else { 0 };
 
     let ergs_passed = (ergs_passed)
         .checked_add(stipend)
         .expect("stipend must not cause overflow");
 
-    let (program_code, was_decommited) = state.decommit(code_key);
+    let (program_code, was_decommited) = state.decommit(code_key, storage);
 
     let program_code = program_code.ok_or(StorageError::KeyNotPresent)?;
     if !was_decommited {

src/op_handlers/log.rs

@@ -6,7 +6,7 @@ use crate::{
         VmStatistics, STORAGE_READ_STORAGE_APPLICATION_CYCLES,
         STORAGE_WRITE_STORAGE_APPLICATION_CYCLES,
     },
-    store::StorageKey,
+    store::{Storage, StorageKey},
     value::TaggedValue,
     Opcode,
 };
@@ -16,6 +16,7 @@ pub fn storage_write(
     opcode: &Opcode,
     state: &mut VMState,
     statistics: &mut VmStatistics,
+    storage: &mut dyn Storage,
 ) -> Result<(), EraVmError> {
     let key_for_contract_storage = vm.get_register(opcode.src0_index).value;
     let address = vm.current_context()?.contract_address;
@@ -24,7 +25,7 @@ pub fn storage_write(
         statistics.storage_application_cycles += STORAGE_WRITE_STORAGE_APPLICATION_CYCLES;
     }
     let value = vm.get_register(opcode.src1_index).value;
-    let refund = state.storage_write(key, value);
+    let refund = state.storage_write(key, value, storage);
     vm.increase_gas(refund)?;
     Ok(())
 }
@@ -34,6 +35,7 @@ pub fn storage_read(
     opcode: &Opcode,
     state: &mut VMState,
     statistics: &mut VmStatistics,
+    storage: &mut dyn Storage,
 ) -> Result<(), EraVmError> {
     let key_for_contract_storage = vm.get_register(opcode.src0_index).value;
     let address = vm.current_context()?.contract_address;
@@ -43,7 +45,7 @@ pub fn storage_read(
     if !state.read_storage_slots().contains(&key) && !state.written_storage_slots().contains(&key) {
         statistics.storage_application_cycles += STORAGE_READ_STORAGE_APPLICATION_CYCLES;
     }
-    let (value, refund) = state.storage_read(key);
+    let (value, refund) = state.storage_read(key, storage);
     vm.increase_gas(refund)?;
     vm.set_register(opcode.dst0_index, TaggedValue::new_raw_integer(value));
     Ok(())

src/op_handlers/opcode_decommit.rs

@@ -6,6 +6,7 @@ use crate::{
     execution::Execution,
     state::VMState,
     statistics::{VmStatistics, STORAGE_READ_STORAGE_APPLICATION_CYCLES},
+    store::Storage,
     value::{FatPointer, TaggedValue},
     Opcode,
 };
@@ -15,6 +16,7 @@ pub fn opcode_decommit(
     opcode: &Opcode,
     state: &mut VMState,
     statistics: &mut VmStatistics,
+    storage: &mut dyn Storage,
 ) -> Result<(), EraVmError> {
     let (src0, src1) = address_operands_read(vm, opcode)?;
 
@@ -34,7 +36,7 @@ pub fn opcode_decommit(
         return Ok(());
     }
 
-    let (code, was_decommited) = state.decommit(code_hash);
+    let (code, was_decommited) = state.decommit(code_hash, storage);
     if was_decommited {
         // refund it
         vm.increase_gas(extra_cost)?;