feat: add state override for gas estimates (#1358)

## What ❔

- Adds state override for gas estimates

## Why ❔

- Solves #947 
- Feature parity with geth
ethereum/go-ethereum#27800

## Checklist

<!-- Check your PR fulfills the following items. -->
<!-- For draft PRs check the boxes as you complete them. -->

- [x] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [x] Tests for the changes have been added / updated.
- [x] Documentation comments have been added / updated.
- [x] Code has been formatted via `zk fmt` and `zk lint`.
- [x] Spellcheck has been run via `zk spellcheck`.
- [x] Linkcheck has been run via `zk linkcheck`.

---------

Co-authored-by: Juan Rigada <[email protected]>
Co-authored-by: Jrigada <[email protected]>
Co-authored-by: Danil <[email protected]>

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,
@@ -29,6 +30,7 @@ pub use self::{
     },
     shadow_storage::ShadowStorage,
     storage_factory::{BatchDiff, PgOrRocksdbStorage, ReadStorageFactory, RocksdbWithMemory},
+    storage_overrides::StorageOverrides,
     storage_view::{StorageView, StorageViewCache, StorageViewMetrics},
     witness::WitnessStorage,
 };
@@ -40,6 +42,7 @@ mod postgres;
 mod rocksdb;
 mod shadow_storage;
 mod storage_factory;
+mod storage_overrides;
 mod storage_view;
 #[cfg(test)]
 mod test_utils;
@@ -89,3 +92,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,150 @@
+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,
+    overridden_factory_deps: HashMap<H256, Vec<u8>>,
+    overridden_account_state: HashMap<AccountTreeId, HashMap<H256, H256>>,
+    overridden_account_state_diff: HashMap<AccountTreeId, HashMap<H256, H256>>,
+    overridden_balance: HashMap<StorageKey, U256>,
+    overridden_nonce: HashMap<StorageKey, U256>,
+    overridden_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,
+            overridden_factory_deps: HashMap::new(),
+            overridden_account_state: HashMap::new(),
+            overridden_account_state_diff: HashMap::new(),
+            overridden_balance: HashMap::new(),
+            overridden_nonce: HashMap::new(),
+            overridden_code: HashMap::new(),
+        }
+    }
+
+    /// Overrides a factory dependency code.
+    pub fn store_factory_dep(&mut self, hash: H256, code: Vec<u8>) {
+        self.overridden_factory_deps.insert(hash, code);
+    }
+
+    /// Overrides an account entire state.
+    pub fn override_account_state(&mut self, account: AccountTreeId, state: HashMap<H256, H256>) {
+        self.overridden_account_state.insert(account, state);
+    }
+
+    /// Overrides an account state diff.
+    pub fn override_account_state_diff(
+        &mut self,
+        account: AccountTreeId,
+        state_diff: HashMap<H256, H256>,
+    ) {
+        self.overridden_account_state_diff
+            .insert(account, state_diff);
+    }
+
+    /// 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.overridden_balance.get(key) {
+            return u256_to_h256(*balance);
+        }
+        if let Some(code) = self.overridden_code.get(key) {
+            return *code;
+        }
+
+        if let Some(nonce) = self.overridden_nonce.get(key) {
+            return u256_to_h256(*nonce);
+        }
+
+        if let Some(account_state) = self.overridden_account_state.get(key.account()) {
+            if let Some(value) = account_state.get(key.key()) {
+                return *value;
+            }
+            return H256::zero();
+        }
+
+        if let Some(account_state_diff) = self.overridden_account_state_diff.get(key.account()) {
+            if let Some(value) = account_state_diff.get(key.key()) {
+                return *value;
+            }
+        }
+
+        self.storage_handle.read_value(key)
+    }
+
+    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.overridden_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.overridden_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.overridden_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.overridden_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
+                            .overridden_account_state_diff
+                            .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)]
@@ -224,6 +224,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
+#[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,14 +8,14 @@ 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(
@@ -52,7 +52,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,
 };
@@ -31,10 +34,20 @@ pub trait EthNamespace {
     async fn chain_id(&self) -> RpcResult<U64>;
 
     #[method(name = "call")]
-    async fn call(&self, req: CallRequest, block: Option<BlockIdVariant>) -> RpcResult<Bytes>;
+    async fn call(
+        &self,
+        req: CallRequest,
+        block: Option<BlockIdVariant>,
+        state_override: Option<StateOverride>,
+    ) -> 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/web3_decl/src/namespaces/zks.rs

@@ -5,8 +5,8 @@ use jsonrpsee::core::RpcResult;
 use jsonrpsee::proc_macros::rpc;
 use zksync_types::{
     api::{
-        BlockDetails, BridgeAddresses, L1BatchDetails, L2ToL1LogProof, Proof, ProtocolVersion,
-        TransactionDetailedResult, TransactionDetails,
+        state_override::StateOverride, BlockDetails, BridgeAddresses, L1BatchDetails,
+        L2ToL1LogProof, Proof, ProtocolVersion, TransactionDetailedResult, TransactionDetails,
     },
     fee::Fee,
     fee_model::{FeeParams, PubdataIndependentBatchFeeModelInput},
@@ -29,10 +29,18 @@ use crate::{
 )]
 pub trait ZksNamespace {
     #[method(name = "estimateFee")]
-    async fn estimate_fee(&self, req: CallRequest) -> RpcResult<Fee>;
+    async fn estimate_fee(
+        &self,
+        req: CallRequest,
+        state_override: Option<StateOverride>,
+    ) -> RpcResult<Fee>;
 
     #[method(name = "estimateGasL1ToL2")]
-    async fn estimate_gas_l1_to_l2(&self, req: CallRequest) -> RpcResult<U256>;
+    async fn estimate_gas_l1_to_l2(
+        &self,
+        req: CallRequest,
+        state_override: Option<StateOverride>,
+    ) -> RpcResult<U256>;
 
     #[method(name = "getBridgehubContract")]
     async fn get_bridgehub_contract(&self) -> RpcResult<Option<Address>>;