feat: add state override for gas estimates

core/lib/types/src/api/mod.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use chrono::{DateTime, Utc};
 use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
 use strum::Display;
@@ -733,3 +735,17 @@ pub struct Proof {
     pub address: Address,
     pub storage_proof: Vec<StorageProof>,
 }
+
+/// Collection of overridden accounts, useful for `eth_estimateGas`.
+pub type StateOverride = HashMap<Address, OverrideAccount>;
+
+/// Account override for `eth_estimateGas`.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct OverrideAccount {
+    pub balance: Option<U256>,
+    pub nonce: Option<U256>,
+    pub code: Option<Bytes>,
+    pub state: Option<HashMap<H256, H256>>,
+    pub state_diff: Option<HashMap<H256, H256>>,
+}

core/lib/web3_decl/src/namespaces/eth.rs

@@ -3,7 +3,7 @@ use jsonrpsee::{
     proc_macros::rpc,
 };
 use zksync_types::{
-    api::{BlockId, BlockIdVariant, BlockNumber, Transaction, TransactionVariant},
+    api::{BlockId, BlockIdVariant, BlockNumber, StateOverride, Transaction, TransactionVariant},
     transaction_request::CallRequest,
     Address, H256,
 };
@@ -36,7 +36,12 @@ pub trait EthNamespace {
     async fn call(&self, req: CallRequest, block: Option<BlockIdVariant>) -> RpcResult<Bytes>;
 
     #[method(name = "estimateGas")]
-    async fn estimate_gas(&self, req: CallRequest, _block: Option<BlockNumber>) -> RpcResult<U256>;
+    async fn estimate_gas(
+        &self,
+        req: CallRequest,
+        _block: Option<BlockNumber>,
+        state_override: Option<StateOverride>,
+    ) -> RpcResult<U256>;
 
     #[method(name = "gasPrice")]
     async fn gas_price(&self) -> RpcResult<U256>;

core/lib/zksync_core/src/api_server/tx_sender/mod.rs

@@ -1,8 +1,9 @@
 //! Helper module to submit transactions into the zkSync Network.
 
-use std::{cmp, sync::Arc, time::Instant};
+use std::{cmp, collections::HashMap, sync::Arc, time::Instant};
 
 use anyhow::Context as _;
+use itertools::Itertools;
 use multivm::{
     interface::VmExecutionResultAndLogs,
     utils::{adjust_pubdata_price_for_tx, derive_base_fee_and_gas_per_pubdata, derive_overhead},
@@ -13,17 +14,18 @@ use zksync_contracts::BaseSystemContracts;
 use zksync_dal::{transactions_dal::L2TxSubmissionResult, ConnectionPool, StorageProcessor};
 use zksync_state::PostgresStorageCaches;
 use zksync_types::{
+    api::StateOverride,
     fee::{Fee, TransactionExecutionMetrics},
     fee_model::BatchFeeInput,
-    get_code_key, get_intrinsic_constants,
+    get_code_key, get_intrinsic_constants, get_nonce_key,
     l1::is_l1_tx_type,
     l2::{error::TxCheckError::TxDuplication, L2Tx},
-    utils::storage_key_for_eth_balance,
+    utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
     AccountTreeId, Address, ExecuteTransactionCommon, L2ChainId, MiniblockNumber, Nonce,
-    PackedEthSignature, ProtocolVersionId, Transaction, VmVersion, H160, H256, MAX_L2_TX_GAS_LIMIT,
-    MAX_NEW_FACTORY_DEPS, U256,
+    PackedEthSignature, ProtocolVersionId, StorageKey, StorageLog, Transaction, VmVersion, H160,
+    H256, MAX_L2_TX_GAS_LIMIT, MAX_NEW_FACTORY_DEPS, U256,
 };
-use zksync_utils::h256_to_u256;
+use zksync_utils::{bytecode::hash_bytecode, h256_to_u256, u256_to_h256};
 
 pub(super) use self::result::SubmitTxError;
 use self::tx_sink::TxSink;
@@ -629,6 +631,7 @@ impl TxSender {
         mut tx: Transaction,
         estimated_fee_scale_factor: f64,
         acceptable_overestimation: u32,
+        state_override: Option<StateOverride>,
     ) -> Result<Fee, SubmitTxError> {
         let estimation_started_at = Instant::now();
 
@@ -690,6 +693,140 @@ impl TxSender {
                 )
             })?;
 
+        if let Some(overrides) = &state_override {
+            for (account, overrides) in overrides {
+                if overrides.state.is_some() && overrides.state_diff.is_some() {
+                    return Err(SubmitTxError::Unexecutable(
+                        "state and stateDiff are mutually exclusive".to_string(),
+                    ));
+                }
+
+                if let Some(balance) = overrides.balance {
+                    let balance_key = storage_key_for_eth_balance(account);
+                    let log = StorageLog::new_write_log(balance_key, u256_to_h256(balance));
+                    self.acquire_replica_connection()
+                        .await?
+                        .storage_logs_dal()
+                        .append_storage_logs(
+                            block_args.resolved_block_number(),
+                            &[(H256::zero(), vec![log])],
+                        )
+                        .await
+                        .map_err(|err| SubmitTxError::Internal(err.into()))?;
+                }
+
+                if let Some(nonce) = overrides.nonce {
+                    let nonce_key = get_nonce_key(account);
+                    let mut connection = self.acquire_replica_connection().await?;
+
+                    let full_nonce = connection
+                        .storage_web3_dal()
+                        .get_value(&nonce_key)
+                        .await
+                        .map_err(|err| SubmitTxError::Internal(err.into()))?;
+                    let (_, deployment_nonce) = decompose_full_nonce(h256_to_u256(full_nonce));
+                    let new_full_nonce = nonces_to_full_nonce(nonce, deployment_nonce);
+
+                    let log = StorageLog::new_write_log(nonce_key, u256_to_h256(new_full_nonce));
+                    connection
+                        .storage_logs_dal()
+                        .append_storage_logs(
+                            block_args.resolved_block_number(),
+                            &[(H256::zero(), vec![log])],
+                        )
+                        .await
+                        .map_err(|err| SubmitTxError::Internal(err.into()))?;
+                }
+
+                if let Some(code) = &overrides.code {
+                    let code_key = get_code_key(account);
+                    let code_hash = hash_bytecode(&code.0);
+
+                    let mut connection = self.acquire_replica_connection().await?;
+
+                    let log = StorageLog::new_write_log(code_key, code_hash);
+                    connection
+                        .storage_logs_dal()
+                        .append_storage_logs(
+                            block_args.resolved_block_number(),
+                            &[(H256::zero(), vec![log])],
+                        )
+                        .await
+                        .map_err(|err| SubmitTxError::Internal(err.into()))?;
+                    connection
+                        .factory_deps_dal()
+                        .insert_factory_deps(
+                            block_args.resolved_block_number(),
+                            &HashMap::from_iter(vec![(code_hash, code.0.clone())]),
+                        )
+                        .await
+                        .map_err(|err| SubmitTxError::Internal(err.into()))?;
+                }
+
+                if let Some(state) = &overrides.state {
+                    let mut connection = self.acquire_replica_connection().await?;
+
+                    // First find storage keys to clear
+                    let keys_to_clear = connection
+                        .storage_logs_dal()
+                        .dump_all_storage_logs_for_tests()
+                        .await
+                        .iter()
+                        .filter(|log| &log.address == account)
+                        .map(|log| log.key)
+                        .unique()
+                        .collect::<Vec<_>>();
+
+                    // Now clear keys and add new values
+                    let mut logs: Vec<StorageLog> = keys_to_clear
+                        .iter()
+                        .map(|key| {
+                            StorageLog::new_write_log(
+                                StorageKey::new(AccountTreeId::new(*account), *key),
+                                H256::zero(),
+                            )
+                        })
+                        .collect();
+                    logs.extend(state.iter().map(|(key, value)| {
+                        StorageLog::new_write_log(
+                            StorageKey::new(AccountTreeId::new(*account), *key),
+                            *value,
+                        )
+                    }));
+
+                    connection
+                        .storage_logs_dal()
+                        .append_storage_logs(
+                            block_args.resolved_block_number(),
+                            &[(H256::zero(), logs)],
+                        )
+                        .await
+                        .map_err(|err| SubmitTxError::Internal(err.into()))?;
+                }
+
+                if let Some(state_diff) = &overrides.state_diff {
+                    let logs = state_diff
+                        .iter()
+                        .map(|(key, value)| {
+                            StorageLog::new_write_log(
+                                StorageKey::new(AccountTreeId::new(*account), *key),
+                                *value,
+                            )
+                        })
+                        .collect();
+                    self.acquire_replica_connection()
+                        .await?
+                        .storage_logs_dal()
+                        .append_storage_logs(
+                            block_args.resolved_block_number(),
+                            &[(H256::zero(), logs)],
+                        )
+                        .await
+                        .map_err(|err| SubmitTxError::Internal(err.into()))?;
+                }
+            }
+        }
+
         if !tx.is_l1()
             && account_code_hash == H256::zero()
             && tx.execute.value > self.get_balance(&tx.initiator_account()).await?

core/lib/zksync_core/src/api_server/web3/backend_jsonrpsee/namespaces/eth.rs

@@ -1,7 +1,7 @@
 use zksync_types::{
     api::{
-        Block, BlockId, BlockIdVariant, BlockNumber, Log, Transaction, TransactionId,
-        TransactionReceipt, TransactionVariant,
+        Block, BlockId, BlockIdVariant, BlockNumber, Log, StateOverride, Transaction,
+        TransactionId, TransactionReceipt, TransactionVariant,
     },
     transaction_request::CallRequest,
     web3::types::{FeeHistory, Index, SyncState},
@@ -33,8 +33,13 @@ impl EthNamespaceServer for EthNamespace {
             .map_err(|err| self.current_method().map_err(err))
     }
 
-    async fn estimate_gas(&self, req: CallRequest, block: Option<BlockNumber>) -> RpcResult<U256> {
-        self.estimate_gas_impl(req, block)
+    async fn estimate_gas(
+        &self,
+        req: CallRequest,
+        block: Option<BlockNumber>,
+        state_override: Option<StateOverride>,
+    ) -> RpcResult<U256> {
+        self.estimate_gas_impl(req, block, state_override)
             .await
             .map_err(|err| self.current_method().map_err(err))
     }

core/lib/zksync_core/src/api_server/web3/namespaces/eth.rs

@@ -2,8 +2,8 @@ use anyhow::Context as _;
 use zksync_system_constants::DEFAULT_L2_TX_GAS_PER_PUBDATA_BYTE;
 use zksync_types::{
     api::{
-        BlockId, BlockNumber, GetLogsFilter, Transaction, TransactionId, TransactionReceipt,
-        TransactionVariant,
+        BlockId, BlockNumber, GetLogsFilter, StateOverride, Transaction, TransactionId,
+        TransactionReceipt, TransactionVariant,
     },
     l2::{L2Tx, TransactionType},
     transaction_request::CallRequest,
@@ -92,6 +92,7 @@ impl EthNamespace {
         &self,
         request: CallRequest,
         _block: Option<BlockNumber>,
+        state_override: Option<StateOverride>,
     ) -> Result<U256, Web3Error> {
         let mut request_with_gas_per_pubdata_overridden = request;
         self.state
@@ -133,7 +134,12 @@ impl EthNamespace {
         let fee = self
             .state
             .tx_sender
-            .get_txs_fee_in_wei(tx.into(), scale_factor, acceptable_overestimation)
+            .get_txs_fee_in_wei(
+                tx.into(),
+                scale_factor,
+                acceptable_overestimation,
+                state_override,
+            )
             .await?;
         Ok(fee.gas_limit)
     }

core/lib/zksync_core/src/api_server/web3/namespaces/zks.rs

@@ -106,7 +106,7 @@ impl ZksNamespace {
         Ok(self
             .state
             .tx_sender
-            .get_txs_fee_in_wei(tx, scale_factor, acceptable_overestimation)
+            .get_txs_fee_in_wei(tx, scale_factor, acceptable_overestimation, None)
             .await?)
     }