[DO NOT MERGE] Zkos dev

Cargo.toml

@@ -28,6 +28,7 @@ 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 = [

crates/api_decl/src/lib.rs

@@ -2,6 +2,7 @@ mod namespaces;
 
 pub use namespaces::{
     AnvilNamespaceServer, ConfigNamespaceServer, EthTestNamespaceServer, EvmNamespaceServer,
+    ZKOSNamespaceServer,
 };
 
 // Re-export available namespaces from zksync-era

crates/api_decl/src/namespaces/mod.rs

@@ -2,8 +2,9 @@ mod anvil;
 mod config;
 mod eth_test;
 mod evm;
+mod zkos;
 
 pub use self::{
     anvil::AnvilNamespaceServer, config::ConfigNamespaceServer, eth_test::EthTestNamespaceServer,
-    evm::EvmNamespaceServer,
+    evm::EvmNamespaceServer, zkos::ZKOSNamespaceServer,
 };

crates/api_decl/src/namespaces/zkos.rs

@@ -0,0 +1,13 @@
+use jsonrpsee::core::RpcResult;
+use jsonrpsee::proc_macros::rpc;
+
+/// API bindings for the `zkos` experimental namespace.
+#[rpc(server, namespace = "zkos")]
+pub trait ZKOSNamespace {
+    /// Returns the witness for a given batch.
+    ///
+    /// # Returns
+    /// Bytes with the witness that can be passed to proving system.
+    #[method(name = "getWitness")]
+    async fn get_witness(&self, batch: u32) -> RpcResult<Option<String>>;
+}

crates/api_server/src/impls/mod.rs

@@ -6,10 +6,11 @@ mod eth_test;
 mod evm;
 mod net;
 mod web3;
+mod zkos;
 mod zks;
 
 pub use self::{
     anvil::AnvilNamespace, config::ConfigNamespace, debug::DebugNamespace, eth::EthNamespace,
     eth_test::EthTestNamespace, evm::EvmNamespace, net::NetNamespace, web3::Web3Namespace,
-    zks::ZksNamespace,
+    zkos::ZKOSNamespace, zks::ZksNamespace,
 };

crates/api_server/src/impls/zkos.rs

@@ -0,0 +1,19 @@
+use anvil_zksync_core::node::zkos_get_batch_witness;
+use jsonrpsee::core::{async_trait, RpcResult};
+
+use anvil_zksync_api_decl::ZKOSNamespaceServer;
+
+pub struct ZKOSNamespace {}
+
+impl ZKOSNamespace {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+#[async_trait]
+impl ZKOSNamespaceServer for ZKOSNamespace {
+    async fn get_witness(&self, batch: u32) -> RpcResult<Option<String>> {
+        Ok(zkos_get_batch_witness(&batch).map(hex::encode))
+    }
+}

crates/api_server/src/server.rs

@@ -1,11 +1,12 @@
+use crate::impls::ZKOSNamespace;
 use crate::{
     AnvilNamespace, ConfigNamespace, DebugNamespace, EthNamespace, EthTestNamespace, EvmNamespace,
     NetNamespace, Web3Namespace, ZksNamespace,
 };
 use anvil_zksync_api_decl::{
     AnvilNamespaceServer, ConfigNamespaceServer, DebugNamespaceServer, EthNamespaceServer,
     EthTestNamespaceServer, EvmNamespaceServer, NetNamespaceServer, Web3NamespaceServer,
-    ZksNamespaceServer,
+    ZKOSNamespaceServer, ZksNamespaceServer,
 };
 use anvil_zksync_core::node::InMemoryNode;
 use http::Method;
@@ -58,6 +59,7 @@ impl NodeServerBuilder {
         rpc.merge(ConfigNamespace::new(node.clone()).into_rpc())
             .unwrap();
         rpc.merge(ZksNamespace::new(node).into_rpc()).unwrap();
+        rpc.merge(ZKOSNamespace::new().into_rpc()).unwrap();
         rpc.merge(Web3Namespace.into_rpc()).unwrap();
         rpc
     }

crates/cli/src/cli.rs

@@ -4,7 +4,7 @@ use anvil_zksync_config::constants::{
     DEFAULT_DISK_CACHE_DIR, DEFAULT_MNEMONIC, TEST_NODE_NETWORK_ID,
 };
 use anvil_zksync_config::types::{
-    AccountGenerator, CacheConfig, CacheType, Genesis, SystemContractsOptions,
+    AccountGenerator, CacheConfig, CacheType, Genesis, SystemContractsOptions, ZKOSConfig,
 };
 use anvil_zksync_config::TestNodeConfig;
 use anvil_zksync_core::{
@@ -166,6 +166,10 @@ pub struct Cli {
     /// Enables EVM emulation. Requires local system contracts.
     pub emulate_evm: bool,
 
+    #[clap(flatten)]
+    /// ZKOS detailed config.
+    pub zkos_config: ZKOSConfig,
+
     // Logging Configuration
     #[arg(long, help_heading = "Logging Configuration")]
     /// Log level (default: info).
@@ -444,6 +448,7 @@ impl Cli {
             .set_config_out(self.config_out)
             .with_host(self.host)
             .with_evm_emulator(if self.emulate_evm { Some(true) } else { None })
+            .with_zkos_config(self.zkos_config)
             .with_health_check_endpoint(if self.health_check_endpoint {
                 Some(true)
             } else {

crates/cli/src/main.rs

@@ -221,9 +221,9 @@ async fn main() -> anyhow::Result<()> {
     let system_contracts = SystemContracts::from_options(
         &config.system_contracts_options,
         config.use_evm_emulator,
-        config.use_zkos,
+        config.zkos_config.clone(),
     );
-    let storage_key_layout = if config.use_zkos {
+    let storage_key_layout = if config.zkos_config.use_zkos {
         StorageKeyLayout::ZkOs
     } else {
         StorageKeyLayout::ZkEra

crates/config/src/config.rs

@@ -70,8 +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,
+    /// ZKOS configuration
+    pub zkos_config: ZKOSConfig,
     /// Optional chain ID for the node
     pub chain_id: Option<u32>,
     /// L1 gas price (optional override)
@@ -157,7 +157,7 @@ impl Default for TestNodeConfig {
             system_contracts_options: Default::default(),
             override_bytecodes_dir: None,
             use_evm_emulator: false,
-            use_zkos: false,
+            zkos_config: Default::default(),
             chain_id: None,
 
             // Gas configuration defaults
@@ -367,12 +367,22 @@ impl TestNodeConfig {
         );
         tracing::info!(
             "ZK OS:              {}",
-            if self.use_zkos {
+            if self.zkos_config.use_zkos {
                 "Enabled".green()
             } else {
                 "Disabled".red()
             }
         );
+        if self.zkos_config.use_zkos {
+            tracing::info!(
+                "ZK bin:             {}",
+                if let Some(path) = self.zkos_config.zkos_bin_path.as_ref() {
+                    path.green()
+                } else {
+                    "Not set".red()
+                }
+            );
+        }
 
         println!("\n");
         tracing::info!("========================================");
@@ -515,6 +525,13 @@ impl TestNodeConfig {
         self
     }
 
+    /// ZKOS configuration
+    #[must_use]
+    pub fn with_zkos_config(mut self, config: ZKOSConfig) -> Self {
+        self.zkos_config = config;
+        self
+    }
+
     /// Get the EVM emulation status
     pub fn is_evm_emulator_enabled(&self) -> bool {
         self.use_evm_emulator

crates/config/src/types/mod.rs

@@ -1,12 +1,14 @@
 mod account_generator;
 mod cache;
 mod genesis;
+mod zkos;
 
 pub use account_generator::AccountGenerator;
 pub use cache::{CacheConfig, CacheType};
 use clap::ValueEnum;
 pub use genesis::Genesis;
 use serde::Deserialize;
+pub use zkos::ZKOSConfig;
 
 #[derive(Deserialize, Default, Debug, Copy, Clone, PartialEq, ValueEnum)]
 pub enum SystemContractsOptions {

crates/config/src/types/zkos.rs

@@ -0,0 +1,14 @@
+use clap::Parser;
+use serde::Deserialize;
+
+/// Genesis
+#[derive(Deserialize, Clone, Debug, Parser, Default)]
+pub struct ZKOSConfig {
+    #[arg(long, help_heading = "Experimental Configuration")]
+    /// Enables zkos (experimental).
+    pub use_zkos: bool,
+
+    #[arg(long, help_heading = "Experimental Configuration")]
+    /// Path to zkos binary (if you need to compute witnesses).
+    pub zkos_bin_path: Option<String>,
+}

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,24 @@ flate2.workspace = true
 thiserror.workspace = true
 async-trait.workspace = true
 
+
+# ZK os dependencies
+
+forward_system = { git = "https://github.com/matter-labs/zk_ee.git", branch = "main"}
+basic_system = { git = "https://github.com/matter-labs/zk_ee.git", branch = "main"}
+zk_ee = { git = "https://github.com/matter-labs/zk_ee.git", branch = "main"}
+system_hooks = { git = "https://github.com/matter-labs/zk_ee.git", branch = "main"}
+zkos_api = { git = "https://github.com/matter-labs/zk_ee.git", branch = "main" }
+
+#forward_system = { path = "../../../zk_ee/forward_system" }
+#basic_system = { path = "../../../zk_ee/basic_system" }
+#zk_ee = { path = "../../../zk_ee/zk_ee" }
+#system_hooks = { path = "../../../zk_ee/system_hooks" }
+#zkos_api = { path = "../../../zk_ee/api" }
+
+
+ruint = { version = "1.12.3", default-features = false, features = ["alloc"]}
+
 [dev-dependencies]
 maplit.workspace = true
 ethers.workspace = true

crates/core/src/lib.rs

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

crates/core/src/node/in_memory.rs

@@ -392,13 +392,14 @@ impl InMemoryNode {
 
         let storage = StorageView::new(inner.read_storage()).into_rc_ptr();
 
-        let mut vm = if self.system_contracts.use_zkos {
+        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,
+                &self.system_contracts.zkos_config,
             ))
         } else {
             AnvilVM::ZKSync(Vm::new(batch_env, system_env, storage))
@@ -651,9 +652,9 @@ impl InMemoryNode {
         let system_contracts = SystemContracts::from_options(
             &config.system_contracts_options,
             config.use_evm_emulator,
-            config.use_zkos,
+            config.zkos_config.clone(),
         );
-        let storage_key_layout = if config.use_zkos {
+        let storage_key_layout = if config.zkos_config.use_zkos {
             StorageKeyLayout::ZkOs
         } else {
             StorageKeyLayout::ZkEra

crates/core/src/node/inner/in_memory_inner.rs

@@ -22,6 +22,7 @@ use crate::system_contracts::SystemContracts;
 use crate::utils::create_debug_output;
 use crate::{delegate_vm, formatter, utils};
 use anvil_zksync_config::constants::NON_FORK_FIRST_BLOCK_TIMESTAMP;
+use anvil_zksync_config::types::ZKOSConfig;
 use anvil_zksync_config::TestNodeConfig;
 use anvil_zksync_types::{ShowCalls, ShowGasDetails, ShowStorageLogs, ShowVMDetails};
 use anyhow::Context;
@@ -558,13 +559,14 @@ impl InMemoryNodeInner {
     ) -> Vec<TransactionResult> {
         let storage = StorageView::new(self.fork_storage.clone()).into_rc_ptr();
 
-        let mut vm = if self.system_contracts.use_zkos {
+        let mut vm = if self.system_contracts.use_zkos() {
             AnvilVM::ZKOs(ZKOsVM::<_, HistoryEnabled>::new(
                 batch_env.clone(),
                 system_env,
                 storage.clone(),
                 // TODO: this might be causing a deadlock.. check..
                 &self.fork_storage.inner.read().unwrap().raw_storage,
+                &self.system_contracts.zkos_config,
             ))
         } else {
             AnvilVM::ZKSync(Vm::new(batch_env.clone(), system_env, storage.clone()))
@@ -830,7 +832,7 @@ impl InMemoryNodeInner {
                 batch_env.clone(),
                 system_env.clone(),
                 &self.fork_storage,
-                self.system_contracts.use_zkos,
+                &self.system_contracts.zkos_config,
             );
 
             if result.statistics.pubdata_published > MAX_VM_PUBDATA_PER_BATCH.try_into().unwrap() {
@@ -868,7 +870,7 @@ impl InMemoryNodeInner {
                 batch_env.clone(),
                 system_env.clone(),
                 &self.fork_storage,
-                self.system_contracts.use_zkos,
+                &self.system_contracts.zkos_config,
             );
 
             if estimate_gas_result.result.is_failed() {
@@ -900,7 +902,7 @@ impl InMemoryNodeInner {
             batch_env,
             system_env,
             &self.fork_storage,
-            self.system_contracts.use_zkos,
+            &self.system_contracts.zkos_config,
         );
 
         let overhead = derive_overhead(
@@ -1016,7 +1018,7 @@ impl InMemoryNodeInner {
         batch_env: L1BatchEnv,
         system_env: SystemEnv,
         fork_storage: &ForkStorage,
-        is_zkos: bool,
+        zkos_config: &ZKOSConfig,
     ) -> VmExecutionResultAndLogs {
         let tx: Transaction = l2_tx.clone().into();
 
@@ -1057,13 +1059,14 @@ impl InMemoryNodeInner {
             .borrow_mut()
             .set_value(balance_key, u256_to_h256(current_balance));
 
-        let mut vm = if is_zkos {
+        let mut vm = if zkos_config.use_zkos {
             let mut vm = ZKOsVM::<_, HistoryDisabled>::new(
                 batch_env,
                 system_env,
                 storage,
                 // TODO: this might be causing a deadlock.. check..
                 &fork_storage.inner.read().unwrap().raw_storage,
+                zkos_config,
             );
             // Temporary hack - as we update the 'storage' just above, but zkos loads its full
             // state from fork_storage (that is not updated).
@@ -1398,9 +1401,9 @@ impl InMemoryNodeInner {
         let system_contracts = SystemContracts::from_options(
             &config.system_contracts_options,
             config.use_evm_emulator,
-            config.use_zkos,
+            config.zkos_config.clone(),
         );
-        let storage_key_layout = if config.use_zkos {
+        let storage_key_layout = if config.zkos_config.use_zkos {
             StorageKeyLayout::ZkOs
         } else {
             StorageKeyLayout::ZkEra

crates/core/src/node/mod.rs

@@ -26,3 +26,4 @@ pub use self::{
 pub use in_memory::*;
 pub use inner::{blockchain, fork, node_executor, time};
 pub use inner::{InMemoryNodeInner, TxExecutionOutput};
+pub use zkos::zkos_get_batch_witness;