Add zk gas per pubdata flag

crates/cli/src/opts/build/zksync.rs

@@ -133,6 +133,10 @@ pub struct ZkSyncArgs {
     )]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub suppressed_errors: Option<Vec<ZkSolcError>>,
+
+    /// Gas per pubdata
+    #[clap(long = "zk-gas-per-pubdata", value_name = "GAS_PER_PUBDATA")]
+    pub gas_per_pubdata: Option<u64>,
 }
 
 impl ZkSyncArgs {
@@ -172,6 +176,7 @@ impl ZkSyncArgs {
         let suppressed_errors =
             self.suppressed_errors.clone().map(|values| values.into_iter().collect::<HashSet<_>>());
         set_if_some!(suppressed_errors, zksync.suppressed_errors);
+        set_if_some!(self.gas_per_pubdata, zksync.gas_per_pubdata);
 
         zksync
     }

crates/config/src/zksync.rs

@@ -74,6 +74,9 @@ pub struct ZkSyncConfig {
 
     // zksolc suppressed errors.
     pub suppressed_errors: HashSet<ZkSolcError>,
+
+    // Gas per pubdata
+    pub gas_per_pubdata: Option<u64>,
 }
 
 impl Default for ZkSyncConfig {
@@ -93,6 +96,7 @@ impl Default for ZkSyncConfig {
             optimizer_details: Default::default(),
             suppressed_errors: Default::default(),
             suppressed_warnings: Default::default(),
+            gas_per_pubdata: Default::default(),
         }
     }
 }

crates/evm/evm/src/executors/strategy.rs

@@ -110,6 +110,13 @@ pub trait ExecutorStrategyExt {
     ) {
     }
 
