feat: add state override for gas estimates

core/lib/state/src/lib.rs

@@ -22,6 +22,7 @@ mod in_memory;
 mod postgres;
 mod rocksdb;
 mod shadow_storage;
+mod storage_overrides;
 mod storage_view;
 #[cfg(test)]
 mod test_utils;
@@ -32,6 +33,7 @@ pub use self::{
     postgres::{PostgresStorage, PostgresStorageCaches},
     rocksdb::{RocksbStorageBuilder, RocksdbStorage},
     shadow_storage::ShadowStorage,
+    storage_overrides::StorageOverrides,
     storage_view::{StorageView, StorageViewMetrics},
     witness::WitnessStorage,
 };

core/lib/state/src/storage_overrides.rs

@@ -0,0 +1,145 @@
+use std::{cell::RefCell, collections::HashMap, fmt, rc::Rc};
+
+use zksync_types::{
+    api::{OverrideState, StateOverride},
+    get_code_key, get_nonce_key,
+    utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
+    AccountTreeId, StorageKey, StorageValue, H256,
+};
+use zksync_utils::{bytecode::hash_bytecode, h256_to_u256, u256_to_h256};
+
+use crate::{ReadStorage, StorageView, StorageViewMetrics, WriteStorage};
+
+/// A storage view that allows to override some of the storage values.
+#[derive(Debug)]
+pub struct StorageOverrides<S> {
+    storage_view: StorageView<S>,
+    overrided_factory_deps: HashMap<H256, Vec<u8>>,
+    overrided_account_state: HashMap<AccountTreeId, HashMap<H256, H256>>,
+}
+impl<S: ReadStorage + fmt::Debug> StorageOverrides<S> {
+    /// Applies the state override to the storage.
+    pub 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.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 = nonces_to_full_nonce(nonce, deployment_nonce);
+
+                self.set_value(nonce_key, u256_to_h256(new_full_nonce));
+            }
+
+            if let Some(code) = &overrides.code {
+                let code_key = get_code_key(account);
+                let code_hash = hash_bytecode(&code.0);
+
+                self.set_value(code_key, code_hash);
+                self.store_factory_dep(code_hash, code.0.clone());
+            }
+
+            match &overrides.state {
+                Some(OverrideState::State(state)) => {
+                    self.override_account_state(AccountTreeId::new(account.clone()), state.clone());
+                }
+                Some(OverrideState::StateDiff(state_diff)) => {
+                    for (key, value) in state_diff.iter() {
+                        self.set_value(
+                            StorageKey::new(AccountTreeId::new(account.clone()), key.clone()),
+                            value.clone(),
+                        );
+                    }
+                }
+                None => {}
+            }
+        }
+    }
+}
+
+impl<S: ReadStorage + fmt::Debug> StorageOverrides<S> {
+    /// Creates a new storage view based on the underlying storage.
+    pub fn new(storage_view: StorageView<S>) -> Self {
+        Self {
+            storage_view,
+            overrided_factory_deps: HashMap::new(),
+            overrided_account_state: HashMap::new(),
+        }
+    }
+
+    /// Overrides a factory dep code.
+    pub fn store_factory_dep(&mut self, hash: H256, code: Vec<u8>) {
+        self.overrided_factory_deps.insert(hash, code);
+    }
+
+    /// Overrides an account entire state.
+    pub fn override_account_state(&mut self, account: AccountTreeId, state: HashMap<H256, H256>) {
+        self.overrided_account_state.insert(account, state);
+    }
+
+    /// Returns the current metrics.
+    pub fn metrics(&self) -> StorageViewMetrics {
+        self.storage_view.metrics()
+    }
+
+    /// Make a Rc RefCell ptr to the storage
+    pub fn to_rc_ptr(self) -> Rc<RefCell<Self>> {
+        Rc::new(RefCell::new(self))
+    }
+}
+
+impl<S: ReadStorage + fmt::Debug> Into<Rc<RefCell<Self>>> for StorageOverrides<S> {
+    fn into(self) -> Rc<RefCell<Self>> {
+        self.to_rc_ptr()
+    }
+}
+
+impl<S: ReadStorage + fmt::Debug> ReadStorage for StorageOverrides<S> {
+    fn read_value(&mut self, key: &StorageKey) -> StorageValue {
+        // If the account is overrided, return the overrided value if any or zero.
+        // Otherwise, return the value from the underlying storage.
+        self.overrided_account_state
+            .get(key.account())
+            .and_then(|state| match state.get(key.key()) {
+                Some(v) => Some(v.clone()),
+                None => Some(H256::zero()),
+            })
+            .unwrap_or_else(|| self.storage_view.read_value(key))
+    }
+
+    /// Only keys contained in the underlying storage will return `false`. If a key was
+    /// inserted using [`Self::set_value()`], it will still return `true`.
+    fn is_write_initial(&mut self, key: &StorageKey) -> bool {
+        self.storage_view.is_write_initial(key)
+    }
+
+    fn load_factory_dep(&mut self, hash: H256) -> Option<Vec<u8>> {
+        self.overrided_factory_deps
+            .get(&hash)
+            .cloned()
+            .or_else(|| self.storage_view.load_factory_dep(hash))
+    }
+
+    fn get_enumeration_index(&mut self, key: &StorageKey) -> Option<u64> {
+        self.storage_view.get_enumeration_index(key)
+    }
+}
+
+impl<S: ReadStorage + fmt::Debug> WriteStorage for StorageOverrides<S> {
+    fn set_value(&mut self, key: StorageKey, value: StorageValue) -> StorageValue {
+        self.storage_view.set_value(key, value)
+    }
+
+    fn modified_storage_keys(&self) -> &HashMap<StorageKey, StorageValue> {
+        self.storage_view.modified_storage_keys()
+    }
+
+    fn missed_storage_invocations(&self) -> usize {
+        self.storage_view.missed_storage_invocations()
+    }
+}

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

