feat: add state override for gas estimates

core/lib/state/src/lib.rs

@@ -12,6 +12,7 @@
 use std::{cell::RefCell, collections::HashMap, fmt, rc::Rc};
 
 use zksync_types::{
+    api::state_override::StateOverride,
     get_known_code_key,
     storage::{StorageKey, StorageValue},
     H256,
@@ -24,6 +25,7 @@ mod postgres;
 mod rocksdb;
 mod shadow_storage;
 mod storage_factory;
+mod storage_overrides;
 mod storage_view;
 #[cfg(test)]
 mod test_utils;
@@ -40,6 +42,7 @@ pub use self::{
     },
     shadow_storage::ShadowStorage,
     storage_factory::{BatchDiff, PgOrRocksdbStorage, ReadStorageFactory, RocksdbWithMemory},
+    storage_overrides::StorageOverrides,
     storage_view::{StorageView, StorageViewMetrics},
 };
 
@@ -87,3 +90,9 @@ pub trait WriteStorage: ReadStorage {
 
 /// Smart pointer to [`WriteStorage`].
 pub type StoragePtr<S> = Rc<RefCell<S>>;
+
+/// Functionality to override the storage state.
+pub trait OverrideStorage {
+    /// Apply state override to the storage.
+    fn apply_state_override(&mut self, overrides: &StateOverride);
+}

core/lib/state/src/storage_overrides.rs

@@ -0,0 +1,129 @@
+use std::{cell::RefCell, collections::HashMap, fmt, rc::Rc};
+
+use zksync_types::{
+    api::state_override::{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, U256,
+};
+use zksync_utils::{bytecode::hash_bytecode, h256_to_u256, u256_to_h256};
+
+use crate::{OverrideStorage, ReadStorage};
+
+/// A storage view that allows to override some of the storage values.
+#[derive(Debug)]
+pub struct StorageOverrides<S> {
+    storage_handle: S,
+    overrided_factory_deps: HashMap<H256, Vec<u8>>,
+    overrided_account_state: HashMap<AccountTreeId, HashMap<H256, H256>>,
+    overrided_balance: HashMap<StorageKey, U256>,
+    overrided_nonce: HashMap<StorageKey, U256>,
+    overrided_code: HashMap<StorageKey, H256>,
+}
+
+impl<S: ReadStorage + fmt::Debug> StorageOverrides<S> {
+    /// Creates a new storage view based on the underlying storage.
+    pub fn new(storage: S) -> Self {
+        Self {
+            storage_handle: storage,
+            overrided_factory_deps: HashMap::new(),
+            overrided_account_state: HashMap::new(),
+            overrided_balance: HashMap::new(),
+            overrided_nonce: HashMap::new(),
+            overrided_code: HashMap::new(),
+        }
+    }
+
+    /// Overrides a factory dependency 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);
+    }
+
+    /// 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> ReadStorage for StorageOverrides<S> {
+    fn read_value(&mut self, key: &StorageKey) -> StorageValue {
+        if let Some(balance) = self.overrided_balance.get(key) {
+            return u256_to_h256(*balance);
+        }
+        if let Some(code) = self.overrided_code.get(key) {
+            return *code;
+        }
+
+        if let Some(nonce) = self.overrided_nonce.get(key) {
+            return u256_to_h256(*nonce);
+        }
+        // If the account is overridden, return the overridden value if any or zero.
+        // Otherwise, return the value from the underlying storage.
+        self.overrided_account_state.get(key.account()).map_or_else(
+            || self.storage_handle.read_value(key),
+            |state| state.get(key.key()).copied().unwrap_or_else(H256::zero),
+        )
+    }
+
+    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.overrided_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)
+    }
+}
+
+impl<S: ReadStorage + fmt::Debug> OverrideStorage for StorageOverrides<S> {
+    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.overrided_balance.insert(balance_key, 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.overrided_nonce.insert(nonce_key, new_full_nonce);
+            }
+
+            if let Some(code) = &overrides.code {
+                let code_key = get_code_key(account);
+                let code_hash = hash_bytecode(&code.0);
+                self.overrided_code.insert(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), state.clone());
+                }
+                Some(OverrideState::StateDiff(state_diff)) => {
+                    for (key, value) in state_diff {
+                        let account_state = self
+                            .overrided_account_state
+                            .entry(AccountTreeId::new(*account))
+                            .or_default();
+                        account_state.insert(*key, *value);
+                    }
+                }
+                None => {}
+            }
+        }
+    }
+}

core/lib/state/src/storage_view.rs

