feat: [zkos] compute and return the witness (#545)

* ugly hacks to return the witness

* compute witness correctly

* passing zkconfig along

* moved zkos into zkos_config

* moved zkos namespace to proper place

* added docs

* put dependency back on main branch

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<Vec<u8>>>;
+}

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<Vec<u8>>> {
+        Ok(zkos_get_batch_witness(&batch))
+    }
+}

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,9 +166,9 @@ pub struct Cli {
     /// Enables EVM emulation. Requires local system contracts.
     pub emulate_evm: bool,
 
-    #[arg(long, help_heading = "System Configuration")]
-    /// Enables zkos (experimental).
-    pub use_zkos: bool,
+    #[clap(flatten)]
+    /// ZKOS detailed config.
+    pub zkos_config: ZKOSConfig,
 
     // Logging Configuration
     #[arg(long, help_heading = "Logging Configuration")]
@@ -448,7 +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(if self.use_zkos { 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,7 +221,7 @@ 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 (node_inner, _fork_storage, blockchain, time) = InMemoryNodeInner::init(

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,12 +525,10 @@ impl TestNodeConfig {
         self
     }
 
-    /// Enable or disable zkos
+    /// ZKOS configuration
     #[must_use]
-    pub fn with_zkos(mut self, enable: Option<bool>) -> Self {
-        if let Some(enable) = enable {
-            self.use_zkos = enable;
-        }
+    pub fn with_zkos_config(mut self, config: ZKOSConfig) -> Self {
+        self.zkos_config = config;
         self
     }
 

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

@@ -50,11 +50,14 @@ forward_system = { git = "https://github.com/matter-labs/zk_ee.git", branch = "m
 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"]}
 

crates/core/src/node/eth.rs

@@ -177,7 +177,7 @@ impl InMemoryNode {
         _block: Option<BlockIdVariant>,
     ) -> anyhow::Result<U256> {
         let balance_key = StorageKeyLayout::get_storage_key_for_base_token(
-            self.system_contracts.use_zkos,
+            self.system_contracts.use_zkos(),
             &address,
         );
 
@@ -274,7 +274,7 @@ impl InMemoryNode {
         _block: Option<BlockIdVariant>,
     ) -> anyhow::Result<U256> {
         let inner = self.inner.read().await;
-        let nonce_key = StorageKeyLayout::get_nonce_key(self.system_contracts.use_zkos, &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()),

crates/core/src/node/in_memory.rs

@@ -385,13 +385,14 @@ impl InMemoryNode {
 
         let storage = StorageView::new(&inner.fork_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))
@@ -652,7 +653,7 @@ impl InMemoryNode {
         let system_contracts = SystemContracts::from_options(
             &config.system_contracts_options,
             config.use_evm_emulator,
-            config.use_zkos,
+            config.zkos_config.clone(),
         );
         let (inner, _, blockchain, time) = InMemoryNodeInner::init(
             fork,

crates/core/src/node/in_memory_ext.rs

@@ -168,7 +168,7 @@ impl InMemoryNode {
     pub async fn set_balance(&self, address: Address, balance: U256) -> bool {
         let writer = self.inner.write().await;
         let balance_key = StorageKeyLayout::get_storage_key_for_base_token(
-            self.system_contracts.use_zkos,
+            self.system_contracts.use_zkos(),
             &address,
         );
         writer
@@ -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 = StorageKeyLayout::get_nonce_key(self.system_contracts.use_zkos, &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 {}",