Move StorageWithOverrides to VM interface

core/lib/vm_interface/src/storage/mod.rs

@@ -5,11 +5,13 @@ use zksync_types::{get_known_code_key, StorageKey, StorageValue, H256};
 pub use self::{
     // Note, that `test_infra` of the bootloader tests relies on this value to be exposed
     in_memory::{InMemoryStorage, IN_MEMORY_STORAGE_DEFAULT_NETWORK_ID},
+    overrides::StorageWithOverrides,
     snapshot::{StorageSnapshot, StorageWithSnapshot},
     view::{ImmutableStorageView, StorageView, StorageViewCache, StorageViewStats},
 };
 
 mod in_memory;
+mod overrides;
 mod snapshot;
 mod view;
 

core/lib/vm_interface/src/storage/overrides.rs

@@ -0,0 +1,70 @@
+//! VM storage functionality specifically used in the VM sandbox.
+
+use std::{
+    collections::{HashMap, HashSet},
+    fmt,
+};
+
+use zksync_types::{AccountTreeId, StorageKey, StorageValue, H256};
+
+use super::ReadStorage;
+
+/// A storage view that allows to override some of the storage values.
+#[derive(Debug)]
+pub struct StorageWithOverrides<S> {
+    storage_handle: S,
+    overridden_slots: HashMap<StorageKey, H256>,
+    overridden_factory_deps: HashMap<H256, Vec<u8>>,
+    empty_accounts: HashSet<AccountTreeId>,
+}
+
+impl<S: ReadStorage> StorageWithOverrides<S> {
+    /// Creates a new storage view based on the underlying storage.
+    pub fn new(storage: S) -> Self {
+        Self {
+            storage_handle: storage,
+            overridden_slots: HashMap::new(),
+            overridden_factory_deps: HashMap::new(),
+            empty_accounts: HashSet::new(),
+        }
+    }
+
+    pub fn set_value(&mut self, key: StorageKey, value: StorageValue) {
+        self.overridden_slots.insert(key, value);
+    }
+
+    pub fn store_factory_dep(&mut self, hash: H256, code: Vec<u8>) {
+        self.overridden_factory_deps.insert(hash, code);
+    }
+
+    pub fn insert_erased_account(&mut self, account: AccountTreeId) {
+        self.empty_accounts.insert(account);
+    }
+}
+
+impl<S: ReadStorage + fmt::Debug> ReadStorage for StorageWithOverrides<S> {
+    fn read_value(&mut self, key: &StorageKey) -> StorageValue {
+        if let Some(value) = self.overridden_slots.get(key) {
+            return *value;
+        }
+        if self.empty_accounts.contains(key.account()) {
+            return H256::zero();
+        }
+        self.storage_handle.read_value(key)
+    }
+
+    fn is_write_initial(&mut self, key: &StorageKey) -> bool {
+        self.storage_handle.is_write_initial(key)
+    }
+
+    fn load_factory_dep(&mut self, hash: H256) -> Option<Vec<u8>> {
+        self.overridden_factory_deps
+            .get(&hash)
+            .cloned()
+            .or_else(|| self.storage_handle.load_factory_dep(hash))
+    }
+
+    fn get_enumeration_index(&mut self, key: &StorageKey) -> Option<u64> {
+        self.storage_handle.get_enumeration_index(key)
+    }
+}

core/node/api_server/src/execution_sandbox/storage.rs

@@ -1,127 +1,67 @@
 //! VM storage functionality specifically used in the VM sandbox.
 
-use std::{
-    collections::{HashMap, HashSet},
-    fmt,
-};
-
-use zksync_multivm::interface::storage::ReadStorage;
+use zksync_multivm::interface::storage::{ReadStorage, StorageWithOverrides};
 use zksync_types::{
     api::state_override::{OverrideState, StateOverride},
     get_code_key, get_known_code_key, get_nonce_key,
     utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
-    AccountTreeId, StorageKey, StorageValue, H256,
+    AccountTreeId, StorageKey, H256,
 };
 use zksync_utils::{h256_to_u256, u256_to_h256};
 
