Move zksync stuff to its own module

crates/cast/bin/cmd/estimate.rs

@@ -2,8 +2,6 @@ use crate::tx::{CastTxBuilder, SenderKind};
 use alloy_primitives::U256;
 use alloy_provider::Provider;
 use alloy_rpc_types::BlockId;
-use alloy_zksync::network::transaction_request::TransactionRequest as ZkTransactionRequest;
-use cast::ZkTransactionOpts;
 use clap::Parser;
 use eyre::Result;
 use foundry_cli::{
@@ -14,6 +12,8 @@ use foundry_common::ens::NameOrAddress;
 use foundry_config::Config;
 use std::str::FromStr;
 
+mod zksync;
+
 /// CLI arguments for `cast estimate`.
 #[derive(Debug, Parser)]
 pub struct EstimateArgs {
@@ -44,7 +44,11 @@ pub struct EstimateArgs {
 
     /// Zksync Transaction
     #[command(flatten)]
-    zksync: ZkTransactionOpts,
+    zk_tx: zksync::ZkTransactionOpts,
+
+    /// Force a zksync eip-712 transaction and apply CREATE overrides
+    #[arg(long = "zksync")]
+    zk_force: bool,
 }
 
 #[derive(Debug, Parser)]
@@ -73,7 +77,9 @@ pub enum EstimateSubcommands {
 
 impl EstimateArgs {
     pub async fn run(self) -> Result<()> {
-        let Self { to, mut sig, mut args, mut tx, block, eth, command, zksync } = self;
+        let Self { to, mut sig, mut args, mut tx, block, eth, command, zk_tx, zk_force } = self;
+
+        let mut zk_code = Default::default();
 
         let config = Config::from(&eth);
         let provider = utils::get_provider(&config)?;
@@ -87,6 +93,12 @@ impl EstimateArgs {
             value,
         }) = command
         {
+            if zk_tx.has_zksync_args() || zk_force {
+                // NOTE(zk): `with_code_sig_and_args` decodes the code and appends it to the input
+                // we want the raw decoded constructor input from that function so we keep the code
+                // to encode the CONTRACT_CREATOR call later
+                zk_code = code.clone();
+            }
             sig = create_sig;
             args = create_args;
             if let Some(value) = value {
@@ -106,11 +118,8 @@ impl EstimateArgs {
             .build_raw(sender)
             .await?;
 
-        let gas = if zksync.has_zksync_args() {
-            let zk_provider = utils::get_provider_zksync(&config)?;
-            let mut zk_tx: ZkTransactionRequest = tx.inner.clone().into();
-            zksync.apply_to_tx(&mut zk_tx);
-            zk_provider.estimate_gas(&zk_tx).await?
+        let gas = if zk_tx.has_zksync_args() || zk_force {
+            zksync::estimate_gas(zk_tx, &config, tx, zk_code).await?
         } else {
             provider.estimate_gas(&tx).block(block.unwrap_or_default()).await?
         };

crates/cast/bin/cmd/estimate/zksync.rs

@@ -0,0 +1,92 @@
+use alloy_network::TransactionBuilder;
+use alloy_primitives::{hex, Address, Bytes, TxKind, U256};
+use alloy_provider::Provider;
+use alloy_rpc_types::TransactionRequest;
+use alloy_serde::WithOtherFields;
+use alloy_zksync::network::{
+    transaction_request::TransactionRequest as ZkTransactionRequest,
+    unsigned_tx::eip712::PaymasterParams,
+};
+use clap::{command, Parser};
+use eyre::Result;
+use foundry_cli::utils;
+use foundry_config::Config;
+
+#[derive(Clone, Debug, Parser)]
+#[command(next_help_heading = "Transaction options")]
+pub struct ZkTransactionOpts {
+    /// Paymaster address for the ZKSync transaction
+    #[arg(long = "zk-paymaster-address", requires = "paymaster_input")]
+    pub paymaster_address: Option<Address>,
+
+    /// Paymaster input for the ZKSync transaction
+    #[arg(long = "zk-paymaster-input", requires = "paymaster_address", value_parser = parse_hex_bytes)]
+    pub paymaster_input: Option<Bytes>,
+
+    /// Factory dependencies for the ZKSync transaction
+    #[arg(long = "zk-factory-deps", value_parser = parse_hex_bytes, value_delimiter = ',')]
+    pub factory_deps: Vec<Bytes>,
+
+    /// Custom signature for the ZKSync transaction
+    #[arg(long = "zk-custom-signature", value_parser = parse_hex_bytes)]
+    pub custom_signature: Option<Bytes>,
+
+    /// Gas per pubdata for the ZKSync transaction
+    #[arg(long = "zk-gas-per-pubdata")]
+    pub gas_per_pubdata: Option<U256>,
+}
+
+fn parse_hex_bytes(s: &str) -> Result<Bytes, String> {
+    hex::decode(s).map(Bytes::from).map_err(|e| format!("Invalid hex string: {e}"))
+}
+
+impl ZkTransactionOpts {
+    pub fn has_zksync_args(&self) -> bool {
+        self.paymaster_address.is_some() ||
+            !self.factory_deps.is_empty() ||
+            self.custom_signature.is_some() ||
+            self.gas_per_pubdata.is_some()
+    }
+}
+
+pub async fn estimate_gas(
+    zk_tx: ZkTransactionOpts,
+    config: &Config,
+    evm_tx: WithOtherFields<TransactionRequest>,
+    zk_code: String,
+) -> Result<u64> {
+    let zk_provider = utils::get_provider_zksync(config)?;
+    let is_create = evm_tx.to == Some(TxKind::Create);
+    let mut tx: ZkTransactionRequest = evm_tx.inner.clone().into();
+    if let Some(gas_per_pubdata) = zk_tx.gas_per_pubdata {
+        tx.set_gas_per_pubdata(gas_per_pubdata)
+    }
+
+    if let Some(custom_signature) = &zk_tx.custom_signature {
+        tx.set_custom_signature(custom_signature.clone());
+    }
+
+    if let (Some(paymaster), Some(paymaster_input)) =
+        (zk_tx.paymaster_address, zk_tx.paymaster_input.clone())
+    {
+        tx.set_paymaster_params(PaymasterParams { paymaster, paymaster_input });
+    }
+
+    if is_create {
+        let evm_input: Vec<u8> = tx.input().cloned().map(|bytes| bytes.into()).unwrap_or_default();
+        let zk_code_decoded = hex::decode(zk_code)?;
+        // constructor input gets appended to the bytecode
+        let zk_input = &evm_input[zk_code_decoded.len()..];
+        tx = tx.with_create_params(
+            zk_code_decoded,
+            zk_input.to_vec(),
+            zk_tx.factory_deps.into_iter().map(|v| v.into()).collect(),
+        )?;
+    } else {
+        tx.set_factory_deps(zk_tx.factory_deps.clone());
+    }
+
+    // TODO: Check if alloy calls this for estimate_gas. If so, we do not need this.
+    tx.prep_for_submission();
+    Ok(zk_provider.estimate_gas(&tx).await?)
+}

crates/cast/src/zksync.rs

@@ -1,14 +1,11 @@
 //! Contains zksync specific logic for foundry's `cast` functionality
 
 use alloy_network::AnyNetwork;
-use alloy_primitives::{Address, Bytes, U256};
 use alloy_provider::{PendingTransactionBuilder, Provider};
 use alloy_transport::Transport;
 use alloy_zksync::network::{
-    transaction_request::TransactionRequest as ZkTransactionRequest,
-    unsigned_tx::eip712::PaymasterParams, Zksync,
+    transaction_request::TransactionRequest as ZkTransactionRequest, Zksync,
 };
-use clap::Parser;
 use eyre::Result;
 
 use crate::Cast;
@@ -66,59 +63,3 @@ where
         Ok(res)
     }
 }
-
-#[derive(Clone, Debug, Parser)]
-#[command(next_help_heading = "Transaction options")]
-pub struct ZkTransactionOpts {
-    // /// Use ZKSync
-    // #[arg(long, default_value_ifs([("paymaster_address", ArgPredicate::IsPresent,
-    // "true"),("paymaster_input", ArgPredicate::IsPresent, "true")]))] pub zksync: bool,
-    /// Paymaster address for the ZKSync transaction
-    #[arg(long = "zk-paymaster-address", requires = "paymaster_input")]
-    pub paymaster_address: Option<Address>,
-
-    /// Paymaster input for the ZKSync transaction
-    #[arg(long = "zk-paymaster-input", requires = "paymaster_address")]
-    pub paymaster_input: Option<Bytes>,
-
-    /// Factory dependencies for the ZKSync transaction
-    #[arg(long = "zk-factory-deps")]
-    pub factory_deps: Vec<Bytes>,
-
-    /// Custom signature for the ZKSync transaction
-    #[arg(long = "zk-custom-signature")]
-    pub custom_signature: Option<Bytes>,
-
-    /// Gas per pubdata for the ZKSync transaction
-    #[arg(long = "zk-gas-per-pubdata")]
-    pub gas_per_pubdata: Option<U256>,
-}
-
-impl ZkTransactionOpts {
-    pub fn has_zksync_args(&self) -> bool {
-        self.paymaster_address.is_some() ||
-            !self.factory_deps.is_empty() ||
-            self.custom_signature.is_some() ||
-            self.gas_per_pubdata.is_some()
-    }
-
-    pub fn apply_to_tx(&self, tx: &mut ZkTransactionRequest) {
-        if let Some(gas_per_pubdata) = self.gas_per_pubdata {
-            tx.set_gas_per_pubdata(gas_per_pubdata)
-        }
-
-        if !self.factory_deps.is_empty() {
-            tx.set_factory_deps(self.factory_deps.clone());
-        }
-
-        if let Some(custom_signature) = &self.custom_signature {
-            tx.set_custom_signature(custom_signature.clone());
-        }
-
-        if let (Some(paymaster), Some(paymaster_input)) =
-            (self.paymaster_address, self.paymaster_input.clone())
-        {
-            tx.set_paymaster_params(PaymasterParams { paymaster, paymaster_input });
-        }
-    }
-}