+    fn zksync_set_gas_per_pubdata(
+        &self,
+        _ctx: &mut dyn ExecutorStrategyContext,
+        _gas_per_pubdata: u64,
+    ) {
+    }
+
     /// Set the fork environment on the context.
     fn zksync_set_fork_env(
         &self,

crates/forge/tests/it/zk/gas.rs

@@ -3,16 +3,14 @@ use foundry_test_utils::{forgetest_async, util, TestProject};
 use foundry_test_utils::{util::OutputExt, ZkSyncNode};
 
 forgetest_async!(zk_script_execution_with_gas_price_specified_by_user, |prj, cmd| {
+    // Setup
     setup_gas_prj(&mut prj);
-
     let node = ZkSyncNode::start().await;
     let url = node.url();
-
     cmd.forge_fuse();
+    let private_key = get_rich_wallet_key();
 
-    let private_key =
-        ZkSyncNode::rich_wallets().next().map(|(_, pk, _)| pk).expect("No rich wallets available");
-
+    // Create script args with gas price parameters
     let script_args = vec![
         "--zk-startup",
         "./script/Gas.s.sol",
@@ -31,11 +29,12 @@ forgetest_async!(zk_script_execution_with_gas_price_specified_by_user, |prj, cmd
         "123123",
     ];
 
+    // Execute script and verify success
     cmd.arg("script").args(&script_args);
-
     let stdout = cmd.assert_success().get_output().stdout_lossy();
     assert!(stdout.contains("ONCHAIN EXECUTION COMPLETE & SUCCESSFUL"));
 
+    // Verify transaction details from broadcast artifacts
     let run_latest = foundry_common::fs::json_files(prj.root().join("broadcast").as_path())
         .find(|file| file.ends_with("run-latest.json"))
         .expect("No broadcast artifacts");
@@ -45,6 +44,7 @@ forgetest_async!(zk_script_execution_with_gas_price_specified_by_user, |prj, cmd
 
     assert_eq!(json["transactions"].as_array().expect("broadcastable txs").len(), 1);
 
+    // Verify gas prices in transaction
     let transaction_hash = json["receipts"][0]["transactionHash"].as_str().unwrap();
     let stdout = cmd
         .cast_fuse()
@@ -60,6 +60,86 @@ forgetest_async!(zk_script_execution_with_gas_price_specified_by_user, |prj, cmd
     assert!(stdout.contains("maxPriorityFeePerGas 123123"));
 });
 
+forgetest_async!(zk_script_execution_with_gas_multiplier, |prj, cmd| {
+    // Setup
+    setup_gas_prj(&mut prj);
+    let node = ZkSyncNode::start().await;
+    let url = node.url();
+    cmd.forge_fuse();
+    let private_key = get_rich_wallet_key();
+
+    // Test with insufficient gas multiplier (should fail)
+    let insufficient_multiplier_args =
+        create_script_args(&private_key, &url, "--gas-estimate-multiplier", "1");
+    cmd.arg("script").args(&insufficient_multiplier_args);
+    cmd.assert_failure();
+    cmd.forge_fuse();
+
+    // Test with sufficient gas multiplier (should succeed)
+    let sufficient_multiplier_args =
+        create_script_args(&private_key, &url, "--gas-estimate-multiplier", "100");
+    cmd.arg("script").args(&sufficient_multiplier_args);
+    let stdout = cmd.assert_success().get_output().stdout_lossy();
+    assert!(stdout.contains("ONCHAIN EXECUTION COMPLETE & SUCCESSFUL"));
+});
+
+forgetest_async!(zk_script_execution_with_gas_per_pubdata, |prj, cmd| {
+    // Setup
+    setup_gas_prj(&mut prj);
+    let node = ZkSyncNode::start().await;
+    let url = node.url();
+    cmd.forge_fuse();
+    let private_key = get_rich_wallet_key();
+
+    // Test with unacceptable gas per pubdata (should fail)
+    let zero_pubdata_args = create_script_args(&private_key, &url, "--zk-gas-per-pubdata", "1");
+    cmd.arg("script").args(&zero_pubdata_args);
+    cmd.assert_failure();
+    cmd.forge_fuse();
+
+    // Test with sufficient gas per pubdata (should succeed)
+    let sufficient_pubdata_args =
+        create_script_args(&private_key, &url, "--zk-gas-per-pubdata", "3000");
+    cmd.arg("script").args(&sufficient_pubdata_args);
+    let stdout = cmd.assert_success().get_output().stdout_lossy();
+    assert!(stdout.contains("ONCHAIN EXECUTION COMPLETE & SUCCESSFUL"));
+});
+
+// Helper function to get rich wallet private key
+fn get_rich_wallet_key() -> String {
+    ZkSyncNode::rich_wallets()
+        .next()
+        .map(|(_, pk, _)| pk)
+        .expect("No rich wallets available")
+        .to_owned()
+}
+
+// Helper function to create script arguments
+fn create_script_args<'a>(
+    private_key: &'a str,
+    url: &'a str,
+    gas_param: &'a str,
+    gas_value: &'a str,
+) -> Vec<&'a str> {
+    vec![
+        "--zk-startup",
+        "./script/Gas.s.sol",
+        "--private-key",
+        private_key,
+        "--chain",
+        "260",
+        "--rpc-url",
+        url,
+        "--slow",
+        "-vvvvv",
+        "--broadcast",
+        "--timeout",
+        "3",
+        gas_param,
+        gas_value,
+    ]
+}
+
 fn setup_gas_prj(prj: &mut TestProject) {
     util::initialize(prj.root());
     prj.add_script("Gas.s.sol", include_str!("../../fixtures/zk/Gas.s.sol")).unwrap();

crates/script/src/broadcast.rs

@@ -141,6 +141,11 @@ pub async fn send_transaction(
                         },
                     );
                 }
+
+                if let Some(gas_per_pubdata) = zk_tx_meta.gas_per_pubdata {
+                    zk_tx.set_gas_per_pubdata(alloy_primitives::Uint::from(gas_per_pubdata));
+                }
+
                 foundry_zksync_core::estimate_fee(&mut zk_tx, &zk_provider, estimate_multiplier)
                     .await?;
 

crates/script/src/lib.rs

@@ -638,6 +638,11 @@ impl ScriptConfig {
                 strategy.context.as_mut(),
                 dual_compiled_contracts,
             );
+            if let Some(gas_per_pubdata) = self.config.zksync.gas_per_pubdata {
+                strategy
+                    .runner
+                    .zksync_set_gas_per_pubdata(strategy.context.as_mut(), gas_per_pubdata);
+            }
 
             if let Some(fork_url) = &self.evm_opts.fork_url {
                 strategy.runner.zksync_set_fork_env(strategy.context.as_mut(), fork_url, &env)?;

crates/strategy/zksync/src/cheatcode/runner/mod.rs

@@ -126,6 +126,8 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
             paymaster_input: paymaster_data.input.to_vec(),
         });
 
+        let gas_per_pubdata = ctx.zk_env.gas_per_pubdata;
+
         let rpc = ecx_inner.db.active_fork_url();
 
         let injected_factory_deps = ctx
@@ -165,6 +167,7 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
                 serde_json::to_value(ZkTransactionMetadata::new(
                     factory_deps,
                     paymaster_params.clone(),
+                    gas_per_pubdata,
                 ))
                 .expect("failed encoding json"),
             );
@@ -189,8 +192,12 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
         });
         tx.other.insert(
             ZKSYNC_TRANSACTION_OTHER_FIELDS_KEY.to_string(),
-            serde_json::to_value(ZkTransactionMetadata::new(zk_tx_factory_deps, paymaster_params))
-                .expect("failed encoding json"),
+            serde_json::to_value(ZkTransactionMetadata::new(
+                zk_tx_factory_deps,
+                paymaster_params,
+                gas_per_pubdata,
+            ))
+            .expect("failed encoding json"),
         );
         broadcastable_transactions.push_back(BroadcastableTransaction {
             rpc,
@@ -256,15 +263,15 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
             paymaster: paymaster_data.address.to_h160(),
             paymaster_input: paymaster_data.input.to_vec(),
         });
