feat(sdk-rs): add VaultTransactionMessageExt

sdk/rs/src/client.rs

@@ -313,22 +313,25 @@ pub fn spending_limit_use(
 /// ```
 /// use squads_multisig::anchor_lang::AnchorSerialize;
 /// use squads_multisig::solana_program::pubkey::Pubkey;
-/// use squads_multisig::solana_program::system_program;
+/// use squads_multisig::solana_program::{system_instruction, system_program};
 /// use squads_multisig::client::{
 ///     VaultTransactionCreateAccounts,
 ///     VaultTransactionCreateArgs,
 ///     vault_transaction_create,
 /// };
+/// use squads_multisig::pda::get_vault_pda;
+/// use squads_multisig::vault_transaction_message::VaultTransactionMessageExt;
 /// use squads_multisig_program::TransactionMessage;
 ///
-/// let message = TransactionMessage {
-///     num_signers: 1,
-///     num_writable_signers: 1,
-///     num_writable_non_signers: 2,
-///     account_keys: vec![].into(),
-///     instructions: vec![].into(),
-///     address_table_lookups: vec![].into(),
-/// }.try_to_vec().unwrap();
+/// // Default vault (index 0).
+/// let vault_pda = get_vault_pda(&Pubkey::new_unique(), 0, None).0;
+///
+/// // Create a vault transaction message that includes 1 instruction - SOL transfer from the default vault.
+/// let message = TransactionMessage::try_compile(
+///   &vault_pda,
+///   &[system_instruction::transfer(&vault_pda, &Pubkey::new_unique(), 1_000_000)],
+///   &[]
+/// ).unwrap().try_to_vec().unwrap();
 ///
 /// let ix = vault_transaction_create(
 ///     VaultTransactionCreateAccounts {

sdk/rs/src/lib.rs

@@ -7,6 +7,7 @@ pub use squads_multisig_program::anchor_lang::solana_program;
 
 pub mod client;
 pub mod pda;
+pub mod vault_transaction_message;
 
 pub mod error {
     use thiserror::Error;

sdk/rs/src/vault_transaction_message/compiled_keys.rs

@@ -0,0 +1,176 @@
+use std::collections::BTreeMap;
+
+use crate::solana_program::address_lookup_table_account::AddressLookupTableAccount;
+use crate::solana_program::instruction::Instruction;
+use crate::solana_program::message::v0::{LoadedAddresses, MessageAddressTableLookup};
+use crate::solana_program::message::{CompileError, MessageHeader};
+
+use crate::solana_program::pubkey::Pubkey;
+
+/// A helper struct to collect pubkeys compiled for a set of instructions
+///
+/// NOTE: The only difference between this and the original implementation from `solana_program` is that we don't mark the instruction programIds as invoked.
+// /// It makes sense to do because the instructions will be called via CPI, so the programIds can come from Address Lookup Tables.
+// /// This allows to compress the message size and avoid hitting the tx size limit during `vault_transaction_create` instruction calls.
+#[derive(Default, Debug, Clone, PartialEq, Eq)]
+pub(crate) struct CompiledKeys {
+    payer: Option<Pubkey>,
+    key_meta_map: BTreeMap<Pubkey, CompiledKeyMeta>,
+}
+
+#[derive(Default, Debug, Clone, PartialEq, Eq)]
+struct CompiledKeyMeta {
+    is_signer: bool,
+    is_writable: bool,
+    is_invoked: bool,
+}
+
+impl CompiledKeys {
+    /// Compiles the pubkeys referenced by a list of instructions and organizes by
+    /// signer/non-signer and writable/readonly.
+    pub(crate) fn compile(instructions: &[Instruction], payer: Option<Pubkey>) -> Self {
+        let mut key_meta_map = BTreeMap::<Pubkey, CompiledKeyMeta>::new();
+        for ix in instructions {
+            let mut meta = key_meta_map.entry(ix.program_id).or_default();
+            // NOTE: This is the only difference from the original.
+            // meta.is_invoked = true;
+            meta.is_invoked = false;
+            for account_meta in &ix.accounts {
+                let meta = key_meta_map.entry(account_meta.pubkey).or_default();
+                meta.is_signer |= account_meta.is_signer;
+                meta.is_writable |= account_meta.is_writable;
+            }
+        }
+        if let Some(payer) = &payer {
+            let mut meta = key_meta_map.entry(*payer).or_default();
+            meta.is_signer = true;
+            meta.is_writable = true;
+        }
+        Self {
+            payer,
+            key_meta_map,
+        }
+    }
+
+    pub(crate) fn try_into_message_components(
+        self,
+    ) -> Result<(MessageHeader, Vec<Pubkey>), CompileError> {
+        let try_into_u8 = |num: usize| -> Result<u8, CompileError> {
+            u8::try_from(num).map_err(|_| CompileError::AccountIndexOverflow)
+        };
+
+        let Self {
+            payer,
+            mut key_meta_map,
+        } = self;
+
+        if let Some(payer) = &payer {
+            key_meta_map.remove_entry(payer);
+        }
+
+        let writable_signer_keys: Vec<Pubkey> = payer
+            .into_iter()
+            .chain(
+                key_meta_map
+                    .iter()
+                    .filter_map(|(key, meta)| (meta.is_signer && meta.is_writable).then(|| *key)),
+            )
+            .collect();
+        let readonly_signer_keys: Vec<Pubkey> = key_meta_map
+            .iter()
+            .filter_map(|(key, meta)| (meta.is_signer && !meta.is_writable).then(|| *key))
+            .collect();
+        let writable_non_signer_keys: Vec<Pubkey> = key_meta_map
+            .iter()
+            .filter_map(|(key, meta)| (!meta.is_signer && meta.is_writable).then(|| *key))
+            .collect();
+        let readonly_non_signer_keys: Vec<Pubkey> = key_meta_map
+            .iter()
+            .filter_map(|(key, meta)| (!meta.is_signer && !meta.is_writable).then(|| *key))
+            .collect();
+
+        let signers_len = writable_signer_keys
+            .len()
+            .saturating_add(readonly_signer_keys.len());
+
+        let header = MessageHeader {
+            num_required_signatures: try_into_u8(signers_len)?,
+            num_readonly_signed_accounts: try_into_u8(readonly_signer_keys.len())?,
+            num_readonly_unsigned_accounts: try_into_u8(readonly_non_signer_keys.len())?,
+        };
+
+        let static_account_keys = std::iter::empty()
+            .chain(writable_signer_keys)
+            .chain(readonly_signer_keys)
+            .chain(writable_non_signer_keys)
+            .chain(readonly_non_signer_keys)
+            .collect();
+
+        Ok((header, static_account_keys))
+    }
+
+    #[cfg(not(target_os = "solana"))]
+    pub(crate) fn try_extract_table_lookup(
+        &mut self,
+        lookup_table_account: &AddressLookupTableAccount,
+    ) -> Result<Option<(MessageAddressTableLookup, LoadedAddresses)>, CompileError> {
+        let (writable_indexes, drained_writable_keys) = self
+            .try_drain_keys_found_in_lookup_table(&lookup_table_account.addresses, |meta| {
+                !meta.is_signer && !meta.is_invoked && meta.is_writable
+            })?;
+        let (readonly_indexes, drained_readonly_keys) = self
+            .try_drain_keys_found_in_lookup_table(&lookup_table_account.addresses, |meta| {
+                !meta.is_signer && !meta.is_invoked && !meta.is_writable
+            })?;
+
+        // Don't extract lookup if no keys were found
+        if writable_indexes.is_empty() && readonly_indexes.is_empty() {
+            return Ok(None);
+        }
+
+        Ok(Some((
+            MessageAddressTableLookup {
+                account_key: lookup_table_account.key,
+                writable_indexes,
+                readonly_indexes,
+            },
+            LoadedAddresses {
+                writable: drained_writable_keys,
+                readonly: drained_readonly_keys,
+            },
+        )))
+    }
+
+    #[cfg(not(target_os = "solana"))]
+    fn try_drain_keys_found_in_lookup_table(
+        &mut self,
+        lookup_table_addresses: &[Pubkey],
+        key_meta_filter: impl Fn(&CompiledKeyMeta) -> bool,
+    ) -> Result<(Vec<u8>, Vec<Pubkey>), CompileError> {
+        let mut lookup_table_indexes = Vec::new();
+        let mut drained_keys = Vec::new();
+
+        for search_key in self
+            .key_meta_map
+            .iter()
+            .filter_map(|(key, meta)| key_meta_filter(meta).then(|| key))
+        {
+            for (key_index, key) in lookup_table_addresses.iter().enumerate() {
+                if key == search_key {
+                    let lookup_table_index = u8::try_from(key_index)
+                        .map_err(|_| CompileError::AddressTableLookupIndexOverflow)?;
+
+                    lookup_table_indexes.push(lookup_table_index);
+                    drained_keys.push(*search_key);
+                    break;
+                }
+            }
+        }
+
+        for key in &drained_keys {
+            self.key_meta_map.remove_entry(key);
+        }
+
+        Ok((lookup_table_indexes, drained_keys))
+    }
+}

sdk/rs/src/vault_transaction_message/mod.rs

@@ -0,0 +1,72 @@
+use squads_multisig_program::{CompiledInstruction, MessageAddressTableLookup, TransactionMessage};
+
+use crate::solana_program::address_lookup_table_account::AddressLookupTableAccount;
+use crate::solana_program::instruction::Instruction;
+use crate::solana_program::message::{AccountKeys, CompileError};
+use crate::solana_program::pubkey::Pubkey;
+use crate::vault_transaction_message::compiled_keys::CompiledKeys;
+
+mod compiled_keys;
+
+pub trait VaultTransactionMessageExt {
+    /// This implementation is mostly a copy-paste from `solana_program::message::v0::Message::try_compile()`,
+    /// but it constructs a `TransactionMessage` meant to be passed to `vault_transaction_create`.
+    fn try_compile(
+        vault_key: &Pubkey,
+        instructions: &[Instruction],
+        address_lookup_table_accounts: &[AddressLookupTableAccount],
+    ) -> Result<TransactionMessage, CompileError> {
+        let mut compiled_keys = CompiledKeys::compile(instructions, Some(*vault_key));
+
+        let mut address_table_lookups = Vec::with_capacity(address_lookup_table_accounts.len());
+        let mut loaded_addresses_list = Vec::with_capacity(address_lookup_table_accounts.len());
+        for lookup_table_account in address_lookup_table_accounts {
+            if let Some((lookup, loaded_addresses)) =
+                compiled_keys.try_extract_table_lookup(lookup_table_account)?
+            {
+                address_table_lookups.push(lookup);
+                loaded_addresses_list.push(loaded_addresses);
+            }
+        }
+
+        let (header, static_keys) = compiled_keys.try_into_message_components()?;
+        let dynamic_keys = loaded_addresses_list.into_iter().collect();
+        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
+        let instructions = account_keys.try_compile_instructions(instructions)?;
+
+        let num_static_keys: u8 = static_keys
+            .len()
+            .try_into()
+            .map_err(|_| CompileError::AccountIndexOverflow)?;
+
+        Ok(TransactionMessage {
+            num_signers: header.num_required_signatures,
+            num_writable_signers: header.num_required_signatures
+                - header.num_readonly_signed_accounts,
+            num_writable_non_signers: num_static_keys
+                - header.num_required_signatures
+                - header.num_readonly_unsigned_accounts,
+            account_keys: static_keys.into(),
+            instructions: instructions
+                .into_iter()
+                .map(|ix| CompiledInstruction {
+                    program_id_index: ix.program_id_index,
+                    account_indexes: ix.accounts.into(),
+                    data: ix.data.into(),
+                })
+                .collect::<Vec<CompiledInstruction>>()
+                .into(),
+            address_table_lookups: address_table_lookups
+                .into_iter()
+                .map(|lookup| MessageAddressTableLookup {
+                    account_key: lookup.account_key,
+                    writable_indexes: lookup.writable_indexes.into(),
+                    readonly_indexes: lookup.readonly_indexes.into(),
+                })
+                .collect::<Vec<MessageAddressTableLookup>>()
+                .into(),
+        })
+    }
+}
+
+impl VaultTransactionMessageExt for TransactionMessage {}