fix!: Revise VM storage interface

.github/workflows/ci.yml

@@ -49,7 +49,7 @@ jobs:
         run: |
           # Check the main library with non-test features (needs to be tested in isolation since the fuzzing crate enables test features)
           cargo clippy -p zksync_vm2 --all-targets -- -D warnings
-          # The benches in `vm2` don't compile with fuzzing enabled 
+          # The benches in `vm2` don't compile with fuzzing enabled
           cargo clippy --workspace --all-features --lib --bins --tests -- -D warnings
 
       - name: Check formatting
@@ -58,6 +58,7 @@ jobs:
 
       - name: Run tests
         run: |
+          PROPTEST_CASES=10000 \
           cargo test -p zksync_vm2_interface -p zksync_vm2 --all-targets
 
       - name: Run doc tests

crates/vm2/src/lib.rs

@@ -49,15 +49,39 @@ mod tracing;
 mod vm;
 mod world_diff;
 
+/// Storage slot information returned from [`StorageInterface::read_storage()`].
+#[derive(Debug, Clone, Copy)]
+pub struct StorageSlot {
+    /// Value of the storage slot.
+    pub value: U256,
+    /// Whether a write to the slot would be considered an initial write. This influences refunds.
+    pub is_write_initial: bool,
+}
+
+impl StorageSlot {
+    /// Represents an empty storage slot.
+    pub const EMPTY: Self = Self {
+        value: U256([0; 4]),
+        is_write_initial: true,
+    };
+}
+
 /// VM storage access operations.
 pub trait StorageInterface {
     /// Reads the specified slot from the storage.
     ///
     /// There is no write counterpart; [`WorldDiff::get_storage_changes()`] gives a list of all storage changes.
-    fn read_storage(&mut self, contract: H160, key: U256) -> Option<U256>;
+    fn read_storage(&mut self, contract: H160, key: U256) -> StorageSlot;
+
+    /// Same as [`Self::read_storage()`], but doesn't request the initialness flag for the read slot.
+    ///
+    /// The default implementation uses `read_storage()`.
+    fn read_storage_value(&mut self, contract: H160, key: U256) -> U256 {
+        self.read_storage(contract, key).value
+    }
 
     /// Computes the cost of writing a storage slot.
-    fn cost_of_writing_storage(&mut self, initial_value: Option<U256>, new_value: U256) -> u32;
+    fn cost_of_writing_storage(&mut self, initial_slot: StorageSlot, new_value: U256) -> u32;
 
     /// Returns if the storage slot is free both in terms of gas and pubdata.
     fn is_free_storage_slot(&self, contract: &H160, key: &U256) -> bool;

crates/vm2/src/single_instruction_test/into_zk_evm.rs

@@ -206,9 +206,7 @@ impl Storage for MockWorldWrapper {
             query.read_value = if query.aux_byte == TRANSIENT_STORAGE_AUX_BYTE {
                 U256::zero()
             } else {
-                self.0
-                    .read_storage(query.address, query.key)
-                    .unwrap_or_default()
+                self.0.read_storage_value(query.address, query.key)
             };
             (query, PubdataCost(0))
         }

crates/vm2/src/single_instruction_test/world.rs

@@ -3,7 +3,7 @@ use primitive_types::{H160, U256};
 use zksync_vm2_interface::Tracer;
 
 use super::mock_array::MockRead;
-use crate::{Program, StorageInterface, World};
+use crate::{Program, StorageInterface, StorageSlot, World};
 
 #[derive(Debug, Arbitrary, Clone)]
 pub struct MockWorld {
@@ -21,11 +21,15 @@ impl<T: Tracer> World<T> for MockWorld {
 }
 
 impl StorageInterface for MockWorld {
-    fn read_storage(&mut self, contract: H160, key: U256) -> Option<U256> {
-        *self.storage_slot.get((contract, key))
+    fn read_storage(&mut self, contract: H160, key: U256) -> StorageSlot {
+        let value = *self.storage_slot.get((contract, key));
+        StorageSlot {
+            value: value.unwrap_or_default(),
+            is_write_initial: value.is_none(),
+        }
     }
 
-    fn cost_of_writing_storage(&mut self, _: Option<U256>, _: U256) -> u32 {
+    fn cost_of_writing_storage(&mut self, _: StorageSlot, _: U256) -> u32 {
         50
     }
 

crates/vm2/src/testonly.rs

@@ -11,7 +11,9 @@ use zkevm_opcode_defs::{
 };
 use zksync_vm2_interface::Tracer;
 
-use crate::{instruction_handlers::address_into_u256, Program, StorageInterface, World};
+use crate::{
+    instruction_handlers::address_into_u256, Program, StorageInterface, StorageSlot, World,
+};
 
 /// Test [`World`] implementation.
 #[derive(Debug)]
@@ -76,23 +78,26 @@ impl<T: Tracer> World<T> for TestWorld<T> {
 }
 
 impl<T> StorageInterface for TestWorld<T> {
-    fn read_storage(&mut self, contract: H160, key: U256) -> Option<U256> {
+    fn read_storage(&mut self, contract: H160, key: U256) -> StorageSlot {
         let deployer_system_contract_address =
             Address::from_low_u64_be(DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW.into());
 
         if contract == deployer_system_contract_address {
-            Some(
-                self.address_to_hash
-                    .get(&key)
-                    .copied()
-                    .unwrap_or(U256::zero()),
-            )
+            let value = self
+                .address_to_hash
+                .get(&key)
+                .copied()
+                .unwrap_or_else(U256::zero);
+            StorageSlot {
+                value,
+                is_write_initial: false,
+            }
         } else {
-            None
+            StorageSlot::EMPTY
         }
     }
 
-    fn cost_of_writing_storage(&mut self, _initial_value: Option<U256>, _new_value: U256) -> u32 {
+    fn cost_of_writing_storage(&mut self, _initial_slot: StorageSlot, _new_value: U256) -> u32 {
         50
     }
 
@@ -108,9 +113,8 @@ impl<T> StorageInterface for TestWorld<T> {
 pub fn initial_decommit<T: Tracer, W: World<T>>(world: &mut W, address: H160) -> Program<T, W> {
     let deployer_system_contract_address =
         Address::from_low_u64_be(DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW.into());
-    let code_info = world
-        .read_storage(deployer_system_contract_address, address_into_u256(address))
-        .unwrap_or_default();
+    let code_info =
+        world.read_storage_value(deployer_system_contract_address, address_into_u256(address));
 
     let mut code_info_bytes = [0; 32];
     code_info.to_big_endian(&mut code_info_bytes);