feat: Adding (experimental) zkos support (#534)

* first version - dirty

* moved nonce key and base token key to separate file

* moved allow_no_target to separate method

* cleanup inconsistent keys, align arguments

* compiles but dirty

* tracers are fixed now

* cleanup warnings

* cleanup warnings and factory deps

* cleanup prints & logs

* small cleanups

* cleaning up rust version

* cleanups after merge

* move deps to optional

* fixed toml typo

* marked zkos implemetnation as 'todo' until we finish opensourcing

* clippy

* revert basic types

* got rid of feature flag.

crates/cli/src/main.rs

@@ -218,8 +218,11 @@ async fn main() -> anyhow::Result<()> {
 
     let fee_input_provider = TestNodeFeeInputProvider::from_fork(fork_details.as_ref());
     let filters = Arc::new(RwLock::new(EthFilters::default()));
-    let system_contracts =
-        SystemContracts::from_options(&config.system_contracts_options, config.use_evm_emulator);
+    let system_contracts = SystemContracts::from_options(
+        &config.system_contracts_options,
+        config.use_evm_emulator,
+        config.use_zkos,
+    );
 
     let (node_inner, _fork_storage, blockchain, time) = InMemoryNodeInner::init(
         fork_details,

crates/config/src/config.rs

@@ -70,6 +70,8 @@ pub struct TestNodeConfig {
     pub override_bytecodes_dir: Option<String>,
     /// Enables EVM emulation mode
     pub use_evm_emulator: bool,
+    /// Enables ZKOS mode (experimental)
+    pub use_zkos: bool,
     /// Optional chain ID for the node
     pub chain_id: Option<u32>,
     /// L1 gas price (optional override)
@@ -155,6 +157,7 @@ impl Default for TestNodeConfig {
             system_contracts_options: Default::default(),
             override_bytecodes_dir: None,
             use_evm_emulator: false,
+            use_zkos: false,
             chain_id: None,
 
             // Gas configuration defaults
@@ -362,6 +365,15 @@ impl TestNodeConfig {
                 "Disabled".red()
             }
         );
+        tracing::info!(
+            "ZK OS:              {}",
+            if self.use_zkos {
+                "Enabled".green()
+            } else {
+                "Disabled".red()
+            }
+        );
+
         println!("\n");
         tracing::info!("========================================");
         for host in &self.host {

crates/core/Cargo.toml

@@ -49,4 +49,4 @@ httptest.workspace = true
 tempdir.workspace = true
 zksync-web3-rs.workspace = true
 test-case.workspace = true
-backon.workspace = true
+backon.workspace = true

crates/core/src/node/debug.rs

@@ -57,9 +57,12 @@ impl InMemoryNode {
             )));
         }
 
-        let allow_no_target = system_contracts.evm_emulator.is_some();
-        let mut l2_tx = L2Tx::from_request(request.into(), MAX_TX_SIZE, allow_no_target)
-            .map_err(Web3Error::SerializationError)?;
+        let mut l2_tx = L2Tx::from_request(
+            request.into(),
+            MAX_TX_SIZE,
+            self.system_contracts.allow_no_target(),
+        )
+        .map_err(Web3Error::SerializationError)?;
         let execution_mode = zksync_multivm::interface::TxExecutionMode::EthCall;
 
         // init vm

crates/core/src/node/eth.rs

@@ -2,21 +2,20 @@ use std::collections::HashSet;
 
 use anyhow::Context as _;
 use colored::Colorize;
-use zksync_multivm::interface::{ExecutionResult, TxExecutionMode};
+use zksync_multivm::interface::ExecutionResult;
 use zksync_multivm::vm_latest::constants::ETH_CALL_GAS_LIMIT;
 use zksync_types::h256_to_u256;
 use zksync_types::{
     api,
     api::{Block, BlockIdVariant, BlockNumber, TransactionVariant},
-    get_code_key, get_nonce_key,
+    get_code_key,
     l2::L2Tx,
     transaction_request::TransactionRequest,
-    utils::storage_key_for_standard_token_balance,
-    PackedEthSignature, L2_BASE_TOKEN_ADDRESS, MAX_L1_TRANSACTION_GAS_LIMIT,
+    PackedEthSignature, MAX_L1_TRANSACTION_GAS_LIMIT,
 };
 use zksync_types::{
     web3::{self, Bytes},
-    AccountTreeId, Address, H160, H256, U256, U64,
+    Address, H160, H256, U256, U64,
 };
 use zksync_web3_decl::{
     error::Web3Error,
@@ -29,15 +28,19 @@ use crate::{
     utils::{h256_to_u64, TransparentError},
 };
 
+use super::keys::StorageKeyLayout;
+
 impl InMemoryNode {
     pub async fn call_impl(
         &self,
         req: zksync_types::transaction_request::CallRequest,
     ) -> Result<Bytes, Web3Error> {
         let system_contracts = self.system_contracts.contracts_for_l2_call().clone();
-        let allow_no_target = system_contracts.evm_emulator.is_some();
-
-        let mut tx = L2Tx::from_request(req.into(), MAX_TX_SIZE, allow_no_target)?;
+        let mut tx = L2Tx::from_request(
+            req.into(),
+            MAX_TX_SIZE,
+            self.system_contracts.allow_no_target(),
+        )?;
         tx.common_data.fee.gas_limit = ETH_CALL_GAS_LIMIT.into();
         let call_result = self
             .run_l2_call(tx, system_contracts)
@@ -78,13 +81,8 @@ impl InMemoryNode {
         let chain_id = self.inner.read().await.fork_storage.chain_id;
 
         let (tx_req, hash) = TransactionRequest::from_bytes(&tx_bytes.0, chain_id)?;
-        // Impersonation does not matter in this context so we assume the tx is not impersonated:
-        // system contracts here are fetched solely to check for EVM emulator.
-        let system_contracts = self
-            .system_contracts
-            .contracts(TxExecutionMode::VerifyExecute, false);
-        let allow_no_target = system_contracts.evm_emulator.is_some();
-        let mut l2_tx = L2Tx::from_request(tx_req, MAX_TX_SIZE, allow_no_target)?;
+        let mut l2_tx =
+            L2Tx::from_request(tx_req, MAX_TX_SIZE, self.system_contracts.allow_no_target())?;
 
         l2_tx.set_input(tx_bytes.0, hash);
         if hash != l2_tx.hash() {
@@ -144,13 +142,8 @@ impl InMemoryNode {
             27,
         ))?;
 
-        // Impersonation does not matter in this context so we assume the tx is not impersonated:
-        // system contracts here are fetched solely to check for EVM emulator.
-        let system_contracts = self
-            .system_contracts
-            .contracts(TxExecutionMode::VerifyExecute, false);
-        let allow_no_target = system_contracts.evm_emulator.is_some();
-        let mut l2_tx: L2Tx = L2Tx::from_request(tx_req, MAX_TX_SIZE, allow_no_target)?;
+        let mut l2_tx: L2Tx =
+            L2Tx::from_request(tx_req, MAX_TX_SIZE, self.system_contracts.allow_no_target())?;
 
         // `v` was overwritten with 0 during converting into l2 tx
         let mut signature = vec![0u8; 65];
@@ -183,8 +176,8 @@ impl InMemoryNode {
         // TODO: Support
         _block: Option<BlockIdVariant>,
     ) -> anyhow::Result<U256> {
-        let balance_key = storage_key_for_standard_token_balance(
-            AccountTreeId::new(L2_BASE_TOKEN_ADDRESS),
+        let balance_key = StorageKeyLayout::get_storage_key_for_base_token(
+            self.system_contracts.use_zkos,
             &address,
         );
 
@@ -281,7 +274,7 @@ impl InMemoryNode {
         _block: Option<BlockIdVariant>,
     ) -> anyhow::Result<U256> {
         let inner = self.inner.read().await;
-        let nonce_key = get_nonce_key(&address);
+        let nonce_key = StorageKeyLayout::get_nonce_key(self.system_contracts.use_zkos, &address);
 
         match inner.fork_storage.read_value_internal(&nonce_key) {
             Ok(result) => Ok(h256_to_u64(result).into()),
@@ -619,7 +612,7 @@ mod tests {
         utils::deployed_address_create,
         Bloom, K256PrivateKey, L2BlockNumber, StorageKey, EMPTY_UNCLES_HASH,
     };
-    use zksync_types::{u256_to_h256, web3, Nonce};
+    use zksync_types::{u256_to_h256, web3, AccountTreeId, Nonce};
     use zksync_web3_decl::types::{SyncState, ValueOrArray};
 
     async fn test_node(url: &str) -> InMemoryNode {

crates/core/src/node/in_memory.rs

@@ -2,9 +2,9 @@
 use super::inner::fork::ForkDetails;
 use super::inner::node_executor::NodeExecutorHandle;
 use super::inner::InMemoryNodeInner;
+use super::vm::AnvilVM;
 use crate::deps::{storage_view::StorageView, InMemoryStorage};
 use crate::filters::EthFilters;
-use crate::formatter;
 use crate::node::call_error_tracer::CallErrorTracer;
 use crate::node::error::LoadStateError;
 use crate::node::fee_model::TestNodeFeeInputProvider;
@@ -16,6 +16,7 @@ use crate::node::state::VersionedState;
 use crate::node::{BlockSealer, BlockSealerMode, NodeExecutor, TxPool};
 use crate::observability::Observability;
 use crate::system_contracts::SystemContracts;
+use crate::{delegate_vm, formatter};
 use anvil_zksync_config::constants::{
     LEGACY_RICH_WALLETS, NON_FORK_FIRST_BLOCK_TIMESTAMP, RICH_WALLETS, TEST_NODE_NETWORK_ID,
 };
@@ -36,13 +37,15 @@ use std::sync::Arc;
 use tokio::sync::RwLock;
 use zksync_contracts::BaseSystemContracts;
 use zksync_multivm::interface::storage::{ReadStorage, StoragePtr};
+use zksync_multivm::interface::VmFactory;
 use zksync_multivm::interface::{
-    ExecutionResult, InspectExecutionMode, L1BatchEnv, L2BlockEnv, TxExecutionMode, VmFactory,
-    VmInterface,
+    ExecutionResult, InspectExecutionMode, L1BatchEnv, L2BlockEnv, TxExecutionMode, VmInterface,
 };
 use zksync_multivm::tracers::CallTracer;
 use zksync_multivm::utils::{get_batch_base_fee, get_max_batch_gas_limit};
-use zksync_multivm::vm_latest::{HistoryDisabled, ToTracerPointer, Vm};
+use zksync_multivm::vm_latest::Vm;
+
+use zksync_multivm::vm_latest::{HistoryDisabled, ToTracerPointer};
 use zksync_multivm::VmVersion;
 use zksync_types::api::{Block, DebugCall, TransactionReceipt, TransactionVariant};
 use zksync_types::block::unpack_block_info;
@@ -381,23 +384,37 @@ impl InMemoryNode {
         let system_env = inner.create_system_env(base_contracts, execution_mode);
 
         let storage = StorageView::new(&inner.fork_storage).into_rc_ptr();
-        let mut vm: Vm<_, HistoryDisabled> = Vm::new(batch_env, system_env, storage);
+
+        let mut vm = if self.system_contracts.use_zkos {
+            AnvilVM::ZKOs(super::zkos::ZKOsVM::<_, HistoryDisabled>::new(
+                batch_env,
+                system_env,
+                storage,
+                // TODO: this might be causing a deadlock.. check..
+                &inner.fork_storage.inner.read().unwrap().raw_storage,
+            ))
+        } else {
+            AnvilVM::ZKSync(Vm::new(batch_env, system_env, storage))
+        };
 
         // We must inject *some* signature (otherwise bootloader code fails to generate hash).
         if l2_tx.common_data.signature.is_empty() {
             l2_tx.common_data.signature = PackedEthSignature::default().serialize_packed().into();
         }
 
         let tx: Transaction = l2_tx.into();
-        vm.push_transaction(tx.clone());
+        delegate_vm!(vm, push_transaction(tx.clone()));
 
         let call_tracer_result = Arc::new(OnceCell::default());
 
         let tracers = vec![
             CallErrorTracer::new().into_tracer_pointer(),
             CallTracer::new(call_tracer_result.clone()).into_tracer_pointer(),
         ];
-        let tx_result = vm.inspect(&mut tracers.into(), InspectExecutionMode::OneTx);
+        let tx_result = delegate_vm!(
+            vm,
+            inspect(&mut tracers.into(), InspectExecutionMode::OneTx)
+        );
 
         let call_traces = Arc::try_unwrap(call_tracer_result)
             .unwrap()
@@ -635,6 +652,7 @@ impl InMemoryNode {
         let system_contracts = SystemContracts::from_options(
             &config.system_contracts_options,
             config.use_evm_emulator,
+            config.use_zkos,
         );
         let (inner, _, blockchain, time) = InMemoryNodeInner::init(
             fork,

crates/core/src/node/in_memory_ext.rs

@@ -2,17 +2,14 @@ use super::inner::fork::ForkDetails;
 use super::pool::TxBatch;
 use super::sealer::BlockSealerMode;
 use super::InMemoryNode;
+use crate::node::keys::StorageKeyLayout;
 use crate::utils::bytecode_to_factory_dep;
 use anvil_zksync_types::api::{DetailedTransaction, ResetRequest};
 use anyhow::anyhow;
 use std::time::Duration;
 use zksync_types::api::{Block, TransactionVariant};
 use zksync_types::u256_to_h256;
-use zksync_types::{
-    get_code_key, get_nonce_key,
-    utils::{nonces_to_full_nonce, storage_key_for_eth_balance},
-    L2BlockNumber, StorageKey,
-};
+use zksync_types::{get_code_key, utils::nonces_to_full_nonce, L2BlockNumber, StorageKey};
 use zksync_types::{AccountTreeId, Address, H256, U256, U64};
 
 type Result<T> = anyhow::Result<T>;
@@ -170,7 +167,10 @@ impl InMemoryNode {
 
     pub async fn set_balance(&self, address: Address, balance: U256) -> bool {
         let writer = self.inner.write().await;
-        let balance_key = storage_key_for_eth_balance(&address);
+        let balance_key = StorageKeyLayout::get_storage_key_for_base_token(
+            self.system_contracts.use_zkos,
+            &address,
+        );
         writer
             .fork_storage
             .set_value(balance_key, u256_to_h256(balance));
@@ -184,7 +184,7 @@ impl InMemoryNode {
 
     pub async fn set_nonce(&self, address: Address, nonce: U256) -> bool {
         let writer = self.inner.write().await;
-        let nonce_key = get_nonce_key(&address);
+        let nonce_key = StorageKeyLayout::get_nonce_key(self.system_contracts.use_zkos, &address);
         let enforced_full_nonce = nonces_to_full_nonce(nonce, nonce);
         tracing::info!(
             "👷 Nonces for address {:?} have been set to {}",
@@ -323,17 +323,7 @@ impl InMemoryNode {
             .strip_prefix("0x")
             .ok_or_else(|| anyhow!("code must be 0x-prefixed"))?;
         let code_bytes = hex::decode(code_slice)?;
-        let hashcode = bytecode_to_factory_dep(code_bytes)?;
-        let hash = u256_to_h256(hashcode.0);
-        let code = hashcode
-            .1
-            .iter()
-            .flat_map(|entry| {
-                let mut bytes = vec![0u8; 32];
-                entry.to_big_endian(&mut bytes);
-                bytes.to_vec()
-            })
-            .collect();
+        let (hash, code) = bytecode_to_factory_dep(code_bytes)?;
         writer.fork_storage.store_factory_dep(hash, code);
         writer.fork_storage.set_value(code_key, hash);
         Ok(())