[Solana]: Add SolanaTransaction functions to get/set compute unit p…

…rice/limit

rust/chains/tw_solana/src/modules/compiled_keys.rs

@@ -11,6 +11,12 @@ use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use tw_coin_entry::error::prelude::*;
 
+pub fn try_into_u8(num: usize) -> SigningResult<u8> {
+    u8::try_from(num)
+        .tw_err(|_| SigningErrorType::Error_tx_too_big)
+        .context("There are too many accounts in the transaction")
+}
+
 #[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
 struct CompiledKeyMeta {
     is_signer: bool,
@@ -67,10 +73,6 @@ impl CompiledKeys {
     }
 
     pub fn try_into_message_components(self) -> SigningResult<(MessageHeader, Vec<SolanaAddress>)> {
-        let try_into_u8 = |num: usize| -> SigningResult<u8> {
-            u8::try_from(num).tw_err(|_| SigningErrorType::Error_tx_too_big)
-        };
-
         let Self {
             ordered_keys,
             key_meta_map,

rust/chains/tw_solana/src/modules/insert_instruction.rs

@@ -0,0 +1,104 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Copyright © 2017 Trust Wallet.
+
+use crate::address::SolanaAddress;
+use crate::modules::compiled_keys::try_into_u8;
+use crate::transaction::v0::MessageAddressTableLookup;
+use crate::transaction::{CompiledInstruction, MessageHeader};
+use std::iter;
+use tw_coin_entry::error::prelude::*;
+use tw_memory::Data;
+
+pub trait InsertInstruction {
+    /// Pushes a simple instruction that doesn't have accounts.
+    fn push_simple_instruction(
+        &mut self,
+        program_id: SolanaAddress,
+        data: Data,
+    ) -> SigningResult<()> {
+        let insert_at = self.instructions_mut().len();
+        self.insert_simple_instruction(insert_at, program_id, data)
+    }
+
+    /// Inserts a simple instruction that doesn't have accounts at the given `insert_at` index.
+    fn insert_simple_instruction(
+        &mut self,
+        insert_at: usize,
+        program_id: SolanaAddress,
+        data: Data,
+    ) -> SigningResult<()> {
+        if insert_at > self.instructions_mut().len() {
+            return SigningError::err(SigningErrorType::Error_internal)
+                .context(format!("Unable to add '{program_id}' instruction at the '{insert_at}' index. Number of existing instructions: {}", self.instructions_mut().len()));
+        }
+
+        // Step 1 - find or add the `program_id` in the accounts list.
+        let program_id_index = match self
+            .account_keys_mut()
+            .iter()
+            .position(|acc| *acc == program_id)
+        {
+            Some(pos) => try_into_u8(pos)?,
+            None => self.push_readonly_unsigned_account(program_id)?,
+        };
+
+        // Step 2 - Create a `CompiledInstruction` based on the `program_id` index and instruction `data`.
+        let new_compiled_ix = CompiledInstruction {
+            program_id_index,
+            accounts: Vec::default(),
+            data,
+        };
+
+        // Step 3 - Insert the created instruction at the given `insert_at` index.
+        self.instructions_mut().insert(insert_at, new_compiled_ix);
+
+        Ok(())
+    }
+
+    fn push_readonly_unsigned_account(&mut self, account: SolanaAddress) -> SigningResult<u8> {
+        debug_assert!(
+            !self.account_keys_mut().contains(&account),
+            "Account must not be in the account list yet"
+        );
+
+        self.account_keys_mut().push(account);
+        self.message_header_mut().num_readonly_unsigned_accounts += 1;
+
+        let account_added_at = try_into_u8(self.account_keys_mut().len() - 1)?;
+
+        // There is no need to update instruction account ids that point to [`Message::account_keys`] list
+        // as we pushed the account to the end of the list.
+        // But we must update it in case if there are instruction account ids that point to [`address_table_lookups`].
+
+        match self.address_table_lookups() {
+            Some(lookups) if !lookups.is_empty() => (),
+            // No address table lookups, no need to update the indexes.
+            _ => return Ok(account_added_at),
+        }
+
+        self.instructions_mut()
+            .iter_mut()
+            .flat_map(|ix| {
+                ix.accounts
+                    .iter_mut()
+                    .chain(iter::once(&mut ix.program_id_index))
+            })
+            // Update every instruction account id that points to the address table lookups.
+            .filter(|ix_account_id| **ix_account_id >= account_added_at)
+            .for_each(|ix_account_id| *ix_account_id += 1);
+
+        Ok(account_added_at)
+    }
+
+    /// Returns ALT (Address Lookup Tables) if supported by the message version.
+    fn address_table_lookups(&self) -> Option<&[MessageAddressTableLookup]> {
+        None
+    }
+
+    fn account_keys_mut(&mut self) -> &mut Vec<SolanaAddress>;
+
+    fn message_header_mut(&mut self) -> &mut MessageHeader;
+
+    fn instructions_mut(&mut self) -> &mut Vec<CompiledInstruction>;
+}

rust/chains/tw_solana/src/modules/instruction_builder/compute_budget_instruction.rs

@@ -5,6 +5,7 @@
 use crate::defined_addresses::COMPUTE_BUDGET_ADDRESS;
 use crate::instruction::Instruction;
 use borsh::{BorshDeserialize, BorshSerialize};
+use tw_encoding::{EncodingError, EncodingResult};
 
 pub type UnitLimit = u32;
 pub type UnitPrice = u64;
@@ -18,6 +19,12 @@ pub enum ComputeBudgetInstruction {
     SetLoadedAccountsDataSizeLimit(u32),
 }
 
+impl ComputeBudgetInstruction {
+    pub fn try_from_borsh(data: &[u8]) -> EncodingResult<Self> {
+        borsh::from_slice(data).map_err(|_| EncodingError::InvalidInput)
+    }
+}
+
 pub struct ComputeBudgetInstructionBuilder;
 
 impl ComputeBudgetInstructionBuilder {

rust/chains/tw_solana/src/modules/instruction_builder/system_instruction.rs

@@ -6,6 +6,7 @@ use crate::address::SolanaAddress;
 use crate::defined_addresses::*;
 use crate::instruction::{AccountMeta, Instruction};
 use serde::{Deserialize, Serialize};
+use tw_encoding::{EncodingError, EncodingResult};
 
 /// An instruction to the system program.
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
@@ -181,6 +182,12 @@ pub enum SystemInstruction {
     UpgradeNonceAccount,
 }
 
+impl SystemInstruction {
+    pub fn try_from_bincode(data: &[u8]) -> EncodingResult<Self> {
+        bincode::deserialize(data).map_err(|_| EncodingError::InvalidInput)
+    }
+}
+
 pub struct SystemInstructionBuilder;
 
 impl SystemInstructionBuilder {

rust/chains/tw_solana/src/modules/message_decompiler.rs

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Copyright © 2017 Trust Wallet.
+
+use crate::address::SolanaAddress;
+use crate::transaction::versioned::VersionedMessage;
+use crate::transaction::CompiledInstruction;
+use tw_coin_entry::error::prelude::{OrTWError, ResultContext, SigningErrorType, SigningResult};
+use tw_memory::Data;
+
+/// [`Instruction`] without `accounts` field.
+pub struct InstructionWithoutAccounts {
+    /// Pubkey of the program that executes this instruction.
+    pub program_id: SolanaAddress,
+    /// Opaque data passed to the program for its own interpretation.
+    pub data: Data,
+}
+
+pub struct MessageDecompiler;
+
+impl MessageDecompiler {
+    pub fn decompile_partly(
+        message: &VersionedMessage,
+    ) -> SigningResult<Vec<InstructionWithoutAccounts>> {
+        match message {
+            VersionedMessage::Legacy(legacy) => {
+                Self::decompile_partly_impl(&legacy.instructions, &legacy.account_keys)
+            },
+            VersionedMessage::V0(v0) => {
+                Self::decompile_partly_impl(&v0.instructions, &v0.account_keys)
+            },
+        }
+    }
+
+    fn decompile_partly_impl(
+        instructions: &[CompiledInstruction],
+        account_keys: &[SolanaAddress],
+    ) -> SigningResult<Vec<InstructionWithoutAccounts>> {
+        instructions
+            .iter()
+            .map(|ix| Self::decompile_instruction_partly(ix, account_keys))
+            .collect()
+    }
+
+    fn decompile_instruction_partly(
+        ix: &CompiledInstruction,
+        account_keys: &[SolanaAddress],
+    ) -> SigningResult<InstructionWithoutAccounts> {
+        // Program ID should always be in the transaction's accounts list even if AddressLookupTable is used:
+        // https://solana.stackexchange.com/questions/16122/using-program-ids-in-address-lookup-tables-missing-documentation-about-luts
+        let program_id = *account_keys
+            .get(ix.program_id_index as usize)
+            .or_tw_err(SigningErrorType::Error_invalid_params)
+            .context("Program ID not found in the accounts list")?;
+        Ok(InstructionWithoutAccounts {
+            program_id,
+            data: ix.data.clone(),
+        })
+    }
+}

rust/chains/tw_solana/src/modules/mod.rs

@@ -8,8 +8,10 @@ use tw_keypair::ed25519;
 
 pub mod compiled_instructions;
 pub mod compiled_keys;
+pub mod insert_instruction;
 pub mod instruction_builder;
 pub mod message_builder;
+pub mod message_decompiler;
 pub mod proto_builder;
 pub mod transaction_decoder;
 pub mod transaction_util;

rust/chains/tw_solana/src/modules/utils.rs

@@ -2,6 +2,13 @@
 //
 // Copyright © 2017 Trust Wallet.
 
+use crate::defined_addresses::{COMPUTE_BUDGET_ADDRESS, SYSTEM_PROGRAM_ID_ADDRESS};
+use crate::modules::insert_instruction::InsertInstruction;
+use crate::modules::instruction_builder::compute_budget_instruction::{
+    ComputeBudgetInstruction, ComputeBudgetInstructionBuilder, UnitLimit, UnitPrice,
+};
+use crate::modules::instruction_builder::system_instruction::SystemInstruction;
+use crate::modules::message_decompiler::{InstructionWithoutAccounts, MessageDecompiler};
 use crate::modules::proto_builder::ProtoBuilder;
 use crate::modules::tx_signer::TxSigner;
 use crate::modules::PubkeySignatureMap;
@@ -75,4 +82,113 @@ impl SolanaTransaction {
             ..Proto::SigningOutput::default()
         })
     }
+
+    pub fn get_compute_unit_price(encoded_tx: &str) -> SigningResult<Option<UnitPrice>> {
+        let tx = VersionedTransaction::from_base64(encoded_tx)?;
+        let instructions = MessageDecompiler::decompile_partly(&tx.message)?;
+        Ok(instructions
+            .iter()
+            .find_map(try_instruction_as_set_unit_price))
+    }
+
+    pub fn get_compute_unit_limit(encoded_tx: &str) -> SigningResult<Option<UnitLimit>> {
+        let tx = VersionedTransaction::from_base64(encoded_tx)?;
+        let instructions = MessageDecompiler::decompile_partly(&tx.message)?;
+        Ok(instructions
+            .iter()
+            .find_map(try_instruction_as_set_unit_limit))
+    }
+
+    pub fn set_compute_unit_price(encoded_tx: &str, price: UnitPrice) -> SigningResult<()> {
+        let tx_bytes = base64::decode(encoded_tx, STANDARD)?;
+        let mut tx: VersionedTransaction =
+            bincode::deserialize(&tx_bytes).map_err(|_| SigningErrorType::Error_input_parse)?;
+        let instructions = MessageDecompiler::decompile_partly(&tx.message)?;
+
+        let set_price_ix = ComputeBudgetInstructionBuilder::set_compute_unit_price(price);
+
+        // First, try to find a `ComputeBudgetInstruction::SetComputeUnitPrice` instruction.
+        let ix_position = instructions
+            .iter()
+            .position(|ix| try_instruction_as_set_unit_price(ix).is_some());
+        // If it presents already, it's enough to update the instruction data only.
+        if let Some(pos) = ix_position {
+            tx.message.instructions_mut()[pos].data = set_price_ix.data;
+            return Ok(());
+        }
+
+        // `ComputeBudgetInstruction::SetComputeUnitPrice` can be pushed to the end of the instructions list.
+        tx.message
+            .push_simple_instruction(set_price_ix.program_id, set_price_ix.data)?;
+
+        Ok(())
+    }
+
+    pub fn set_compute_unit_limit(encoded_tx: &str, limit: UnitLimit) -> SigningResult<()> {
+        let tx_bytes = base64::decode(encoded_tx, STANDARD)?;
+        let mut tx: VersionedTransaction =
+            bincode::deserialize(&tx_bytes).map_err(|_| SigningErrorType::Error_input_parse)?;
+        let instructions = MessageDecompiler::decompile_partly(&tx.message)?;
+
+        let set_limit_ix = ComputeBudgetInstructionBuilder::set_compute_unit_limit(limit);
+
+        // First, try to find a `ComputeBudgetInstruction::SetComputeUnitLimit` instruction.
+        let ix_position = instructions
+            .iter()
+            .position(|ix| try_instruction_as_set_unit_limit(ix).is_some());
+        // If it presents already, it's enough to update the instruction data only.
+        if let Some(pos) = ix_position {
+            tx.message.instructions_mut()[pos].data = set_limit_ix.data;
+            return Ok(());
+        }
+
+        // `ComputeBudgetInstruction::SetComputeUnitLimit` should be at the beginning of the instructions list.
+        // However `SystemInstruction::AdvanceNonceAccount` must be the first instruction.
+        // So in case if the advance nonce instruction presents, we should insert unit limit as the second instruction.
+        let insert_at = match instructions.first() {
+            Some(first_ix) if is_instruction_advance_nonce_account(first_ix) => 1,
+            _ => 0,
+        };
+
+        tx.message.insert_simple_instruction(
+            insert_at,
+            set_limit_ix.program_id,
+            set_limit_ix.data,
+        )?;
+
+        Ok(())
+    }
+}
+
+fn try_instruction_as_compute_budget(
+    ix: &InstructionWithoutAccounts,
+) -> Option<ComputeBudgetInstruction> {
+    if ix.program_id != *COMPUTE_BUDGET_ADDRESS {
+        return None;
+    }
+    ComputeBudgetInstruction::try_from_borsh(&ix.data).ok()
+}
+
+fn try_instruction_as_set_unit_price(ix: &InstructionWithoutAccounts) -> Option<UnitPrice> {
+    match try_instruction_as_compute_budget(ix)? {
+        ComputeBudgetInstruction::SetComputeUnitPrice(price) => Some(price),
+        _ => None,
+    }
+}
+
+fn try_instruction_as_set_unit_limit(ix: &InstructionWithoutAccounts) -> Option<UnitLimit> {
+    match try_instruction_as_compute_budget(ix)? {
+        ComputeBudgetInstruction::SetComputeUnitLimit(limit) => Some(limit),
+        _ => None,
+    }
+}
+
+fn is_instruction_advance_nonce_account(ix: &InstructionWithoutAccounts) -> bool {
+    if ix.program_id != *SYSTEM_PROGRAM_ID_ADDRESS {
+        return false;
+    }
+    let Ok(system_ix) = SystemInstruction::try_from_bincode(&ix.data) else {
+        return false;
+    };
+    system_ix == SystemInstruction::AdvanceNonceAccount
 }

rust/chains/tw_solana/src/transaction/legacy.rs

@@ -3,6 +3,7 @@
 // Copyright © 2017 Trust Wallet.
 
 use crate::address::SolanaAddress;
+use crate::modules::insert_instruction::InsertInstruction;
 use crate::transaction::{short_vec, CompiledInstruction, MessageHeader, Signature};
 use serde::{Deserialize, Serialize};
 use tw_hash::{as_byte_sequence, H256};
@@ -28,6 +29,20 @@ pub struct Message {
     pub instructions: Vec<CompiledInstruction>,
 }
 
+impl InsertInstruction for Message {
+    fn account_keys_mut(&mut self) -> &mut Vec<SolanaAddress> {
+        &mut self.account_keys
+    }
+
+    fn message_header_mut(&mut self) -> &mut MessageHeader {
+        &mut self.header
+    }
+
+    fn instructions_mut(&mut self) -> &mut Vec<CompiledInstruction> {
+        &mut self.instructions
+    }
+}
+
 #[derive(Debug, PartialEq, Default, Eq, Clone, Serialize, Deserialize)]
 pub struct Transaction {
     /// A set of signatures of a serialized [`Message`], signed by the first