feat: Adding (experimental) zkos support

Cargo.toml

@@ -28,12 +28,14 @@ categories = ["cryptography"]
 #########################
 zksync_multivm = { git = "https://github.com/matter-labs/zksync-era.git", rev = "core-v25.4.0" }
 zksync_contracts = { git = "https://github.com/matter-labs/zksync-era.git", rev = "core-v25.4.0" }
+zksync_basic_types = { git = "https://github.com/matter-labs/zksync-era.git", rev = "core-v25.4.0" }
 zksync_types = { git = "https://github.com/matter-labs/zksync-era.git", rev = "core-v25.4.0" }
 zksync_vm_interface = { git = "https://github.com/matter-labs/zksync-era.git", rev = "core-v25.4.0" }
 zksync_web3_decl = { git = "https://github.com/matter-labs/zksync-era.git", rev = "core-v25.4.0", features = [
     "server",
 ] }
 
+
 #########################
 # External dependencies #
 #########################

crates/cli/Cargo.toml

@@ -37,3 +37,9 @@ flate2.workspace = true
 
 [dev-dependencies]
 tempdir.workspace = true
+
+[features]
+
+zkos = ["anvil_zksync_core/zkos"]
+
+default = []

crates/config/Cargo.toml

@@ -25,3 +25,6 @@ rand.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 tracing.workspace = true
+
+[features]
+zkos = []

crates/config/src/config.rs

@@ -362,6 +362,9 @@ impl TestNodeConfig {
                 "Disabled".red()
             }
         );
+        #[cfg(feature = "zkos")]
+        tracing::info!("ZK OS:              {}", "Enabled".green());
+
         println!("\n");
         tracing::info!("========================================");
         for host in &self.host {

crates/core/Cargo.toml

@@ -17,6 +17,7 @@ zksync_multivm.workspace = true
 zksync_contracts.workspace = true
 zksync_types.workspace = true
 zksync_web3_decl.workspace = true
+zksync_basic_types.workspace = true
 
 anyhow.workspace = true
 tokio.workspace = true
@@ -42,6 +43,7 @@ flate2.workspace = true
 thiserror.workspace = true
 async-trait.workspace = true
 
+
 [dev-dependencies]
 maplit.workspace = true
 ethers.workspace = true
@@ -50,3 +52,8 @@ tempdir.workspace = true
 zksync-web3-rs.workspace = true
 test-case.workspace = true
 backon.workspace = true
+
+
+[features]
+zkos = []
+default = []

crates/core/src/lib.rs

@@ -1,3 +1,4 @@
+#![cfg_attr(feature = "zkos", feature(allocator_api))]
 //! anvil-zksync
 //!
 //! The `anvil-zksync` crate provides an in-memory node designed primarily for local testing.

crates/core/src/node/debug.rs

@@ -57,8 +57,7 @@ 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)
+        let mut l2_tx = L2Tx::from_request(request.into(), MAX_TX_SIZE, self.allow_no_target())
             .map_err(Web3Error::SerializationError)?;
         let execution_mode = zksync_multivm::interface::TxExecutionMode::EthCall;
 

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,15 @@ 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.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 +77,7 @@ 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.allow_no_target())?;
 
         l2_tx.set_input(tx_bytes.0, hash);
         if hash != l2_tx.hash() {
@@ -144,13 +137,7 @@ 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.allow_no_target())?;
 
         // `v` was overwritten with 0 during converting into l2 tx
         let mut signature = vec![0u8; 65];
@@ -183,10 +170,7 @@ 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),
-            &address,
-        );
+        let balance_key = StorageKeyLayout::get_storage_key_for_base_token(&address);
 
         let inner_guard = self.inner.read().await;
         match inner_guard.fork_storage.read_value_internal(&balance_key) {
@@ -281,7 +265,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(&address);
 
         match inner.fork_storage.read_value_internal(&nonce_key) {
             Ok(result) => Ok(h256_to_u64(result).into()),
@@ -619,7 +603,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

@@ -36,13 +36,17 @@ use std::sync::Arc;
 use tokio::sync::RwLock;
 use zksync_contracts::BaseSystemContracts;
 use zksync_multivm::interface::storage::{ReadStorage, StoragePtr};
+#[cfg(not(feature = "zkos"))]
+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};
+#[cfg(not(feature = "zkos"))]
+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;
@@ -306,6 +310,18 @@ impl InMemoryNode {
         Ok(())
     }
 
+    /// Whether it accepts the transactions that have 'null' as target.
+    /// This is used only when EVM emulator is enabled, or we're running in zkos mode.
+    pub fn allow_no_target(&self) -> bool {
+        #[cfg(feature = "zkos")]
+        return true;
+        #[cfg(not(feature = "zkos"))]
+        self.system_contracts
+            .contracts_for_l2_call()
+            .evm_emulator
+            .is_some()
+    }
+
     /// Applies multiple transactions across multiple blocks. All transactions are expected to be
     /// executable. Note that on error this method may leave node in partially applied state (i.e.
     /// some txs have been applied while others have not).
@@ -381,7 +397,17 @@ impl InMemoryNode {
         let system_env = inner.create_system_env(base_contracts, execution_mode);
 
         let storage = StorageView::new(&inner.fork_storage).into_rc_ptr();
+
+        #[cfg(not(feature = "zkos"))]
         let mut vm: Vm<_, HistoryDisabled> = Vm::new(batch_env, system_env, storage);
+        #[cfg(feature = "zkos")]
+        let mut vm: super::zkos::ZKOsVM<_, HistoryDisabled> = super::zkos::ZKOsVM::new(
+            batch_env,
+            system_env,
+            storage,
+            // TODO: this might be causing a deadlock.. check..
+            &inner.fork_storage.inner.read().unwrap().raw_storage,
+        );
 
         // We must inject *some* signature (otherwise bootloader code fails to generate hash).
         if l2_tx.common_data.signature.is_empty() {

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,7 @@ 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(&address);
         writer
             .fork_storage
             .set_value(balance_key, u256_to_h256(balance));
@@ -184,7 +181,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(&address);
         let enforced_full_nonce = nonces_to_full_nonce(nonce, nonce);
         tracing::info!(
             "👷 Nonces for address {:?} have been set to {}",
@@ -323,17 +320,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(())