Repository: FuelLabs/fuel-core (Rust)

Issue:
Update `submit` and `submitAndAwaitStatus` GraphQL End-Point
The `submit` and `submitAndAwaitStatus` endpoints should change to allow for more flexibility for predicates.

## Requirements

These end-points will receive a transaction almost on ready state, the only thing that can be missing is the predicate estimation. In this case, what should happen is:

- If input is a predicate and the `predicate_gas_used` is **zero**, it should also estimate the `predicate_gas_used`.
- If input is a predicate and the `predicate_gas_used` is grater than the final `predicate_gas_used` it should not throw.



Fix (patch):
diff --git a/.changes/added/2841.md b/.changes/added/2841.md
new file mode 100644
index 00000000000..af648a9f138
--- /dev/null
+++ b/.changes/added/2841.md
@@ -0,0 +1,6 @@
+Following endpoints allow estimation of predicates on submission of the transaction via new `estimatePredicates` argument:
+- `submit`
+- `submit_and_await`
+- `submit_and_await_status`
+
+The change is backward compatible with all SDKs. The change is not forward-compatible with Rust SDK in the case of the `estiamte_predicates` flag set.
\ No newline at end of file
diff --git a/crates/client/assets/schema.sdl b/crates/client/assets/schema.sdl
index d974add1173..c31e3938515 100644
--- a/crates/client/assets/schema.sdl
+++ b/crates/client/assets/schema.sdl
@@ -793,7 +793,7 @@ type Mutation {
 	
 	Returns submitted transaction if the transaction is included in the `TxPool` without problems.
 	"""
-	submit(tx: HexString!): Transaction!
+	submit(tx: HexString!, estimatePredicates: Boolean): Transaction!
 	"""
 	Sequentially produces `blocks_to_produce` blocks. The first block starts with
 	`start_timestamp`. If the block production in the [`crate::service::Config`] is
@@ -1339,13 +1339,13 @@ type Subscription {
 	"""
 	Submits transaction to the `TxPool` and await either confirmation or failure.
 	"""
-	submitAndAwait(tx: HexString!): TransactionStatus!
+	submitAndAwait(tx: HexString!, estimatePredicates: Boolean): TransactionStatus!
 	"""
 	Submits the transaction to the `TxPool` and returns a stream of events.
 	Compared to the `submitAndAwait`, the stream also contains `
 	SubmittedStatus` as an intermediate state.
 	"""
-	submitAndAwaitStatus(tx: HexString!): TransactionStatus!
+	submitAndAwaitStatus(tx: HexString!, estimatePredicates: Boolean): TransactionStatus!
 	contractStorageSlots(contractId: ContractId!): StorageSlot!
 	contractStorageBalances(contractId: ContractId!): ContractBalance!
 }
diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs
index 11038bb867d..89a04d068f9 100644
--- a/crates/client/src/client.rs
+++ b/crates/client/src/client.rs
@@ -12,7 +12,10 @@ use crate::{
             gas_price::EstimateGasPrice,
             message::MessageStatusArgs,
             relayed_tx::RelayedTransactionStatusArgs,
-            tx::DryRunArg,
+            tx::{
+                DryRunArg,
+                TxWithEstimatedPredicatesArg,
+            },
             Tai64Timestamp,
             TransactionId,
         },
@@ -825,31 +828,55 @@ impl FuelClient {
     pub async fn submit(
         &self,
         tx: &Transaction,
+    ) -> io::Result<types::primitives::TransactionId> {
+        self.submit_opt(tx, None).await
+    }
+
+    pub async fn submit_opt(
+        &self,
+        tx: &Transaction,
+        estimate_predicates: Option<bool>,
     ) -> io::Result<types::primitives::TransactionId> {
         let tx = tx.clone().to_bytes();
-        let query = schema::tx::Submit::build(TxArg {
+        let query = schema::tx::Submit::build(TxWithEstimatedPredicatesArg {
             tx: HexString(Bytes(tx)),
+            estimate_predicates,
         });
 
         let id = self.query(query).await.map(|r| r.submit)?.id.into();
         Ok(id)
     }
 
+    /// Similar to [`Self::submit_and_await_commit_opt`], but with default options.
+    #[cfg(feature = "subscriptions")]
+    pub async fn submit_and_await_commit(
+        &self,
+        tx: &Transaction,
+    ) -> io::Result<TransactionStatus> {
+        self.submit_and_await_commit_opt(tx, None).await
+    }
+
     /// Submit the transaction and wait for it either to be included in
     /// a block or removed from `TxPool`.
     ///
+    /// If `estimate_predicates` is set, the predicates will be estimated before
+    /// the transaction is inserted into transaction pool.
+    ///
     /// This will wait forever if needed, so consider wrapping this call
     /// with a `tokio::time::timeout`.
     #[cfg(feature = "subscriptions")]
-    pub async fn submit_and_await_commit(
+    pub async fn submit_and_await_commit_opt(
         &self,
         tx: &Transaction,
+        estimate_predicates: Option<bool>,
     ) -> io::Result<TransactionStatus> {
         use cynic::SubscriptionBuilder;
         let tx = tx.clone().to_bytes();
-        let s = schema::tx::SubmitAndAwaitSubscription::build(TxArg {
-            tx: HexString(Bytes(tx)),
-        });
+        let s =
+            schema::tx::SubmitAndAwaitSubscription::build(TxWithEstimatedPredicatesArg {
+                tx: HexString(Bytes(tx)),
+                estimate_predicates,
+            });
 
         let mut stream = self.subscribe(s).await?.map(
             |r: io::Result<schema::tx::SubmitAndAwaitSubscription>| {
@@ -871,12 +898,25 @@ impl FuelClient {
     pub async fn submit_and_await_commit_with_tx(
         &self,
         tx: &Transaction,
+    ) -> io::Result<StatusWithTransaction> {
+        self.submit_and_await_commit_with_tx_opt(tx, None).await
+    }
+
+    /// Similar to [`Self::submit_and_await_commit_opt`], but the status also contains transaction.
+    #[cfg(feature = "subscriptions")]
+    pub async fn submit_and_await_commit_with_tx_opt(
+        &self,
+        tx: &Transaction,
+        estimate_predicates: Option<bool>,
     ) -> io::Result<StatusWithTransaction> {
         use cynic::SubscriptionBuilder;
         let tx = tx.clone().to_bytes();
-        let s = schema::tx::SubmitAndAwaitSubscriptionWithTransaction::build(TxArg {
-            tx: HexString(Bytes(tx)),
-        });
+        let s = schema::tx::SubmitAndAwaitSubscriptionWithTransaction::build(
+            TxWithEstimatedPredicatesArg {
+                tx: HexString(Bytes(tx)),
+                estimate_predicates,
+            },
+        );
 
         let mut stream = self.subscribe(s).await?.map(
             |r: io::Result<schema::tx::SubmitAndAwaitSubscriptionWithTransaction>| {
@@ -893,19 +933,30 @@ impl FuelClient {
         Ok(status)
     }
 
-    /// Submits the transaction to the `TxPool` and returns a stream of events.
-    /// Compared to the `submit_and_await_commit`, the stream also contains
-    /// `SubmittedStatus` as an intermediate state.
+    /// Similar to [`Self::submit_and_await_commit`], but includes all intermediate states.
     #[cfg(feature = "subscriptions")]
     pub async fn submit_and_await_status<'a>(
         &'a self,
         tx: &'a Transaction,
+    ) -> io::Result<impl Stream<Item = io::Result<TransactionStatus>> + 'a> {
+        self.submit_and_await_status_opt(tx, None).await
+    }
+
+    /// Similar to [`Self::submit_and_await_commit_opt`], but includes all intermediate states.
+    #[cfg(feature = "subscriptions")]
+    pub async fn submit_and_await_status_opt<'a>(
+        &'a self,
+        tx: &'a Transaction,
+        estimate_predicates: Option<bool>,
     ) -> io::Result<impl Stream<Item = io::Result<TransactionStatus>> + 'a> {
         use cynic::SubscriptionBuilder;
         let tx = tx.clone().to_bytes();
-        let s = schema::tx::SubmitAndAwaitStatusSubscription::build(TxArg {
-            tx: HexString(Bytes(tx)),
-        });
+        let s = schema::tx::SubmitAndAwaitStatusSubscription::build(
+            TxWithEstimatedPredicatesArg {
+                tx: HexString(Bytes(tx)),
+                estimate_predicates,
+            },
+        );
 
         let stream = self.subscribe(s).await?.map(
             |r: io::Result<schema::tx::SubmitAndAwaitStatusSubscription>| {
diff --git a/crates/client/src/client/schema/tx.rs b/crates/client/src/client/schema/tx.rs
index 87b83ba6f35..3258d156980 100644
--- a/crates/client/src/client/schema/tx.rs
+++ b/crates/client/src/client/schema/tx.rs
@@ -491,6 +491,13 @@ pub struct TxArg {
     pub tx: HexString,
 }
 
+#[derive(cynic::QueryVariables)]
+pub struct TxWithEstimatedPredicatesArg {
+    pub tx: HexString,
+    #[cynic(skip_serializing_if = "Option::is_none")]
+    pub estimate_predicates: Option<bool>,
+}
+
 #[derive(cynic::QueryFragment, Clone, Debug)]
 #[cynic(
     schema_path = "./assets/schema.sdl",
@@ -604,10 +611,10 @@ pub struct DryRun {
 #[cynic(
     schema_path = "./assets/schema.sdl",
     graphql_type = "Mutation",
-    variables = "TxArg"
+    variables = "TxWithEstimatedPredicatesArg"
 )]
 pub struct Submit {
-    #[arguments(tx: $tx)]
+    #[arguments(tx: $tx, estimatePredicates: $estimate_predicates)]
     pub submit: TransactionIdFragment,
 }
 
@@ -615,10 +622,10 @@ pub struct Submit {
 #[cynic(
     schema_path = "./assets/schema.sdl",
     graphql_type = "Subscription",
-    variables = "TxArg"
+    variables = "TxWithEstimatedPredicatesArg"
 )]
 pub struct SubmitAndAwaitSubscription {
-    #[arguments(tx: $tx)]
+    #[arguments(tx: $tx, estimatePredicates: $estimate_predicates)]
     pub submit_and_await: TransactionStatus,
 }
 
@@ -626,10 +633,10 @@ pub struct SubmitAndAwaitSubscription {
 #[cynic(
     schema_path = "./assets/schema.sdl",
     graphql_type = "Subscription",
-    variables = "TxArg"
+    variables = "TxWithEstimatedPredicatesArg"
 )]
 pub struct SubmitAndAwaitSubscriptionWithTransaction {
-    #[arguments(tx: $tx)]
+    #[arguments(tx: $tx, estimatePredicates: $estimate_predicates)]
     pub submit_and_await: StatusWithTransaction,
 }
 
@@ -637,10 +644,10 @@ pub struct SubmitAndAwaitSubscriptionWithTransaction {
 #[cynic(
     schema_path = "./assets/schema.sdl",
     graphql_type = "Subscription",
-    variables = "TxArg"
+    variables = "TxWithEstimatedPredicatesArg"
 )]
 pub struct SubmitAndAwaitStatusSubscription {
-    #[arguments(tx: $tx)]
+    #[arguments(tx: $tx, estimatePredicates: $estimate_predicates)]
     pub submit_and_await_status: TransactionStatus,
 }
 
@@ -723,8 +730,9 @@ pub mod tests {
     fn submit_tx_gql_output() {
         use cynic::MutationBuilder;
         let tx = fuel_tx::Transaction::default_test_tx();
-        let query = Submit::build(TxArg {
+        let query = Submit::build(TxWithEstimatedPredicatesArg {
             tx: HexString(Bytes(tx.to_bytes())),
+            estimate_predicates: Some(true),
         });
         insta::assert_snapshot!(query.query)
     }
diff --git a/crates/fuel-core/src/schema/tx.rs b/crates/fuel-core/src/schema/tx.rs
index 349ebd7cad4..b626fe5aa7d 100644
--- a/crates/fuel-core/src/schema/tx.rs
+++ b/crates/fuel-core/src/schema/tx.rs
@@ -63,6 +63,7 @@ use fuel_core_storage::{
     iter::IterDirection,
     Error as StorageError,
     IsNotFound,
+    PredicateStorageRequirements,
     Result as StorageResult,
 };
 use fuel_core_tx_status_manager::TxStatusMessage;
@@ -91,6 +92,7 @@ use futures::{
 };
 use std::{
     borrow::Cow,
+    future::Future,
     iter,
     sync::Arc,
 };
@@ -431,24 +433,15 @@ impl TxQuery {
     ) -> async_graphql::Result<Transaction> {
         let query = ctx.read_view()?.into_owned();
 
-        let mut tx = FuelTx::from_bytes(&tx.0)?;
+        let tx = FuelTx::from_bytes(&tx.0)?;
 
-        let params = ctx
+        let tx = ctx.estimate_predicates(tx, query).await?;
+        let chain_id = ctx
             .data_unchecked::<ChainInfoProvider>()
-            .current_consensus_params();
-
-        let memory_pool = ctx.data_unchecked::<SharedMemoryPool>();
-        let memory = memory_pool.get_memory().await;
-
-        let parameters = CheckPredicateParams::from(params.as_ref());
-        let tx = tokio_rayon::spawn_fifo(move || {
-            let result = tx.estimate_predicates(&parameters, memory, &query);
-            result.map(|_| tx)
-        })
-        .await
-        .map_err(|err| anyhow::anyhow!("{:?}", err))?;
+            .current_consensus_params()
+            .chain_id();
 
-        Ok(Transaction::from_tx(tx.id(&params.chain_id()), tx))
+        Ok(Transaction::from_tx(tx.id(&chain_id), tx))
     }
 
     #[cfg(feature = "test-helpers")]
@@ -583,18 +576,26 @@ impl TxMutation {
         &self,
         ctx: &Context<'_>,
         tx: HexString,
+        estimate_predicates: Option<bool>,
     ) -> async_graphql::Result<Transaction> {
         let txpool = ctx.data_unchecked::<TxPool>();
-        let params = ctx
-            .data_unchecked::<ChainInfoProvider>()
-            .current_consensus_params();
-        let tx = FuelTx::from_bytes(&tx.0)?;
+        let mut tx = FuelTx::from_bytes(&tx.0)?;
+
+        if estimate_predicates.unwrap_or(false) {
+            let query = ctx.read_view()?.into_owned();
+            tx = ctx.estimate_predicates(tx, query).await?;
+        }
 
         txpool
             .insert(tx.clone())
             .await
             .map_err(|e| anyhow::anyhow!(e))?;
-        let id = tx.id(&params.chain_id());
+
+        let chain_id = ctx
+            .data_unchecked::<ChainInfoProvider>()
+            .current_consensus_params()
+            .chain_id();
+        let id = tx.id(&chain_id);
 
         let tx = Transaction(tx, id);
         Ok(tx)
@@ -646,11 +647,14 @@ impl TxStatusSubscription {
         &self,
         ctx: &'a Context<'a>,
         tx: HexString,
+        estimate_predicates: Option<bool>,
     ) -> async_graphql::Result<
         impl Stream<Item = async_graphql::Result<TransactionStatus>> + 'a,
     > {
         use tokio_stream::StreamExt;
-        let subscription = submit_and_await_status(ctx, tx).await?;
+        let subscription =
+            submit_and_await_status(ctx, tx, estimate_predicates.unwrap_or(false))
+                .await?;
 
         Ok(subscription
             .skip_while(|event| matches!(event, Ok(TransactionStatus::Submitted(..))))
@@ -665,16 +669,18 @@ impl TxStatusSubscription {
         &self,
         ctx: &'a Context<'a>,
         tx: HexString,
+        estimate_predicates: Option<bool>,
     ) -> async_graphql::Result<
         impl Stream<Item = async_graphql::Result<TransactionStatus>> + 'a,
     > {
-        submit_and_await_status(ctx, tx).await
+        submit_and_await_status(ctx, tx, estimate_predicates.unwrap_or(false)).await
     }
 }
 
 async fn submit_and_await_status<'a>(
     ctx: &'a Context<'a>,
     tx: HexString,
+    estimate_predicates: bool,
 ) -> async_graphql::Result<
     impl Stream<Item = async_graphql::Result<TransactionStatus>> + 'a,
 > {
@@ -684,8 +690,14 @@ async fn submit_and_await_status<'a>(
     let params = ctx
         .data_unchecked::<ChainInfoProvider>()
         .current_consensus_params();
-    let tx = FuelTx::from_bytes(&tx.0)?;
+    let mut tx = FuelTx::from_bytes(&tx.0)?;
     let tx_id = tx.id(&params.chain_id());
+
+    if estimate_predicates {
+        let query = ctx.read_view()?.into_owned();
+        tx = ctx.estimate_predicates(tx, query).await?;
+    }
+
     let subscription = tx_status_manager.tx_update_subscribe(tx_id).await?;
 
     txpool.insert(tx).await?;
@@ -842,7 +854,13 @@ pub trait ContextExt {
     fn try_find_tx(
         &self,
         id: Bytes32,
-    ) -> impl std::future::Future<Output = StorageResult<Option<FuelTx>>> + Send;
+    ) -> impl Future<Output = StorageResult<Option<FuelTx>>> + Send;
+
+    fn estimate_predicates(
+        &self,
+        tx: FuelTx,
+        query: impl PredicateStorageRequirements + Send + Sync + 'static,
+    ) -> impl Future<Output = anyhow::Result<FuelTx>> + Send;
 }
 
 impl<'a> ContextExt for Context<'a> {
@@ -862,4 +880,40 @@ impl<'a> ContextExt for Context<'a> {
             }
         }
     }
+
+    async fn estimate_predicates(
+        &self,
+        mut tx: FuelTx,
+        query: impl PredicateStorageRequirements + Send + Sync + 'static,
+    ) -> anyhow::Result<FuelTx> {
+        let mut has_predicates = false;
+
+        for input in tx.inputs()? {
+            if input.predicate().is_some() {
+                has_predicates = true;
+                break;
+            }
+        }
+
+        if !has_predicates {
+            return Ok(tx);
+        }
+
+        let params = self
+            .data_unchecked::<ChainInfoProvider>()
+            .current_consensus_params();
+
+        let memory_pool = self.data_unchecked::<SharedMemoryPool>();
+        let memory = memory_pool.get_memory().await;
+
+        let parameters = CheckPredicateParams::from(params.as_ref());
+        let tx = tokio_rayon::spawn_fifo(move || {
+            let result = tx.estimate_predicates(&parameters, memory, &query);
+            result.map(|_| tx)
+        })
+        .await
+        .map_err(|err| anyhow::anyhow!("{:?}", err))?;
+
+        Ok(tx)
+    }
 }