[cherry-pick][cli] Add gas estimate feature (#17322) (#17457)

## Description 

Cherry picking feature into `releases/sui-v1.24.0-release` branch for
the next release next week.
This PR adds automatic gas estimation by dry running the transaction if
the gas budget is not provided. That means that the `gas_budget` flag is
now optional.

## Test plan 

Existing tests + new test
```
cd crates/sui && cargo test -- test_gas_estimation
```

---

## Release notes

Check each box that your changes affect. If none of the boxes relate to
your changes, release notes aren't required.

For each box you select, include information after the relevant heading
that describes the impact of your changes that a user might notice and
any actions they must take to implement updates.

- [ ] Protocol: 
- [ ] Nodes (Validators and Full nodes): 
- [ ] Indexer: 
- [ ] JSON-RPC: 
- [ ] GraphQL: 
- [x] CLI: Added automatic gas estimation feature for the Sui CLI. If
gas budget is not provided, the tool will dry run the transaction to get
a gas budget estimate, and then it will execute the transaction. That
means that for all relevant commands the `--gas-budget` flag is now
optional. Please note that this incurs a small cost in performance due
to the additional dry run call.
- [ ] Rust SDK:

crates/sui/src/client_commands.rs

@@ -133,6 +133,9 @@ macro_rules! serialize_or_execute {
     }};
 }
 
