reduce usage of ForkStorage outside inner module

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

@@ -221,7 +221,7 @@ async fn main() -> anyhow::Result<()> {
     let system_contracts =
         SystemContracts::from_options(&config.system_contracts_options, config.use_evm_emulator);
 
-    let (node_inner, _fork_storage, blockchain, time) = InMemoryNodeInner::init(
+    let (node_inner, storage, blockchain, time) = InMemoryNodeInner::init(
         fork_details,
         fee_input_provider.clone(),
         filters,
@@ -245,6 +245,7 @@ async fn main() -> anyhow::Result<()> {
     let node: InMemoryNode = InMemoryNode::new(
         node_inner,
         blockchain,
+        storage,
         node_handle,
         Some(observability),
         time,

crates/core/src/node/debug.rs

@@ -68,7 +68,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

@@ -75,7 +75,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)?;
         // Impersonation does not matter in this context so we assume the tx is not impersonated:
@@ -104,10 +104,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());
@@ -187,9 +184,7 @@ impl InMemoryNode {
             AccountTreeId::new(L2_BASE_TOKEN_ADDRESS),
             &address,
         );
-
-        let inner_guard = self.inner.read().await;
-        match inner_guard.fork_storage.read_value_internal(&balance_key) {
+        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}")),
         }
@@ -258,12 +253,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))
@@ -280,10 +272,8 @@ impl InMemoryNode {
         // TODO: Support
         _block: Option<BlockIdVariant>,
     ) -> anyhow::Result<U256> {
-        let inner = self.inner.read().await;
         let nonce_key = get_nonce_key(&address);
-
-        match inner.fork_storage.read_value_internal(&nonce_key) {
+        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

@@ -2,14 +2,16 @@
 use super::inner::fork::ForkDetails;
 use super::inner::node_executor::NodeExecutorHandle;
 use super::inner::InMemoryNodeInner;
-use crate::deps::{storage_view::StorageView, InMemoryStorage};
+use crate::deps::storage_view::StorageView;
+use crate::deps::InMemoryStorage;
 use crate::filters::EthFilters;
 use crate::formatter;
 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;
@@ -46,17 +48,15 @@ use zksync_multivm::vm_latest::{HistoryDisabled, ToTracerPointer, Vm};
 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).
@@ -244,6 +244,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>>>,
@@ -261,6 +262,7 @@ impl InMemoryNode {
     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>,
@@ -272,6 +274,7 @@ impl InMemoryNode {
         InMemoryNode {
             inner,
             blockchain,
+            storage,
             node_handle,
             snapshots: Default::default(),
             time,
@@ -380,7 +383,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: Vm<_, HistoryDisabled> = Vm::new(batch_env, system_env, storage);
 
         // We must inject *some* signature (otherwise bootloader code fails to generate hash).
@@ -455,22 +458,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> {
@@ -605,6 +596,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)> {
@@ -636,7 +631,7 @@ impl InMemoryNode {
             &config.system_contracts_options,
             config.use_evm_emulator,
         );
-        let (inner, _, blockchain, time) = InMemoryNodeInner::init(
+        let (inner, storage, blockchain, time) = InMemoryNodeInner::init(
             fork,
             fee_provider,
             Arc::new(RwLock::new(Default::default())),
@@ -661,6 +656,7 @@ impl InMemoryNode {
         Self::new(
             inner,
             blockchain,
+            storage,
             node_handle,
             None,
             time,