From 2a574a424986d398991d92a3c029ebd9e2dbb61a Mon Sep 17 00:00:00 2001
From: samkim-crypto <skim13@cs.stanford.edu>
Date: Fri, 12 Jan 2024 12:02:15 +0900
Subject: [PATCH 1/3] add `process_ixs_with_additional_compute_budget`

---
 token/client/src/token.rs | 27 +++++++++++++++++++++++++--
 1 file changed, 25 insertions(+), 2 deletions(-)

diff --git a/token/client/src/token.rs b/token/client/src/token.rs
index 51232ac9593..afa0088fd80 100644
--- a/token/client/src/token.rs
+++ b/token/client/src/token.rs
@@ -508,6 +508,7 @@ where
     async fn construct_tx<S: Signers>(
         &self,
         token_instructions: &[Instruction],
+        additional_compute_budget: Option<u32>,
         signing_keypairs: &S,
     ) -> TokenResult<Transaction> {
         let mut instructions = vec![];
@@ -532,6 +533,14 @@ where
 
         instructions.extend_from_slice(token_instructions);
 
+        if let Some(additional_compute_budget) = additional_compute_budget {
+            instructions.push(
+                solana_sdk::compute_budget::ComputeBudgetInstruction::set_compute_unit_limit(
+                    additional_compute_budget,
+                ),
+            );
+        }
+
         let (message, blockhash) =
             if let (Some(nonce_account), Some(nonce_authority), Some(nonce_blockhash)) = (
                 self.nonce_account,
@@ -581,7 +590,7 @@ where
         signing_keypairs: &S,
     ) -> TokenResult<T::SimulationOutput> {
         let transaction = self
-            .construct_tx(token_instructions, signing_keypairs)
+            .construct_tx(token_instructions, None, signing_keypairs)
             .await?;
 
         self.client
@@ -594,9 +603,23 @@ where
         &self,
         token_instructions: &[Instruction],
         signing_keypairs: &S,
+    ) -> TokenResult<T::Output> {
+        self.process_ixs_with_additional_compute_budget(token_instructions, None, signing_keypairs)
+            .await
+    }
+
+    pub async fn process_ixs_with_additional_compute_budget<S: Signers>(
+        &self,
+        token_instructions: &[Instruction],
+        additional_compute_budget: Option<u32>,
+        signing_keypairs: &S,
     ) -> TokenResult<T::Output> {
         let transaction = self
-            .construct_tx(token_instructions, signing_keypairs)
+            .construct_tx(
+                token_instructions,
+                additional_compute_budget,
+                signing_keypairs,
+            )
             .await?;
 
         self.client

From 23cb9c481963cd2b366732d572a2c2cf81327c86 Mon Sep 17 00:00:00 2001
From: samkim-crypto <skim13@cs.stanford.edu>
Date: Fri, 12 Jan 2024 12:03:55 +0900
Subject: [PATCH 2/3] add extra compute budget for transfer with fee

---
 token/client/src/token.rs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/token/client/src/token.rs b/token/client/src/token.rs
index afa0088fd80..e71f9f6bf84 100644
--- a/token/client/src/token.rs
+++ b/token/client/src/token.rs
@@ -2641,7 +2641,10 @@ where
             .new_decryptable_available_balance(transfer_amount, source_aes_key)
             .map_err(|_| TokenError::AccountDecryption)?;
 
-        self.process_ixs(
+        // additional compute budget required for `VerifyTransferWithFee`
+        const TRANSFER_WITH_FEE_COMPUTE_BUDGET: u32 = 500_000;
+
+        self.process_ixs_with_additional_compute_budget(
             &confidential_transfer::instruction::transfer_with_fee(
                 &self.program_id,
                 source_account,
@@ -2652,6 +2655,7 @@ where
                 &multisig_signers,
                 proof_location,
             )?,
+            Some(TRANSFER_WITH_FEE_COMPUTE_BUDGET),
             signing_keypairs,
         )
         .await

From 34936d94788c073eee95b8f3b092872b66db9538 Mon Sep 17 00:00:00 2001
From: samkim-crypto <skim13@cs.stanford.edu>
Date: Fri, 12 Jan 2024 22:23:09 +0900
Subject: [PATCH 3/3] update `process_ixs_with_additional_compute_budget` to
 take in a regular `u32`

---
 token/client/src/token.rs | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/token/client/src/token.rs b/token/client/src/token.rs
index e71f9f6bf84..bf8f2236794 100644
--- a/token/client/src/token.rs
+++ b/token/client/src/token.rs
@@ -604,20 +604,26 @@ where
         token_instructions: &[Instruction],
         signing_keypairs: &S,
     ) -> TokenResult<T::Output> {
-        self.process_ixs_with_additional_compute_budget(token_instructions, None, signing_keypairs)
+        let transaction = self
+            .construct_tx(token_instructions, None, signing_keypairs)
+            .await?;
+
+        self.client
+            .send_transaction(&transaction)
             .await
+            .map_err(TokenError::Client)
     }
 
     pub async fn process_ixs_with_additional_compute_budget<S: Signers>(
         &self,
         token_instructions: &[Instruction],
-        additional_compute_budget: Option<u32>,
+        additional_compute_budget: u32,
         signing_keypairs: &S,
     ) -> TokenResult<T::Output> {
         let transaction = self
             .construct_tx(
                 token_instructions,
-                additional_compute_budget,
+                Some(additional_compute_budget),
                 signing_keypairs,
             )
             .await?;
@@ -2655,7 +2661,7 @@ where
                 &multisig_signers,
                 proof_location,
             )?,
-            Some(TRANSFER_WITH_FEE_COMPUTE_BUDGET),
+            TRANSFER_WITH_FEE_COMPUTE_BUDGET,
             signing_keypairs,
         )
         .await