+        let gas_per_pubdata = ctx.zk_env.gas_per_pubdata;
         let factory_deps = if call.target_address == DEFAULT_CREATE2_DEPLOYER_ZKSYNC {
             // We shouldn't need factory_deps for CALLs
             factory_deps.clone()
         } else {
             // For this case we use only the injected factory deps
             injected_factory_deps
         };
-        let zk_tx = ZkTransactionMetadata::new(factory_deps, paymaster_params);
-
+        let zk_tx = ZkTransactionMetadata::new(factory_deps, paymaster_params, gas_per_pubdata);
         let mut tx_req = TransactionRequest {
             from: Some(broadcast.new_origin),
             to: Some(TxKind::from(Some(call.target_address))),

crates/strategy/zksync/src/executor/runner.rs

@@ -154,6 +154,15 @@ impl ExecutorStrategyExt for ZksyncExecutorStrategyRunner {
         ctx.dual_compiled_contracts = dual_compiled_contracts;
     }
 
+    fn zksync_set_gas_per_pubdata(
+        &self,
+        ctx: &mut dyn ExecutorStrategyContext,
+        gas_per_pubdata: u64,
+    ) {
+        let ctx = get_context(ctx);
+        ctx.zk_env.gas_per_pubdata = Some(gas_per_pubdata);
+    }
+
     fn zksync_set_fork_env(
         &self,
         ctx: &mut dyn ExecutorStrategyContext,
@@ -173,6 +182,7 @@ impl ExecutorStrategyExt for ZksyncExecutorStrategyRunner {
 
         if let Some(block_details) = maybe_block_details {
             ctx.zk_env = ZkEnv {
+                gas_per_pubdata: ctx.zk_env.gas_per_pubdata,
                 l1_gas_price: block_details
                     .l1_gas_price
                     .try_into()

crates/zksync/core/src/lib.rs

@@ -113,12 +113,18 @@ pub struct ZkTransactionMetadata {
     pub factory_deps: Vec<Vec<u8>>,
     /// Paymaster data for ZK transactions.
     pub paymaster_data: Option<PaymasterParams>,
+    /// Gas per pubdata for ZK transactions.
+    pub gas_per_pubdata: Option<u64>,
 }
 
 impl ZkTransactionMetadata {
     /// Create a new [`ZkTransactionMetadata`] with the given factory deps
-    pub fn new(factory_deps: Vec<Vec<u8>>, paymaster_data: Option<PaymasterParams>) -> Self {
-        Self { factory_deps, paymaster_data }
+    pub fn new(
+        factory_deps: Vec<Vec<u8>>,
+        paymaster_data: Option<PaymasterParams>,
+        gas_per_pubdata: Option<u64>,
+    ) -> Self {
+        Self { factory_deps, paymaster_data, gas_per_pubdata }
     }
 }
 /// Estimated gas from a ZK network.
@@ -143,7 +149,10 @@ pub async fn estimate_fee<P: ZksyncProvider<T>, T: Transport + Clone>(
     let max_priority_fee = tx.max_priority_fee_per_gas().unwrap_or(fee.max_priority_fee_per_gas);
     tx.set_max_fee_per_gas(max_fee);
     tx.set_max_priority_fee_per_gas(max_priority_fee);
-    tx.set_gas_per_pubdata(fee.gas_per_pubdata_limit);
+    match tx.gas_per_pubdata() {
+        Some(v) if !v.is_zero() => (),
+        _ => tx.set_gas_per_pubdata(fee.gas_per_pubdata_limit),
+    }
 
     Ok(())
 }

crates/zksync/core/src/vm/env.rs

@@ -23,13 +23,20 @@ pub struct ZkEnv {
     pub fair_l2_gas_price: u64,
     /// fair pubdata price
     pub fair_pubdata_price: u64,
+    /// gas per pubdata
+    pub gas_per_pubdata: Option<u64>,
 }
 
 impl Default for ZkEnv {
     fn default() -> Self {
         // TODO: fair pubdata price of 0 yields division by 0 error somewhere in
         // some cases. Should investigate this edge case further
-        Self { l1_gas_price: 0, fair_l2_gas_price: 0, fair_pubdata_price: 1000 }
+        Self {
+            l1_gas_price: 0,
+            fair_l2_gas_price: 0,
+            fair_pubdata_price: 1000,
+            gas_per_pubdata: None,
+        }
     }
 }