@@ -6,9 +6,9 @@ use std::{
     time::{Duration, Instant},
 };
 
-use zksync_types::{StorageKey, StorageValue, H256};
+use zksync_types::{api::state_override::StateOverride, StorageKey, StorageValue, H256};
 
-use crate::{ReadStorage, WriteStorage};
+use crate::{OverrideStorage, ReadStorage, WriteStorage};
 
 /// Metrics for [`StorageView`].
 #[derive(Debug, Default, Clone, Copy)]
@@ -52,6 +52,13 @@ pub struct StorageView<S> {
     metrics: StorageViewMetrics,
 }
 
+impl<S> StorageView<S> {
+    /// Returns the modified storage keys
+    pub fn modified_storage_keys(&self) -> &HashMap<StorageKey, StorageValue> {
+        &self.modified_storage_keys
+    }
+}
+
 impl<S> ReadStorage for Box<S>
 where
     S: ReadStorage + ?Sized,
@@ -197,6 +204,12 @@ impl<S: ReadStorage + fmt::Debug> WriteStorage for StorageView<S> {
     }
 }
 
+impl<S: ReadStorage + fmt::Debug + OverrideStorage> OverrideStorage for StorageView<S> {
+    fn apply_state_override(&mut self, state_override: &StateOverride) {
+        self.storage_handle.apply_state_override(state_override);
+    }
+}
+
 #[cfg(test)]
 mod test {
     use zksync_types::{AccountTreeId, Address, H256};

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

@@ -18,6 +18,7 @@ use crate::{
 };
 
 pub mod en;
+pub mod state_override;
 
 /// Block Number
 #[derive(Copy, Clone, Debug, PartialEq, Display)]

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

@@ -0,0 +1,70 @@
+use std::{collections::HashMap, ops::Deref};
+
+use serde::{Deserialize, Deserializer, Serialize};
+use zksync_basic_types::{web3::Bytes, H256, U256};
+
+use crate::Address;
+
+/// 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

@@ -400,7 +400,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/vm_utils/src/lib.rs

@@ -8,22 +8,22 @@ use zksync_multivm::{
     vm_latest::HistoryEnabled,
     VmInstance,
 };
-use zksync_state::{PostgresStorage, StoragePtr, StorageView, WriteStorage};
+use zksync_state::{PostgresStorage, StorageOverrides, StoragePtr, StorageView, WriteStorage};
 use zksync_types::{L1BatchNumber, L2ChainId, Transaction};
 
 use crate::storage::L1BatchParamsProvider;
 
 pub type VmAndStorage<'a> = (
-    VmInstance<StorageView<PostgresStorage<'a>>, HistoryEnabled>,
-    StoragePtr<StorageView<PostgresStorage<'a>>>,
+    VmInstance<StorageView<StorageOverrides<PostgresStorage<'a>>>, HistoryEnabled>,
+    StoragePtr<StorageView<StorageOverrides<PostgresStorage<'a>>>>,
 );
 
 pub fn create_vm(
     rt_handle: Handle,
     l1_batch_number: L1BatchNumber,
     mut connection: Connection<'_, Core>,
     l2_chain_id: L2ChainId,
-) -> anyhow::Result<VmAndStorage> {
+) -> anyhow::Result<VmAndStorage<'_>> {
     let l1_batch_params_provider = rt_handle
         .block_on(L1BatchParamsProvider::new(&mut connection))
         .context("failed initializing L1 batch params provider")?;
@@ -51,7 +51,8 @@ pub fn create_vm(
     let storage_l2_block_number = first_l2_block_in_batch.number() - 1;
     let pg_storage =
         PostgresStorage::new(rt_handle.clone(), connection, storage_l2_block_number, true);
-    let storage_view = StorageView::new(pg_storage).to_rc_ptr();
+    let storage_overrides = StorageOverrides::new(pg_storage);
+    let storage_view = StorageView::new(storage_overrides).to_rc_ptr();
     let vm = VmInstance::new(l1_batch_env, system_env, storage_view.clone());
 
     Ok((vm, storage_view))

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

@@ -2,7 +2,10 @@
 use jsonrpsee::core::RpcResult;
 use jsonrpsee::proc_macros::rpc;
 use zksync_types::{
-    api::{BlockId, BlockIdVariant, BlockNumber, Transaction, TransactionVariant},
+    api::{
+        state_override::StateOverride, BlockId, BlockIdVariant, BlockNumber, Transaction,
+        TransactionVariant,
+    },
     transaction_request::CallRequest,
     Address, H256,
 };
@@ -34,7 +37,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>;