refactor: reduce usage of ForkStorage outside inner module (#539)

* reduce usage of `ForkStorage` outside `inner` module

* adapt for zkos

crates/api_server/src/impls/anvil.rs

@@ -155,11 +155,19 @@ impl AnvilNamespaceServer for AnvilNamespace {
     }
 
     async fn set_balance(&self, address: Address, balance: U256) -> RpcResult<bool> {
-        Ok(self.node.set_balance(address, balance).await)
+        Ok(self
+            .node
+            .set_balance(address, balance)
+            .await
+            .map_err(RpcError::from)?)
     }
 
     async fn set_nonce(&self, address: Address, nonce: U256) -> RpcResult<bool> {
-        Ok(self.node.set_nonce(address, nonce).await)
+        Ok(self
+            .node
+            .set_nonce(address, nonce)
+            .await
+            .map_err(RpcError::from)?)
     }
 
     async fn anvil_mine(&self, num_blocks: Option<U64>, interval: Option<U64>) -> RpcResult<()> {
@@ -203,7 +211,11 @@ impl AnvilNamespaceServer for AnvilNamespace {
     }
 
     async fn set_storage_at(&self, address: Address, slot: U256, value: U256) -> RpcResult<bool> {
-        Ok(self.node.set_storage_at(address, slot, value).await)
+        Ok(self
+            .node
+            .set_storage_at(address, slot, value)
+            .await
+            .map_err(RpcError::from)?)
     }
 
     async fn set_chain_id(&self, id: u32) -> RpcResult<()> {

crates/cli/src/bytecode_override.rs

@@ -41,7 +41,7 @@ pub async fn override_bytecodes(node: &InMemoryNode, bytecodes_dir: String) -> a
                 let bytecode = Vec::from_hex(contract.bytecode.object)
                     .with_context(|| format!("Failed to parse hex from {:?}", path))?;
 
-                node.override_bytecode(&address, &bytecode)
+                node.override_bytecode(address, bytecode)
                     .await
                     .expect("Failed to override bytecode");
                 tracing::info!("+++++ Replacing bytecode at address {:?} +++++", address);

crates/cli/src/main.rs

@@ -13,7 +13,7 @@ use anvil_zksync_core::filters::EthFilters;
 use anvil_zksync_core::node::fork::ForkDetails;
 use anvil_zksync_core::node::{
     BlockSealer, BlockSealerMode, ImpersonationManager, InMemoryNode, InMemoryNodeInner,
-    NodeExecutor, TestNodeFeeInputProvider, TxPool,
+    NodeExecutor, StorageKeyLayout, TestNodeFeeInputProvider, TxPool,
 };
 use anvil_zksync_core::observability::Observability;
 use anvil_zksync_core::system_contracts::SystemContracts;
@@ -223,18 +223,27 @@ async fn main() -> anyhow::Result<()> {
         config.use_evm_emulator,
         config.use_zkos,
     );
+    let storage_key_layout = if config.use_zkos {
+        StorageKeyLayout::ZkOs
+    } else {
+        StorageKeyLayout::ZkEra
+    };
 
-    let (node_inner, _fork_storage, blockchain, time) = InMemoryNodeInner::init(
+    let (node_inner, storage, blockchain, time) = InMemoryNodeInner::init(
         fork_details,
         fee_input_provider.clone(),
         filters,
         config.clone(),
         impersonation.clone(),
         system_contracts.clone(),
+        storage_key_layout,
     );
 
-    let (node_executor, node_handle) =
-        NodeExecutor::new(node_inner.clone(), system_contracts.clone());
+    let (node_executor, node_handle) = NodeExecutor::new(
+        node_inner.clone(),
+        system_contracts.clone(),
+        storage_key_layout,
+    );
     let sealing_mode = if config.no_mining {
         BlockSealerMode::noop()
     } else if let Some(block_time) = config.block_time {
@@ -248,13 +257,15 @@ async fn main() -> anyhow::Result<()> {
     let node: InMemoryNode = InMemoryNode::new(
         node_inner,
         blockchain,
+        storage,
         node_handle,
         Some(observability),
         time,
         impersonation,
         pool,
         block_sealer_state,
         system_contracts,
+        storage_key_layout,
     );
 
     if let Some(ref bytecodes_dir) = config.override_bytecodes_dir {

crates/core/src/node/debug.rs

@@ -71,7 +71,7 @@ impl InMemoryNode {
         // update the enforced_base_fee within l1_batch_env to match the logic in zksync_core
         l1_batch_env.enforced_base_fee = Some(l2_tx.common_data.fee.max_fee_per_gas.as_u64());
         let system_env = inner.create_system_env(system_contracts.clone(), execution_mode);
-        let storage = StorageView::new(&inner.fork_storage).into_rc_ptr();
+        let storage = StorageView::new(inner.read_storage()).into_rc_ptr();
         let mut vm: Vm<_, HistoryDisabled> = Vm::new(l1_batch_env, system_env, storage);
 
         // We must inject *some* signature (otherwise bootloader code fails to generate hash).

crates/core/src/node/eth.rs

@@ -28,8 +28,6 @@ use crate::{
     utils::{h256_to_u64, TransparentError},
 };
 
-use super::keys::StorageKeyLayout;
-
 impl InMemoryNode {
     pub async fn call_impl(
         &self,
@@ -78,7 +76,7 @@ impl InMemoryNode {
     }
 
     pub async fn send_raw_transaction_impl(&self, tx_bytes: Bytes) -> Result<H256, Web3Error> {
-        let chain_id = self.inner.read().await.fork_storage.chain_id;
+        let chain_id = self.chain_id().await;
 
         let (tx_req, hash) = TransactionRequest::from_bytes(&tx_bytes.0, chain_id)?;
         let mut l2_tx =
@@ -102,10 +100,7 @@ impl InMemoryNode {
     ) -> Result<H256, Web3Error> {
         let (chain_id, l2_gas_price) = {
             let reader = self.inner.read().await;
-            (
-                reader.fork_storage.chain_id,
-                reader.fee_input_provider.gas_price(),
-            )
+            (self.chain_id().await, reader.fee_input_provider.gas_price())
         };
 
         let mut tx_req = TransactionRequest::from(tx.clone());
@@ -176,13 +171,10 @@ impl InMemoryNode {
         // TODO: Support
         _block: Option<BlockIdVariant>,
     ) -> anyhow::Result<U256> {
-        let balance_key = StorageKeyLayout::get_storage_key_for_base_token(
-            self.system_contracts.use_zkos,
-            &address,
-        );
-
-        let inner_guard = self.inner.read().await;
-        match inner_guard.fork_storage.read_value_internal(&balance_key) {
+        let balance_key = self
+            .storage_key_layout
+            .get_storage_key_for_base_token(&address);
+        match self.storage.read_value_alt(&balance_key).await {
             Ok(balance) => Ok(h256_to_u256(balance)),
             Err(error) => Err(anyhow::anyhow!("failed to read account balance: {error}")),
         }
@@ -251,12 +243,9 @@ impl InMemoryNode {
         // TODO: Support
         _block: Option<BlockIdVariant>,
     ) -> anyhow::Result<Bytes> {
-        let inner = self.inner.write().await;
-
         let code_key = get_code_key(&address);
-
-        match inner.fork_storage.read_value_internal(&code_key) {
-            Ok(code_hash) => match inner.fork_storage.load_factory_dep_internal(code_hash) {
+        match self.storage.read_value_alt(&code_key).await {
+            Ok(code_hash) => match self.storage.load_factory_dep_alt(code_hash).await {
                 Ok(raw_code) => {
                     let code = raw_code.unwrap_or_default();
                     Ok(Bytes::from(code))
@@ -273,10 +262,8 @@ impl InMemoryNode {
         // TODO: Support
         _block: Option<BlockIdVariant>,
     ) -> anyhow::Result<U256> {
-        let inner = self.inner.read().await;
-        let nonce_key = StorageKeyLayout::get_nonce_key(self.system_contracts.use_zkos, &address);
-
-        match inner.fork_storage.read_value_internal(&nonce_key) {
+        let nonce_key = self.storage_key_layout.get_nonce_key(&address);
+        match self.storage.read_value_alt(&nonce_key).await {
             Ok(result) => Ok(h256_to_u64(result).into()),
             Err(error) => Err(anyhow::anyhow!("failed to read nonce storage: {error}")),
         }

crates/core/src/node/in_memory.rs

@@ -3,13 +3,15 @@ use super::inner::fork::ForkDetails;
 use super::inner::node_executor::NodeExecutorHandle;
 use super::inner::InMemoryNodeInner;
 use super::vm::AnvilVM;
-use crate::deps::{storage_view::StorageView, InMemoryStorage};
+use crate::deps::storage_view::StorageView;
+use crate::deps::InMemoryStorage;
 use crate::filters::EthFilters;
 use crate::node::call_error_tracer::CallErrorTracer;
 use crate::node::error::LoadStateError;
 use crate::node::fee_model::TestNodeFeeInputProvider;
 use crate::node::impersonate::{ImpersonationManager, ImpersonationState};
 use crate::node::inner::blockchain::ReadBlockchain;
+use crate::node::inner::storage::ReadStorageDyn;
 use crate::node::inner::time::ReadTime;
 use crate::node::sealer::BlockSealerState;
 use crate::node::state::VersionedState;
@@ -45,21 +47,20 @@ use zksync_multivm::tracers::CallTracer;
 use zksync_multivm::utils::{get_batch_base_fee, get_max_batch_gas_limit};
 use zksync_multivm::vm_latest::Vm;
 
+use crate::node::keys::StorageKeyLayout;
 use zksync_multivm::vm_latest::{HistoryDisabled, ToTracerPointer};
 use zksync_multivm::VmVersion;
 use zksync_types::api::{Block, DebugCall, TransactionReceipt, TransactionVariant};
 use zksync_types::block::unpack_block_info;
-use zksync_types::bytecode::BytecodeHash;
 use zksync_types::fee_model::BatchFeeInput;
 use zksync_types::l2::L2Tx;
 use zksync_types::storage::{
     EMPTY_UNCLES_HASH, SYSTEM_CONTEXT_ADDRESS, SYSTEM_CONTEXT_BLOCK_INFO_POSITION,
 };
 use zksync_types::web3::{keccak256, Bytes};
-use zksync_types::{get_code_key, h256_to_u256};
 use zksync_types::{
-    AccountTreeId, Address, Bloom, L1BatchNumber, L2BlockNumber, PackedEthSignature, StorageKey,
-    StorageValue, Transaction, H160, H256, H64, U256, U64,
+    h256_to_u256, AccountTreeId, Address, Bloom, L1BatchNumber, L2BlockNumber, L2ChainId,
+    PackedEthSignature, StorageKey, StorageValue, Transaction, H160, H256, H64, U256, U64,
 };
 
 /// Max possible size of an ABI encoded tx (in bytes).
@@ -247,6 +248,7 @@ pub struct InMemoryNode {
     /// A thread safe reference to the [InMemoryNodeInner].
     pub(crate) inner: Arc<RwLock<InMemoryNodeInner>>,
     pub(crate) blockchain: Box<dyn ReadBlockchain>,
+    pub(crate) storage: Box<dyn ReadStorageDyn>,
     pub(crate) node_handle: NodeExecutorHandle,
     /// List of snapshots of the [InMemoryNodeInner]. This is bounded at runtime by [MAX_SNAPSHOTS].
     pub(crate) snapshots: Arc<RwLock<Vec<Snapshot>>>,
@@ -257,24 +259,28 @@ pub struct InMemoryNode {
     pub(crate) pool: TxPool,
     pub(crate) sealer_state: BlockSealerState,
     pub(crate) system_contracts: SystemContracts,
+    pub(crate) storage_key_layout: StorageKeyLayout,
 }
 
 impl InMemoryNode {
     #[allow(clippy::too_many_arguments)]
     pub fn new(
         inner: Arc<RwLock<InMemoryNodeInner>>,
         blockchain: Box<dyn ReadBlockchain>,
+        storage: Box<dyn ReadStorageDyn>,
         node_handle: NodeExecutorHandle,
         observability: Option<Observability>,
         time: Box<dyn ReadTime>,
         impersonation: ImpersonationManager,
         pool: TxPool,
         sealer_state: BlockSealerState,
         system_contracts: SystemContracts,
+        storage_key_layout: StorageKeyLayout,
     ) -> Self {
         InMemoryNode {
             inner,
             blockchain,
+            storage,
             node_handle,
             snapshots: Default::default(),
             time,
@@ -283,6 +289,7 @@ impl InMemoryNode {
             pool,
             sealer_state,
             system_contracts,
+            storage_key_layout,
         }
     }
 
@@ -383,7 +390,7 @@ impl InMemoryNode {
         let (batch_env, _) = inner.create_l1_batch_env().await;
         let system_env = inner.create_system_env(base_contracts, execution_mode);
 
-        let storage = StorageView::new(&inner.fork_storage).into_rc_ptr();
+        let storage = StorageView::new(inner.read_storage()).into_rc_ptr();
 
         let mut vm = if self.system_contracts.use_zkos {
             AnvilVM::ZKOs(super::zkos::ZKOsVM::<_, HistoryDisabled>::new(
@@ -472,22 +479,10 @@ impl InMemoryNode {
     // Forcefully stores the given bytecode at a given account.
     pub async fn override_bytecode(
         &self,
-        address: &Address,
-        bytecode: &[u8],
-    ) -> Result<(), String> {
-        let inner = self.inner.write().await;
-
-        let code_key = get_code_key(address);
-
-        let bytecode_hash = BytecodeHash::for_bytecode(bytecode).value();
-
-        inner
-            .fork_storage
-            .store_factory_dep(bytecode_hash, bytecode.to_owned());
-
-        inner.fork_storage.set_value(code_key, bytecode_hash);
-
-        Ok(())
+        address: Address,
+        bytecode: Vec<u8>,
+    ) -> anyhow::Result<()> {
+        self.node_handle.set_code_sync(address, bytecode).await
     }
 
     pub async fn dump_state(&self, preserve_historical_states: bool) -> anyhow::Result<Bytes> {
@@ -622,6 +617,10 @@ impl InMemoryNode {
         observability.set_logging(directive)?;
         Ok(true)
     }
+
+    pub async fn chain_id(&self) -> L2ChainId {
+        self.inner.read().await.chain_id()
+    }
 }
 
 pub fn load_last_l1_batch<S: ReadStorage>(storage: StoragePtr<S>) -> Option<(u64, u64)> {
@@ -654,16 +653,22 @@ impl InMemoryNode {
             config.use_evm_emulator,
             config.use_zkos,
         );
-        let (inner, _, blockchain, time) = InMemoryNodeInner::init(
+        let storage_key_layout = if config.use_zkos {
+            StorageKeyLayout::ZkOs
+        } else {
+            StorageKeyLayout::ZkEra
+        };
+        let (inner, storage, blockchain, time) = InMemoryNodeInner::init(
             fork,
             fee_provider,
             Arc::new(RwLock::new(Default::default())),
             config,
             impersonation.clone(),
             system_contracts.clone(),
+            storage_key_layout,
         );
         let (node_executor, node_handle) =
-            NodeExecutor::new(inner.clone(), system_contracts.clone());
+            NodeExecutor::new(inner.clone(), system_contracts.clone(), storage_key_layout);
         let pool = TxPool::new(
             impersonation.clone(),
             anvil_zksync_types::TransactionOrder::Fifo,
@@ -679,13 +684,15 @@ impl InMemoryNode {
         Self::new(
             inner,
             blockchain,
+            storage,
             node_handle,
             None,
             time,
             impersonation,
             pool,
             block_sealer_state,
             system_contracts,
+            storage_key_layout,
         )
     }