From a233d44bbe61dc6a758a754c3b78fe4f83e56699 Mon Sep 17 00:00:00 2001 From: Alex Ostrovski Date: Mon, 30 Sep 2024 11:45:03 +0300 Subject: [PATCH] fix!: Revise VM storage interface (#73) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # What ❔ Replaces storage slot data returned by `StorageInterface` with a dedicated type (essentially, `(U256, bool)`). ## Why ❔ Currently, `StorageInterface` reads storage slots as `Option` where `None` corresponds to initial writes. This loses information during sandboxed execution. Indeed, there may be slots which have non-zero value, but still are initial writes (i.e., these slots were first written to during previous blocks / transactions in the sandboxed batch). --- .github/workflows/ci.yml | 3 +- crates/vm2/src/lib.rs | 28 +- .../single_instruction_test/into_zk_evm.rs | 4 +- .../vm2/src/single_instruction_test/world.rs | 12 +- crates/vm2/src/testonly.rs | 30 +- crates/vm2/src/world_diff.rs | 280 ++++++++++++------ 6 files changed, 242 insertions(+), 115 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8a7fa3c9..ae7a9b10 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/crates/vm2/src/lib.rs b/crates/vm2/src/lib.rs index 338bfdab..c49514eb 100644 --- a/crates/vm2/src/lib.rs +++ b/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; + 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, 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; diff --git a/crates/vm2/src/single_instruction_test/into_zk_evm.rs b/crates/vm2/src/single_instruction_test/into_zk_evm.rs index b5863cad..2a6ef7a8 100644 --- a/crates/vm2/src/single_instruction_test/into_zk_evm.rs +++ b/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)) } diff --git a/crates/vm2/src/single_instruction_test/world.rs b/crates/vm2/src/single_instruction_test/world.rs index 500757c1..0eb11883 100644 --- a/crates/vm2/src/single_instruction_test/world.rs +++ b/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 World for MockWorld { } impl StorageInterface for MockWorld { - fn read_storage(&mut self, contract: H160, key: U256) -> Option { - *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) -> u32 { + fn cost_of_writing_storage(&mut self, _: StorageSlot, _: U256) -> u32 { 50 } diff --git a/crates/vm2/src/testonly.rs b/crates/vm2/src/testonly.rs index 88a932b1..c2729a92 100644 --- a/crates/vm2/src/testonly.rs +++ b/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 World for TestWorld { } impl StorageInterface for TestWorld { - fn read_storage(&mut self, contract: H160, key: U256) -> Option { + 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, _new_value: U256) -> u32 { + fn cost_of_writing_storage(&mut self, _initial_slot: StorageSlot, _new_value: U256) -> u32 { 50 } @@ -108,9 +113,8 @@ impl StorageInterface for TestWorld { pub fn initial_decommit>(world: &mut W, address: H160) -> Program { 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); diff --git a/crates/vm2/src/world_diff.rs b/crates/vm2/src/world_diff.rs index 948def92..42f21b87 100644 --- a/crates/vm2/src/world_diff.rs +++ b/crates/vm2/src/world_diff.rs @@ -9,7 +9,7 @@ use zksync_vm2_interface::{CycleStats, Event, L2ToL1Log, Tracer}; use crate::{ rollback::{Rollback, RollbackableLog, RollbackableMap, RollbackablePod, RollbackableSet}, - StorageInterface, + StorageInterface, StorageSlot, }; /// Pending modifications to the global state that are executed at the end of a block. @@ -35,7 +35,7 @@ pub struct WorldDiff { written_storage_slots: RollbackableSet<(H160, U256)>, // This is never rolled back. It is just a cache to avoid asking these from DB every time. - storage_initial_values: BTreeMap<(H160, U256), Option>, + storage_initial_values: BTreeMap<(H160, U256), StorageSlot>, } #[derive(Debug)] @@ -87,7 +87,7 @@ impl WorldDiff { .as_ref() .get(&(contract, key)) .copied() - .unwrap_or_else(|| world.read_storage(contract, key).unwrap_or_default()); + .unwrap_or_else(|| world.read_storage_value(contract, key)); let newly_added = self.read_storage_slots.add((contract, key)); if newly_added { @@ -178,17 +178,23 @@ impl WorldDiff { } /// Gets changes for all touched storage slots. - pub fn get_storage_changes( - &self, - ) -> impl Iterator, U256))> + '_ { + pub fn get_storage_changes(&self) -> impl Iterator + '_ { self.storage_changes .as_ref() .iter() .filter_map(|(key, &value)| { - if self.storage_initial_values[key].unwrap_or_default() == value { + let initial_slot = &self.storage_initial_values[key]; + if initial_slot.value == value { None } else { - Some((*key, (self.storage_initial_values[key], value))) + Some(( + *key, + StorageChange { + before: initial_slot.value, + after: value, + is_initial: initial_slot.is_write_initial, + }, + )) } }) } @@ -206,9 +212,9 @@ impl WorldDiff { ( key, StorageChange { - before: before.or(initial), + before: before.unwrap_or(initial.value), after, - is_initial: initial.is_none(), + is_initial: initial.is_write_initial, }, ) }) @@ -354,8 +360,8 @@ pub struct Snapshot { /// Change in a single storage slot. #[derive(Debug, PartialEq)] pub struct StorageChange { - /// Value before the slot was written to. `None` if the slot was not written to previously. - pub before: Option, + /// Value before the slot was written to. + pub before: U256, /// Value written to the slot. pub after: U256, /// `true` if the slot is not set in the [`World`](crate::World). @@ -369,105 +375,195 @@ const COLD_WRITE_AFTER_WARM_READ_REFUND: u32 = STORAGE_ACCESS_COLD_READ_COST; #[cfg(test)] mod tests { - use proptest::prelude::*; + use proptest::{bits, collection::btree_map, prelude::*}; use super::*; + use crate::StorageSlot; + + fn test_storage_changes( + initial_values: &BTreeMap<(H160, U256), StorageSlot>, + first_changes: BTreeMap<(H160, U256), U256>, + second_changes: BTreeMap<(H160, U256), U256>, + ) { + let mut world_diff = WorldDiff { + storage_initial_values: initial_values.clone(), + ..WorldDiff::default() + }; + + let checkpoint1 = world_diff.snapshot(); + for (key, value) in &first_changes { + world_diff.write_storage(&mut NoWorld, &mut (), key.0, key.1, *value); + } + let actual_changes = world_diff + .get_storage_changes_after(&checkpoint1) + .collect::>(); + let expected_changes = first_changes + .iter() + .map(|(key, value)| { + let before = initial_values + .get(key) + .map_or_else(U256::zero, |slot| slot.value); + let is_initial = initial_values + .get(key) + .map_or(true, |slot| slot.is_write_initial); + ( + *key, + StorageChange { + before, + after: *value, + is_initial, + }, + ) + }) + .collect(); + assert_eq!(actual_changes, expected_changes); + + let checkpoint2 = world_diff.snapshot(); + for (key, value) in &second_changes { + world_diff.write_storage(&mut NoWorld, &mut (), key.0, key.1, *value); + } + let actual_changes = world_diff + .get_storage_changes_after(&checkpoint2) + .collect::>(); + let expected_changes = second_changes + .iter() + .map(|(key, value)| { + let before = first_changes + .get(key) + .or(initial_values.get(key).map(|slot| &slot.value)) + .copied() + .unwrap_or_default(); + let is_initial = initial_values + .get(key) + .map_or(true, |slot| slot.is_write_initial); + ( + *key, + StorageChange { + before, + after: *value, + is_initial, + }, + ) + }) + .collect(); + assert_eq!(actual_changes, expected_changes); + + let mut combined = first_changes + .into_iter() + .filter_map(|(key, value)| { + let initial = initial_values + .get(&key) + .copied() + .unwrap_or(StorageSlot::EMPTY); + (initial.value != value).then_some(( + key, + StorageChange { + before: initial.value, + after: value, + is_initial: initial.is_write_initial, + }, + )) + }) + .collect::>(); + for (key, value) in second_changes { + let initial = initial_values + .get(&key) + .copied() + .unwrap_or(StorageSlot::EMPTY); + if initial.value == value { + combined.remove(&key); + } else { + combined.insert( + key, + StorageChange { + before: initial.value, + after: value, + is_initial: initial.is_write_initial, + }, + ); + } + } + + assert_eq!(combined, world_diff.get_storage_changes().collect()); + } proptest! { #[test] - fn test_storage_changes( - initial_values in arbitrary_storage_changes(), + fn storage_changes_work_as_expected( + initial_values in arbitrary_initial_storage(), first_changes in arbitrary_storage_changes(), second_changes in arbitrary_storage_changes(), ) { - let storage_initial_values = initial_values - .iter() - .map(|(key, value)| (*key, Some(*value))) - .collect(); - let mut world_diff = WorldDiff { - storage_initial_values, - ..WorldDiff::default() - }; - - let checkpoint1 = world_diff.snapshot(); - for (key, value) in &first_changes { - world_diff.write_storage(&mut NoWorld, &mut (), key.0, key.1, *value); - } - assert_eq!( - world_diff - .get_storage_changes_after(&checkpoint1) - .collect::>(), - first_changes - .iter() - .map(|(key, value)| ( - *key, - StorageChange { - before: initial_values.get(key).copied(), - after: *value, - is_initial: !initial_values.contains_key(key), - } - )) - .collect() - ); - - let checkpoint2 = world_diff.snapshot(); - for (key, value) in &second_changes { - world_diff.write_storage(&mut NoWorld, &mut (), key.0, key.1, *value); - } - assert_eq!( - world_diff - .get_storage_changes_after(&checkpoint2) - .collect::>(), - second_changes - .iter() - .map(|(key, value)| ( - *key, - StorageChange { - before: first_changes.get(key).or(initial_values.get(key)).copied(), - after: *value, - is_initial: !initial_values.contains_key(key), - } - )) - .collect() - ); - - let mut combined = first_changes - .into_iter() - .filter_map(|(key, value)| { - let initial = initial_values.get(&key).copied(); - (initial.unwrap_or_default() != value).then_some((key, (initial, value))) - }) - .collect::>(); - for (key, value) in second_changes { - let initial = initial_values.get(&key).copied(); - if initial.unwrap_or_default() == value { - combined.remove(&key); - } else { - combined.insert(key, (initial, value)); - } - } + test_storage_changes(&initial_values, first_changes, second_changes); + } - assert_eq!(combined, world_diff.get_storage_changes().collect()); + #[test] + fn storage_changes_work_with_constrained_changes( + initial_values in constrained_initial_storage(), + first_changes in constrained_storage_changes(), + second_changes in constrained_storage_changes(), + ) { + test_storage_changes(&initial_values, first_changes, second_changes); } } + /// Max items in generated initial storage / changes. + const MAX_ITEMS: usize = 5; + /// Bit mask for bytes in constrained `U256` / `H160` values. + const BIT_MASK: u8 = 0b_1111; + + fn arbitrary_initial_storage() -> impl Strategy> { + btree_map( + any::<([u8; 20], [u8; 32])>() + .prop_map(|(contract, key)| (H160::from(contract), U256::from(key))), + any::<([u8; 32], bool)>().prop_map(|(value, is_write_initial)| StorageSlot { + value: U256::from(value), + is_write_initial, + }), + 0..=MAX_ITEMS, + ) + } + + fn constrained_initial_storage() -> impl Strategy> { + btree_map( + (bits::u8::masked(BIT_MASK), bits::u8::masked(BIT_MASK)) + .prop_map(|(contract, key)| (H160::repeat_byte(contract), U256::from(key))), + (bits::u8::masked(BIT_MASK), any::()).prop_map(|(value, is_write_initial)| { + StorageSlot { + value: U256::from(value), + is_write_initial, + } + }), + 0..=MAX_ITEMS, + ) + } + fn arbitrary_storage_changes() -> impl Strategy> { - any::>().prop_map(|vec| { - vec.into_iter() - .map(|((contract, key), value)| { - ((H160::from(contract), U256::from(key)), U256::from(value)) - }) - .collect() - }) + btree_map( + any::<([u8; 20], [u8; 32])>() + .prop_map(|(contract, key)| (H160::from(contract), U256::from(key))), + any::<[u8; 32]>().prop_map(U256::from), + 0..=MAX_ITEMS, + ) + } + + fn constrained_storage_changes() -> impl Strategy> { + btree_map( + (bits::u8::masked(BIT_MASK), bits::u8::masked(BIT_MASK)) + .prop_map(|(contract, key)| (H160::repeat_byte(contract), U256::from(key))), + bits::u8::masked(BIT_MASK).prop_map(U256::from), + 0..=MAX_ITEMS, + ) } struct NoWorld; + impl StorageInterface for NoWorld { - fn read_storage(&mut self, _: H160, _: U256) -> Option { - None + fn read_storage(&mut self, _: H160, _: U256) -> StorageSlot { + StorageSlot::EMPTY } - fn cost_of_writing_storage(&mut self, _: Option, _: U256) -> u32 { + fn cost_of_writing_storage(&mut self, _: StorageSlot, _: U256) -> u32 { 0 }