@@ -1,3 +1,5 @@
+use std::{collections::HashMap, ops::Deref};
+
 use chrono::{DateTime, Utc};
 use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
 use strum::Display;
@@ -733,3 +735,67 @@ pub struct Proof {
     pub address: Address,
     pub storage_proof: Vec<StorageProof>,
 }
+
+/// Collection of overridden accounts, useful for `eth_estimateGas`.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct 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>,
+    #[serde(flatten, deserialize_with = "state_deserializer")]
+    pub state: Option<OverrideState>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub enum OverrideState {
+    State(HashMap<H256, H256>),
+    StateDiff(HashMap<H256, H256>),
+}
+
+fn state_deserializer<'de, D>(deserializer: D) -> Result<Option<OverrideState>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    let val = serde_json::Value::deserialize(deserializer)?;
+    let state: Option<HashMap<H256, H256>> = match val.get("state") {
+        Some(val) => serde_json::from_value(val.clone()).map_err(serde::de::Error::custom)?,
+        None => None,
+    };
+    let state_diff: Option<HashMap<H256, H256>> = match val.get("stateDiff") {
+        Some(val) => serde_json::from_value(val.clone()).map_err(serde::de::Error::custom)?,
+        None => None,
+    };
+
+    match (state, state_diff) {
+        (Some(state), None) => Ok(Some(OverrideState::State(state))),
+        (None, Some(state_diff)) => Ok(Some(OverrideState::StateDiff(state_diff))),
+        (None, None) => Ok(None),
+        _ => Err(serde::de::Error::custom(
+            "Both 'state' and 'stateDiff' cannot be set simultaneously",
+        )),
+    }
+}
+
+impl StateOverride {
+    pub fn new(state: HashMap<Address, OverrideAccount>) -> Self {
+        Self(state)
+    }
+
+    pub fn get(&self, address: &Address) -> Option<&OverrideAccount> {
+        self.0.get(address)
+    }
+}
+
+impl Deref for StateOverride {
+    type Target = HashMap<Address, OverrideAccount>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}

core/lib/types/src/transaction_request.rs

@@ -384,7 +384,9 @@ impl TransactionRequest {
     }
 
     // returns packed eth signature if it is present