-/// A storage view that allows to override some of the storage values.
-#[derive(Debug)]
-pub(super) struct StorageWithOverrides<S> {
-    storage_handle: S,
-    overridden_slots: HashMap<StorageKey, H256>,
-    overridden_factory_deps: HashMap<H256, Vec<u8>>,
-    overridden_accounts: HashSet<AccountTreeId>,
-}
-
-impl<S: ReadStorage> StorageWithOverrides<S> {
-    /// Creates a new storage view based on the underlying storage.
-    pub(super) fn new(storage: S, state_override: &StateOverride) -> Self {
-        let mut this = Self {
-            storage_handle: storage,
-            overridden_slots: HashMap::new(),
-            overridden_factory_deps: HashMap::new(),
-            overridden_accounts: HashSet::new(),
-        };
-        this.apply_state_override(state_override);
-        this
-    }
-
-    fn apply_state_override(&mut self, state_override: &StateOverride) {
-        for (account, overrides) in state_override.iter() {
-            if let Some(balance) = overrides.balance {
-                let balance_key = storage_key_for_eth_balance(account);
-                self.overridden_slots
-                    .insert(balance_key, u256_to_h256(balance));
-            }
+/// This method is blocking.
+pub(super) fn apply_state_override<S: ReadStorage>(
+    storage: S,
+    state_override: &StateOverride,
+) -> StorageWithOverrides<S> {
+    let mut storage = StorageWithOverrides::new(storage);
+    for (account, overrides) in state_override.iter() {
+        if let Some(balance) = overrides.balance {
+            let balance_key = storage_key_for_eth_balance(account);
+            storage.set_value(balance_key, u256_to_h256(balance));
+        }
 
-            if let Some(nonce) = overrides.nonce {
-                let nonce_key = get_nonce_key(account);
-                let full_nonce = self.read_value(&nonce_key);
-                let (_, deployment_nonce) = decompose_full_nonce(h256_to_u256(full_nonce));
-                let new_full_nonce = u256_to_h256(nonces_to_full_nonce(nonce, deployment_nonce));
-                self.overridden_slots.insert(nonce_key, new_full_nonce);
-            }
+        if let Some(nonce) = overrides.nonce {
+            let nonce_key = get_nonce_key(account);
+            let full_nonce = storage.read_value(&nonce_key);
+            let (_, deployment_nonce) = decompose_full_nonce(h256_to_u256(full_nonce));
+            let new_full_nonce = u256_to_h256(nonces_to_full_nonce(nonce, deployment_nonce));
+            storage.set_value(nonce_key, new_full_nonce);
+        }
 
-            if let Some(code) = &overrides.code {
-                let code_key = get_code_key(account);
-                let code_hash = code.hash();
-                self.overridden_slots.insert(code_key, code_hash);
-                let known_code_key = get_known_code_key(&code_hash);
-                self.overridden_slots
-                    .insert(known_code_key, H256::from_low_u64_be(1));
-                self.store_factory_dep(code_hash, code.clone().into_bytes());
-            }
+        if let Some(code) = &overrides.code {
+            let code_key = get_code_key(account);
+            let code_hash = code.hash();
+            storage.set_value(code_key, code_hash);
+            let known_code_key = get_known_code_key(&code_hash);
+            storage.set_value(known_code_key, H256::from_low_u64_be(1));
+            storage.store_factory_dep(code_hash, code.clone().into_bytes());
+        }
 
-            match &overrides.state {
-                Some(OverrideState::State(state)) => {
-                    let account = AccountTreeId::new(*account);
-                    self.override_account_state_diff(account, state);
-                    self.overridden_accounts.insert(account);
+        match &overrides.state {
+            Some(OverrideState::State(state)) => {
+                let account = AccountTreeId::new(*account);
+                for (&key, &value) in state {
+                    storage.set_value(StorageKey::new(account, key), value);
                 }
-                Some(OverrideState::StateDiff(state_diff)) => {
-                    let account = AccountTreeId::new(*account);
-                    self.override_account_state_diff(account, state_diff);
+                storage.insert_erased_account(account);
+            }
+            Some(OverrideState::StateDiff(state_diff)) => {
+                let account = AccountTreeId::new(*account);
+                for (&key, &value) in state_diff {
+                    storage.set_value(StorageKey::new(account, key), value);
                 }
-                None => { /* do nothing */ }
             }
+            None => { /* do nothing */ }
         }
     }
-
-    fn store_factory_dep(&mut self, hash: H256, code: Vec<u8>) {
-        self.overridden_factory_deps.insert(hash, code);
-    }
-
-    fn override_account_state_diff(
-        &mut self,
-        account: AccountTreeId,
-        state_diff: &HashMap<H256, H256>,
-    ) {
-        let account_slots = state_diff
-            .iter()
-            .map(|(&slot, &value)| (StorageKey::new(account, slot), value));
-        self.overridden_slots.extend(account_slots);
-    }
-}
-
-impl<S: ReadStorage + fmt::Debug> ReadStorage for StorageWithOverrides<S> {
-    fn read_value(&mut self, key: &StorageKey) -> StorageValue {
-        if let Some(value) = self.overridden_slots.get(key) {
-            return *value;
-        }
-        if self.overridden_accounts.contains(key.account()) {
-            return H256::zero();
-        }
-        self.storage_handle.read_value(key)
-    }
-
-    fn is_write_initial(&mut self, key: &StorageKey) -> bool {
-        self.storage_handle.is_write_initial(key)
-    }
-
-    fn load_factory_dep(&mut self, hash: H256) -> Option<Vec<u8>> {
-        self.overridden_factory_deps
-            .get(&hash)
-            .cloned()
-            .or_else(|| self.storage_handle.load_factory_dep(hash))
-    }
-
-    fn get_enumeration_index(&mut self, key: &StorageKey) -> Option<u64> {
-        self.storage_handle.get_enumeration_index(key)
-    }
+    storage
 }
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashMap;
+
     use zksync_multivm::interface::storage::InMemoryStorage;
     use zksync_types::{
         api::state_override::{Bytecode, OverrideAccount},
@@ -184,7 +124,7 @@ mod tests {
         storage.set_value(retained_key, H256::repeat_byte(0xfe));
         let erased_key = StorageKey::new(AccountTreeId::new(Address::repeat_byte(5)), H256::zero());
         storage.set_value(erased_key, H256::repeat_byte(1));
-        let mut storage = StorageWithOverrides::new(storage, &overrides);
+        let mut storage = apply_state_override(storage, &overrides);
 
         let balance = storage.read_value(&storage_key_for_eth_balance(&Address::repeat_byte(1)));
         assert_eq!(balance, H256::from_low_u64_be(1));