Merge branch 'gprusak-general' of http://github.com/matter-labs/zksyn…

…c-era into gprusak-general

.dockerignore

@@ -39,7 +39,7 @@ contracts/.git
 !etc/multivm_bootloaders
 !cargo
 !bellman-cuda
-!prover/vk_setup_data_generator_server_fri/data/
+!prover/crates/bin/vk_setup_data_generator_server_fri/data/
 !.github/release-please/manifest.json
 
 !etc/env/file_based

.gitignore

@@ -108,7 +108,7 @@ hyperchain-*.yml
 /etc/hyperchains/artifacts
 
 # Prover keys that should not be commited
-prover/vk_setup_data_generator_server_fri/data/setup_*
+prover/crates/bin/vk_setup_data_generator_server_fri/data/setup_*
 
 # Zk Toolbox
 chains/era/configs/*

Cargo.toml

@@ -79,7 +79,7 @@ members = [
     "core/tests/vm-benchmark/harness",
 
     # Parts of prover workspace that are needed for Core workspace
-    "prover/prover_dal"
+    "prover/crates/lib/prover_dal"
 ]
 resolver = "2"
 
@@ -222,7 +222,7 @@ zksync_protobuf_build = "=0.1.0-rc.2"
 
 # "Local" dependencies
 zksync_multivm = { version = "0.1.0", path = "core/lib/multivm" }
-zksync_prover_dal = { version = "0.1.0", path = "prover/prover_dal" }
+zksync_prover_dal = { version = "0.1.0", path = "prover/crates/lib/prover_dal" }
 zksync_vlog = { version = "0.1.0", path = "core/lib/vlog" }
 zksync_vm_utils = { version = "0.1.0", path = "core/lib/vm_utils" }
 zksync_vm_benchmark_harness = { version = "0.1.0", path = "core/tests/vm-benchmark/harness" }

core/bin/zksync_server/src/main.rs

@@ -103,7 +103,7 @@ fn main() -> anyhow::Result<()> {
             let mut configs =
                 decode_yaml_repr::<zksync_protobuf_config::proto::general::GeneralConfig>(&yaml)
                     .context("failed decoding general YAML config")?;
-            // Tallback to the consensus_config.yaml file.
+            // Fallback to the consensus_config.yaml file.
             // TODO: remove once we move the consensus config to general config on stage
             if configs.consensus_config.is_none() {
                 configs.consensus_config =

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))