-    fn get_packed_signature(&self) -> Result<PackedEthSignature, SerializationTransactionError> {
+    pub fn get_packed_signature(
+        &self,
+    ) -> Result<PackedEthSignature, SerializationTransactionError> {
         let packed_v = self
             .v
             .ok_or(SerializationTransactionError::IncompleteSignature)?

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/Cargo.toml

@@ -64,6 +64,7 @@ itertools = "0.10.3"
 metrics = "0.21"
 ctrlc = { version = "3.1", features = ["termination"] }
 rand = "0.8"
+rlp = "0.5"
 
 tokio = { version = "1", features = ["time"] }
 futures = { version = "0.3", features = ["compat"] }

core/lib/zksync_core/src/api_server/execution_sandbox/apply.rs

@@ -6,7 +6,11 @@
 //!
 //! This module is intended to be blocking.
 
-use std::time::{Duration, Instant};
+use std::{
+    cell::RefCell,
+    rc::Rc,
+    time::{Duration, Instant},
+};
 
 use anyhow::Context as _;
 use multivm::{
@@ -17,13 +21,15 @@ use multivm::{
 };
 use tokio::runtime::Handle;
 use zksync_dal::{ConnectionPool, StorageProcessor};
-use zksync_state::{PostgresStorage, ReadStorage, StoragePtr, StorageView, WriteStorage};
+use zksync_state::{
+    PostgresStorage, ReadStorage, StorageOverrides, StoragePtr, StorageView, WriteStorage,
+};
 use zksync_system_constants::{
     SYSTEM_CONTEXT_ADDRESS, SYSTEM_CONTEXT_CURRENT_L2_BLOCK_INFO_POSITION,
     SYSTEM_CONTEXT_CURRENT_TX_ROLLING_HASH_POSITION, ZKPORTER_IS_AVAILABLE,
 };
 use zksync_types::{
-    api,
+    api::{self, StateOverride},
     block::{pack_block_info, unpack_block_info, MiniblockHasher},
     fee_model::BatchFeeInput,
     get_nonce_key,
@@ -38,7 +44,7 @@ use super::{
     BlockArgs, TxExecutionArgs, TxSharedArgs, VmPermit,
 };
 
-type BoxedVm<'a> = Box<VmInstance<StorageView<PostgresStorage<'a>>, HistoryDisabled>>;
+type BoxedVm<'a> = Box<VmInstance<StorageOverrides<PostgresStorage<'a>>, HistoryDisabled>>;
 
 #[derive(Debug)]
 struct Sandbox<'a> {
@@ -259,7 +265,10 @@ impl<'a> Sandbox<'a> {
         mut self,
         tx: &Transaction,
         adjust_pubdata_price: bool,
-    ) -> (BoxedVm<'a>, StoragePtr<StorageView<PostgresStorage<'a>>>) {
+    ) -> (
+        BoxedVm<'a>,
+        StoragePtr<StorageOverrides<PostgresStorage<'a>>>,
+    ) {
         self.setup_storage_view(tx);
         let protocol_version = self.system_env.version;
         if adjust_pubdata_price {
@@ -271,14 +280,15 @@ impl<'a> Sandbox<'a> {
             );
         };
 
-        let storage_view = self.storage_view.to_rc_ptr();
+        let storage_overrides: Rc<RefCell<StorageOverrides<PostgresStorage<'_>>>> =
+            StorageOverrides::new(self.storage_view).into();
         let vm = Box::new(VmInstance::new_with_specific_version(
             self.l1_batch_env,
             self.system_env,
-            storage_view.clone(),
+            storage_overrides.clone(),
             protocol_version.into_api_vm_version(),
         ));
-        (vm, storage_view)
+        (vm, storage_overrides)
     }
 }
 
@@ -294,8 +304,9 @@ pub(super) fn apply_vm_in_sandbox<T>(
     connection_pool: &ConnectionPool,
     tx: Transaction,
     block_args: BlockArgs,
+    state_override: Option<StateOverride>,
     apply: impl FnOnce(
-        &mut VmInstance<StorageView<PostgresStorage<'_>>, HistoryDisabled>,
+        &mut VmInstance<StorageOverrides<PostgresStorage<'_>>, HistoryDisabled>,
         Transaction,
     ) -> T,
 ) -> anyhow::Result<T> {
@@ -328,6 +339,13 @@ pub(super) fn apply_vm_in_sandbox<T>(
         tx.initiator_account(),
         tx.nonce().unwrap_or(Nonce(0))
     );
+
+    // Apply state override
+    if let Some(state_override) = state_override {
+        let mut storage_override = storage_view.as_ref().borrow_mut();
+        storage_override.apply_state_override(state_override);
+    }
+
     let execution_latency = SANDBOX_METRICS.sandbox[&SandboxStage::Execution].start();
     let result = apply(&mut vm, tx);
     let vm_execution_took = execution_latency.observe();