+/// Only to be used within CLI
+pub const GAS_SAFE_OVERHEAD: u64 = 1000;
+
 #[derive(Parser)]
 #[clap(rename_all = "kebab-case")]
 pub enum SuiClientCommands {
@@ -613,9 +616,12 @@ pub enum SuiClientCommands {
 /// Global options for most transaction execution related commands
 #[derive(Args, Debug)]
 pub struct Opts {
-    /// Gas budget for this transaction (in MIST)
+    /// An optional gas budget for this transaction (in MIST). If gas budget is not provided, the
+    /// tool will first perform a dry run to estimate the gas cost, and then it will execute the
+    /// transaction. Please note that this incurs a small cost in performance due to the additional
+    /// dry run call.
     #[arg(long)]
-    pub gas_budget: u64,
+    pub gas_budget: Option<u64>,
     /// Perform a dry run of the transaction, without executing it.
     #[arg(long)]
     pub dry_run: bool,
@@ -647,7 +653,7 @@ impl Opts {
     /// Uses the passed gas_budget for the gas budget variable and sets all other flags to false.
     pub fn for_testing(gas_budget: u64) -> Self {
         Self {
-            gas_budget,
+            gas_budget: Some(gas_budget),
             dry_run: false,
             serialize_unsigned_transaction: false,
             serialize_signed_transaction: false,
@@ -657,7 +663,7 @@ impl Opts {
     /// and sets all other flags to false.
     pub fn for_testing_dry_run(gas_budget: u64) -> Self {
         Self {
-            gas_budget,
+            gas_budget: Some(gas_budget),
             dry_run: true,
             serialize_unsigned_transaction: false,
             serialize_signed_transaction: false,
@@ -935,11 +941,27 @@ impl SuiClientCommands {
                     )
                     .await;
                 }
+
+                let gas_budget = match gas_budget {
+                    Some(gas_budget) => gas_budget,
+                    None => {
+                        estimate_gas_budget(
+                            context,
+                            sender,
+                            tx_kind.clone(),
+                            gas_price,
+                            gas.map(|x| vec![x]),
+                            None,
+                        )
+                        .await?
+                    }
+                };
+
                 let data = client
                     .transaction_builder()
                     .tx_data(
                         sender,
-                        tx_kind,
+                        tx_kind.clone(),
                         gas_budget,
                         gas_price,
                         gas.into_iter().collect(),
@@ -1041,6 +1063,21 @@ impl SuiClientCommands {
                     .await;
                 }
 
+                let gas_budget = match gas_budget {
+                    Some(gas_budget) => gas_budget,
+                    None => {
+                        estimate_gas_budget(
+                            context,
+                            sender,
+                            tx_kind.clone(),
+                            gas_price,
+                            gas.map(|x| vec![x]),
+                            None,
+                        )
+                        .await?
+                    }
+                };
+
                 let data = client
                     .transaction_builder()
                     .tx_data(
@@ -1266,6 +1303,20 @@ impl SuiClientCommands {
                     )
                     .await;
                 }
+                let gas_budget = match gas_budget {
+                    Some(gas_budget) => gas_budget,
+                    None => {
+                        estimate_gas_budget(
+                            context,
+                            signer,
+                            tx_kind.clone(),
+                            gas_price,
+                            gas.map(|x| vec![x]),
+                            None,
+                        )
+                        .await?
+                    }
+                };
                 let data = client
                     .transaction_builder()
                     .tx_data(
@@ -1321,6 +1372,20 @@ impl SuiClientCommands {
                     )
                     .await;
                 }
+                let gas_budget = match gas_budget {
+                    Some(gas_budget) => gas_budget,
+                    None => {
+                        estimate_gas_budget(
+                            context,
+                            from,
+                            tx_kind.clone(),
+                            gas_price,
+                            gas.map(|x| vec![x]),
+                            None,
+                        )
+                        .await?
+                    }
+                };
                 let data = client
                     .transaction_builder()
                     .tx_data(
@@ -1372,6 +1437,20 @@ impl SuiClientCommands {
                     )
                     .await;
                 }
+                let gas_budget = match gas_budget {
+                    Some(gas_budget) => gas_budget,
+                    None => {
+                        estimate_gas_budget(
+                            context,
+                            from,
+                            tx_kind.clone(),
+                            gas_price,
+                            Some(vec![object_id]),
+                            None,
+                        )
+                        .await?
+                    }
+                };
                 let data = client
                     .transaction_builder()
                     .tx_data(
@@ -1456,6 +1535,20 @@ impl SuiClientCommands {
                     .await;
                 }
 
+                let gas_budget = match gas_budget {
+                    Some(gas_budget) => gas_budget,
+                    None => {
+                        estimate_gas_budget(
+                            context,
+                            from,
+                            kind.clone(),
+                            gas_price,
+                            gas.map(|x| vec![x]),
+                            None,
+                        )
+                        .await?
+                    }
+                };
                 let data = client
                     .transaction_builder()
                     .tx_data(
@@ -1530,6 +1623,20 @@ impl SuiClientCommands {
                     .await;
                 }
 
+                let gas_budget = match gas_budget {
+                    Some(gas_budget) => gas_budget,
+                    None => {
+                        estimate_gas_budget(
+                            context,
+                            signer,
+                            kind.clone(),
+                            gas_price,
+                            Some(input_coins.clone()),
+                            None,
+                        )
+                        .await?
+                    }
+                };
                 let data = client
                     .transaction_builder()
                     .tx_data(
@@ -1583,6 +1690,20 @@ impl SuiClientCommands {
                     .await;
                 }
 
+                let gas_budget = match gas_budget {
+                    Some(gas_budget) => gas_budget,
+                    None => {
+                        estimate_gas_budget(
+                            context,
+                            signer,
+                            tx_kind.clone(),
+                            gas_price,
+                            Some(input_coins.clone()),
+                            None,
+                        )
+                        .await?
+                    }
+                };
                 let data = client
                     .transaction_builder()
                     .tx_data(signer, tx_kind, gas_budget, gas_price, input_coins, None)
@@ -1735,6 +1856,21 @@ impl SuiClientCommands {
                     .await;
                 }
 
+                let gas_budget = match gas_budget {
+                    Some(gas_budget) => gas_budget,
+                    None => {
+                        estimate_gas_budget(
+                            context,
+                            signer,
+                            tx_kind.clone(),
+                            gas_price,
+                            gas.map(|x| vec![x]),
+                            None,
+                        )
+                        .await?
+                    }
+                };
+
                 let data = client
                     .transaction_builder()
                     .tx_data(
@@ -1788,6 +1924,21 @@ impl SuiClientCommands {
                     )
                     .await;
                 }
+
+                let gas_budget = match gas_budget {
+                    Some(gas_budget) => gas_budget,
+                    None => {
+                        estimate_gas_budget(
+                            context,
+                            signer,
+                            tx_kind.clone(),
+                            gas_price,
+                            gas.map(|x| vec![x]),
+                            None,
+                        )
+                        .await?
+                    }
+                };
                 let data = client
                     .transaction_builder()
                     .tx_data(
@@ -2998,15 +3149,19 @@ fn format_balance(
 }
 
 /// Helper function to reduce code duplication for executing dry run
-async fn execute_dry_run(
+pub async fn execute_dry_run(
     context: &mut WalletContext,
     signer: SuiAddress,
     kind: TransactionKind,
-    gas_budget: u64,
+    gas_budget: Option<u64>,
     gas_price: u64,
     gas_payment: Option<Vec<ObjectID>>,
     sponsor: Option<SuiAddress>,
 ) -> Result<SuiClientCommandResult, anyhow::Error> {
+    let gas_budget = match gas_budget {
+        Some(gas_budget) => gas_budget,
+        None => max_gas_budget(context).await?,
+    };
     let dry_run_tx_data = context
         .get_client()
         .await?
@@ -3022,3 +3177,52 @@ async fn execute_dry_run(
         .map_err(|e| anyhow!("Dry run failed: {e}"))?;
     Ok(SuiClientCommandResult::DryRun(response))
 }
+
+/// Call a dry run with the transaction data to estimate the gas budget.
+/// The estimated gas budget is computed as following:
+/// * the maximum between A and B, where:
+/// A = computation cost + GAS_SAFE_OVERHEAD * reference gas price
+/// B = computation cost + storage cost - storage rebate + GAS_SAFE_OVERHEAD * reference gas price
+/// overhead
+///
+/// This gas estimate is computed exactly as in the TypeScript SDK
+/// <https://github.com/MystenLabs/sui/blob/3c4369270605f78a243842098b7029daf8d883d9/sdk/typescript/src/transactions/TransactionBlock.ts#L845-L858>
+pub async fn estimate_gas_budget(
+    context: &mut WalletContext,
+    signer: SuiAddress,
+    kind: TransactionKind,
+    gas_price: u64,
+    gas_payment: Option<Vec<ObjectID>>,
+    sponsor: Option<SuiAddress>,
+) -> Result<u64, anyhow::Error> {
+    let client = context.get_client().await?;
+    let Ok(SuiClientCommandResult::DryRun(dry_run)) =
+        execute_dry_run(context, signer, kind, None, gas_price, gas_payment, sponsor).await
+    else {
+        bail!("Could not automatically determine the gas budget. Please supply one using the --gas-budget flag.")
+    };
+
+    let safe_overhead = GAS_SAFE_OVERHEAD * client.read_api().get_reference_gas_price().await?;
+    let computation_cost_with_overhead =
+        dry_run.effects.gas_cost_summary().computation_cost + safe_overhead;
+
+    let gas_usage = dry_run.effects.gas_cost_summary().net_gas_usage() + safe_overhead as i64;
+    Ok(computation_cost_with_overhead.max(if gas_usage < 0 { 0 } else { gas_usage as u64 }))
+}
+
+/// Queries the protocol config for the maximum gas allowed in a transaction.
+pub async fn max_gas_budget(context: &mut WalletContext) -> Result<u64, anyhow::Error> {
+    let cfg = context
+        .get_client()
+        .await?
+        .read_api()
+        .get_protocol_config(None)
+        .await?;
+    Ok(match cfg.attributes.get("max_tx_gas") {
+        Some(Some(sui_json_rpc_types::SuiProtocolConfigValue::U64(y))) => *y,
+        _ => bail!(
+            "Could not automatically find the maximum gas allowed in a transaction from the \
+            protocol config. Please provide a gas budget with the --gas-budget flag."
+        ),
+    })
+}

crates/sui/src/client_ptb/ast.rs

@@ -108,7 +108,7 @@ pub struct ProgramMetadata {
     pub gas_object_id: Option<Spanned<ObjectID>>,
     pub json_set: bool,
     pub dry_run_set: bool,
-    pub gas_budget: Spanned<u64>,
+    pub gas_budget: Option<Spanned<u64>>,
 }
 
 /// A parsed module access consisting of the address, module name, and function name.

crates/sui/src/client_ptb/displays/ptb_preview.rs

@@ -24,12 +24,9 @@ impl<'a> Display for PTBPreview<'a> {
                 builder.push_record([command, vals]);
             }
         }
-        // index of horizontal line to draw after commands
-        let line_index = builder.count_rows();
-        builder.push_record([
-            GAS_BUDGET,
-            self.program_metadata.gas_budget.value.to_string().as_str(),
-        ]);
+        if let Some(gas_budget) = self.program_metadata.gas_budget {
+            builder.push_record([GAS_BUDGET, gas_budget.value.to_string().as_str()]);
+        }
         if let Some(gas_coin_id) = self.program_metadata.gas_object_id {
             builder.push_record([GAS_COIN, gas_coin_id.value.to_string().as_str()]);
         }
@@ -53,8 +50,6 @@ impl<'a> Display for PTBPreview<'a> {
         table.with(TableStyle::rounded().horizontals([
             HorizontalLine::new(1, TableStyle::modern().get_horizontal()),
             HorizontalLine::new(2, TableStyle::modern().get_horizontal()),
-            HorizontalLine::new(2, TableStyle::modern().get_horizontal()),
-            HorizontalLine::new(line_index + 2, TableStyle::modern().get_horizontal()),
         ]));
         table.with(tabled::settings::style::BorderSpanCorrection);
 

crates/sui/src/client_ptb/parser.rs

@@ -193,14 +193,6 @@ impl<'a, I: Iterator<Item = &'a str>> ProgramParser<'a, I> {
                 .push(err!(sp, "Trailing {tok} found after the last command",));
         }
 
-        let Some(gas_budget) = self.state.gas_budget else {
-            self.state.errors.push(err!(
-                sp => help: { "Use --gas-budget <u64> to set a gas budget" },
-                "Gas budget not set."
-            ));
-            return Err(self.state.errors);
-        };
-
         if self.state.errors.is_empty() {
             Ok((
                 A::Program {
@@ -215,7 +207,7 @@ impl<'a, I: Iterator<Item = &'a str>> ProgramParser<'a, I> {
                     gas_object_id: self.state.gas_object_id,
                     json_set: self.state.json_set,
                     dry_run_set: self.state.dry_run_set,
-                    gas_budget,
+                    gas_budget: self.state.gas_budget,
                 },
             ))
         } else {