Skip to main content

gmsol_solana_utils/
instruction_group.rs

1use std::{
2    borrow::Cow,
3    collections::{BTreeMap, BTreeSet, HashSet},
4    ops::Deref,
5};
6
7use smallvec::SmallVec;
8use solana_sdk::{
9    hash::Hash,
10    instruction::Instruction,
11    message::{v0, VersionedMessage},
12    pubkey::Pubkey,
13    signature::NullSigner,
14    signer::Signer,
15    transaction::VersionedTransaction,
16};
17
18use crate::{
19    address_lookup_table::AddressLookupTables, compute_budget::ComputeBudget,
20    signer::BoxClonableSigner, transaction_group::TransactionGroupOptions,
21};
22
23const ATOMIC_SIZE: usize = 3;
24const PARALLEL_SIZE: usize = 2;
25
26/// A trait representing types that can be converted into [`AtomicGroup`]s.
27pub trait IntoAtomicGroup {
28    /// Hint.
29    type Hint;
30
31    /// Convert into an [`AtomicGroup`].
32    fn into_atomic_group(self, hint: &Self::Hint) -> crate::Result<AtomicGroup>;
33
34    /// Convert into an [`AtomicGroup`] with RPC client.
35    #[cfg(client_traits)]
36    fn into_atomic_group_with_rpc_client(
37        self,
38        client: &impl crate::client_traits::RpcClient,
39    ) -> impl std::future::Future<Output = crate::Result<AtomicGroup>>
40    where
41        Self: Sized,
42        Self::Hint: crate::client_traits::FromRpcClientWith<Self>,
43    {
44        use crate::client_traits::FromRpcClientWith;
45
46        async move {
47            let hint = Self::Hint::from_rpc_client_with(&self, client).await?;
48            self.into_atomic_group(&hint)
49        }
50    }
51}
52
53/// Options for getting instructions.
54#[derive(Debug, Clone, Default)]
55pub struct GetInstructionsOptions {
56    /// Options for compute budget.
57    pub compute_budget: ComputeBudgetOptions,
58    /// If set, a memo will be included in the final transaction.
59    pub memo: Option<String>,
60    /// If set, the signer list for the memo instruction will be replaced.
61    pub memo_signers: Option<Vec<Pubkey>>,
62    /// Extra compute units.
63    pub extra_compute_units: u32,
64}
65
66/// Options for compute budget.
67#[derive(Debug, Clone, Default)]
68pub struct ComputeBudgetOptions {
69    /// Without compute budget instruction.
70    pub without_compute_budget: bool,
71    /// Compute unit price in micro lamports.
72    pub compute_unit_price_micro_lamports: Option<u64>,
73    /// Compute unit min priority lamports.
74    pub compute_unit_min_priority_lamports: Option<u64>,
75}
76
77/// Options type for [`AtomicGroup`].
78#[derive(Debug, Clone)]
79pub struct AtomicGroupOptions {
80    /// Indicates whether the group is mergeable.
81    pub is_mergeable: bool,
82}
83
84impl Default for AtomicGroupOptions {
85    fn default() -> Self {
86        Self { is_mergeable: true }
87    }
88}
89
90/// A group of instructions that are expected to be executed in the same transaction.
91#[derive(Debug, Clone)]
92pub struct AtomicGroup {
93    payer: Pubkey,
94    signers: BTreeMap<Pubkey, NullSigner>,
95    owned_signers: BTreeMap<Pubkey, BoxClonableSigner<'static>>,
96    instructions: SmallVec<[Instruction; ATOMIC_SIZE]>,
97    compute_budget: ComputeBudget,
98    options: AtomicGroupOptions,
99}
100
101impl AtomicGroup {
102    /// Returns whether the atomic group is mergeable.
103    pub fn is_mergeable(&self) -> bool {
104        self.options().is_mergeable
105    }
106
107    /// Returns the options of the group.
108    pub fn options(&self) -> &AtomicGroupOptions {
109        &self.options
110    }
111
112    /// Create from an iterator of instructions and options.
113    pub fn with_instructions_and_options(
114        payer: &Pubkey,
115        instructions: impl IntoIterator<Item = Instruction>,
116        options: AtomicGroupOptions,
117    ) -> Self {
118        Self {
119            payer: *payer,
120            signers: BTreeMap::from([(*payer, NullSigner::new(payer))]),
121            owned_signers: Default::default(),
122            instructions: SmallVec::from_iter(instructions),
123            compute_budget: Default::default(),
124            options,
125        }
126    }
127
128    /// Create from an iterator of instructions.
129    pub fn with_instructions(
130        payer: &Pubkey,
131        instructions: impl IntoIterator<Item = Instruction>,
132    ) -> Self {
133        Self::with_instructions_and_options(payer, instructions, Default::default())
134    }
135
136    /// Create a new empty group.
137    pub fn new(payer: &Pubkey) -> Self {
138        Self::with_instructions(payer, None)
139    }
140
141    /// Add an instruction.
142    pub fn add(&mut self, instruction: Instruction) -> &mut Self {
143        self.instructions.push(instruction);
144        self
145    }
146
147    /// Add a signer.
148    pub fn add_signer(&mut self, signer: &Pubkey) -> &mut Self {
149        self.signers.insert(*signer, NullSigner::new(signer));
150        self
151    }
152
153    /// Add an owned signer.
154    pub fn add_owned_signer(&mut self, signer: impl Signer + Clone + 'static) -> &mut Self {
155        self.owned_signers
156            .insert(signer.pubkey(), BoxClonableSigner::new(signer));
157        self
158    }
159
160    /// Get compute budget.
161    pub fn compute_budget(&self) -> &ComputeBudget {
162        &self.compute_budget
163    }
164
165    /// Get mutable reference to the compute budget.
166    pub fn compute_budget_mut(&mut self) -> &mut ComputeBudget {
167        &mut self.compute_budget
168    }
169
170    /// Returns the pubkey of the payer.
171    pub fn payer(&self) -> &Pubkey {
172        &self.payer
173    }
174
175    /// Returns signers that need to be provided externally including the payer.
176    pub fn external_signers(&self) -> impl Iterator<Item = &Pubkey> + '_ {
177        self.signers.keys()
178    }
179
180    fn compute_budget_instructions(
181        &self,
182        compute_unit_price_micro_lamports: Option<u64>,
183        compute_unit_min_priority_lamports: Option<u64>,
184        extra_compute_units: u32,
185    ) -> Vec<Instruction> {
186        self.compute_budget
187            .compute_budget_instructions_with_extra_units(
188                compute_unit_price_micro_lamports,
189                compute_unit_min_priority_lamports,
190                extra_compute_units,
191            )
192    }
193
194    /// Returns instructions.
195    pub fn instructions_with_options(
196        &self,
197        options: GetInstructionsOptions,
198    ) -> impl Iterator<Item = Cow<'_, Instruction>> {
199        let compute_budget_instructions = if options.compute_budget.without_compute_budget {
200            Vec::default()
201        } else {
202            self.compute_budget_instructions(
203                options.compute_budget.compute_unit_price_micro_lamports,
204                options.compute_budget.compute_unit_min_priority_lamports,
205                options.extra_compute_units,
206            )
207        };
208        let memo_signers = match options.memo_signers.as_ref() {
209            Some(signers) => signers.iter().collect(),
210            None => Vec::from([&self.payer]),
211        };
212        let memo_instruction = options
213            .memo
214            .as_ref()
215            .map(|s| spl_memo::build_memo(s.as_bytes(), &memo_signers));
216        compute_budget_instructions
217            .into_iter()
218            .chain(memo_instruction)
219            .map(Cow::Owned)
220            .chain(self.instructions.iter().map(Cow::Borrowed))
221    }
222
223    /// Estimates the transaciton size.
224    pub fn transaction_size(
225        &self,
226        is_versioned_transaction: bool,
227        luts: Option<&AddressLookupTables>,
228        options: GetInstructionsOptions,
229    ) -> usize {
230        crate::utils::transaction_size_with_luts(
231            self.payer,
232            &self.instructions_with_options(options).collect::<Vec<_>>(),
233            is_versioned_transaction,
234            luts,
235        )
236    }
237
238    /// Estimates the transaction size after merge.
239    pub fn transaction_size_after_merge(
240        &self,
241        other: &Self,
242        is_versioned_transaction: bool,
243        luts: Option<&AddressLookupTables>,
244        options: GetInstructionsOptions,
245    ) -> usize {
246        crate::utils::transaction_size_with_luts(
247            self.payer,
248            &self
249                .instructions_with_options(options)
250                .chain(other.instructions_with_options(GetInstructionsOptions {
251                    compute_budget: ComputeBudgetOptions {
252                        without_compute_budget: true,
253                        ..Default::default()
254                    },
255                    ..Default::default()
256                }))
257                .collect::<Vec<_>>(),
258            is_versioned_transaction,
259            luts,
260        )
261    }
262
263    /// Merge two [`AtomicGroup`]s.
264    ///
265    /// # Note
266    /// - Merging does not change the payer of the current [`AtomicGroup`].
267    pub fn merge(&mut self, mut other: Self) -> &mut Self {
268        self.signers.append(&mut other.signers);
269        self.owned_signers.append(&mut other.owned_signers);
270        self.instructions.extend(other.instructions);
271        self.compute_budget += other.compute_budget;
272        self
273    }
274
275    fn v0_message_with_blockhash_and_options(
276        &self,
277        recent_blockhash: Hash,
278        options: GetInstructionsOptions,
279        luts: Option<&AddressLookupTables>,
280    ) -> crate::Result<v0::Message> {
281        let instructions = self
282            .instructions_with_options(options)
283            .map(|ix| (*ix).clone())
284            .collect::<Vec<_>>();
285        let luts = luts
286            .map(|t| t.accounts().collect::<Vec<_>>())
287            .unwrap_or_default();
288        Ok(v0::Message::try_compile(
289            self.payer(),
290            &instructions,
291            &luts,
292            recent_blockhash,
293        )?)
294    }
295
296    /// Create versioned message with the given blockhash and options.
297    pub fn message_with_blockhash_and_options(
298        &self,
299        recent_blockhash: Hash,
300        options: GetInstructionsOptions,
301        luts: Option<&AddressLookupTables>,
302    ) -> crate::Result<VersionedMessage> {
303        Ok(VersionedMessage::V0(
304            self.v0_message_with_blockhash_and_options(recent_blockhash, options, luts)?,
305        ))
306    }
307
308    /// Create partially signed transaction with the given blockhash and options.
309    pub fn partially_signed_transaction_with_blockhash_and_options(
310        &self,
311        recent_blockhash: Hash,
312        options: GetInstructionsOptions,
313        luts: Option<&AddressLookupTables>,
314        mut before_sign: impl FnMut(&VersionedMessage) -> crate::Result<()>,
315    ) -> crate::Result<VersionedTransaction> {
316        let mut memo_signers = vec![];
317        if let Some(signers) = options.memo_signers.as_ref() {
318            let signers: BTreeSet<_> = signers.iter().collect();
319            for signer in signers {
320                if !self.signers.contains_key(signer) && !self.owned_signers.contains_key(signer) {
321                    memo_signers.push(NullSigner::new(signer));
322                }
323            }
324        }
325        let message = self.message_with_blockhash_and_options(recent_blockhash, options, luts)?;
326        (before_sign)(&message)?;
327        let signers = self
328            .signers
329            .values()
330            .chain(memo_signers.iter())
331            .map(|s| s as &dyn Signer)
332            .chain(self.owned_signers.values().map(|s| s as &dyn Signer))
333            .collect::<Vec<_>>();
334        Ok(VersionedTransaction::try_new(message, &signers)?)
335    }
336
337    /// Estimates the execution fee of the result transaction.
338    pub fn estimate_execution_fee(
339        &self,
340        compute_unit_price_micro_lamports: Option<u64>,
341        compute_unit_min_priority_lamports: Option<u64>,
342    ) -> u64 {
343        self.estimate_execution_fee_with_extra_units(
344            compute_unit_price_micro_lamports,
345            compute_unit_min_priority_lamports,
346            0,
347        )
348    }
349
350    /// Estimates the execution fee of the result transaction with extra compute units.
351    pub fn estimate_execution_fee_with_extra_units(
352        &self,
353        compute_unit_price_micro_lamports: Option<u64>,
354        compute_unit_min_priority_lamports: Option<u64>,
355        extra_compute_units: u32,
356    ) -> u64 {
357        let ixs = self
358            .instructions_with_options(GetInstructionsOptions {
359                compute_budget: ComputeBudgetOptions {
360                    without_compute_budget: true,
361                    ..Default::default()
362                },
363                ..Default::default()
364            })
365            .collect::<Vec<_>>();
366
367        let num_signers = ixs
368            .iter()
369            .flat_map(|ix| ix.accounts.iter())
370            .filter(|meta| meta.is_signer)
371            .map(|meta| &meta.pubkey)
372            .collect::<HashSet<_>>()
373            .len() as u64;
374        num_signers * 5_000
375            + self.compute_budget.fee_with_extra_units(
376                compute_unit_price_micro_lamports,
377                compute_unit_min_priority_lamports,
378                extra_compute_units,
379            )
380    }
381}
382
383impl Extend<Instruction> for AtomicGroup {
384    fn extend<T: IntoIterator<Item = Instruction>>(&mut self, iter: T) {
385        self.instructions.extend(iter);
386    }
387}
388
389impl Deref for AtomicGroup {
390    type Target = [Instruction];
391
392    fn deref(&self) -> &Self::Target {
393        self.instructions.deref()
394    }
395}
396
397/// The options type for [`ParallelGroup`].
398#[derive(Debug, Clone)]
399pub struct ParallelGroupOptions {
400    /// Indicates whether the [`ParallelGroup`] is mergeable.
401    pub is_mergeable: bool,
402}
403
404impl Default for ParallelGroupOptions {
405    fn default() -> Self {
406        Self { is_mergeable: true }
407    }
408}
409
410/// A group of atomic instructions that can be executed in parallel.
411#[derive(Debug, Clone, Default)]
412pub struct ParallelGroup {
413    groups: SmallVec<[AtomicGroup; PARALLEL_SIZE]>,
414    options: ParallelGroupOptions,
415}
416
417impl ParallelGroup {
418    /// Create a new [`ParallelGroup`] with the given options.
419    pub fn with_options(
420        groups: impl IntoIterator<Item = AtomicGroup>,
421        options: ParallelGroupOptions,
422    ) -> Self {
423        Self {
424            groups: FromIterator::from_iter(groups),
425            options,
426        }
427    }
428
429    /// Returns the options.
430    pub fn options(&self) -> &ParallelGroupOptions {
431        &self.options
432    }
433
434    /// Returns whether the group is mergeable.
435    pub fn is_mergeable(&self) -> bool {
436        self.options().is_mergeable
437    }
438
439    /// Set whether the group is mergeable.
440    pub fn set_is_mergeable(&mut self, is_mergeable: bool) -> &mut Self {
441        self.options.is_mergeable = is_mergeable;
442        self
443    }
444
445    /// Add an [`AtomicGroup`].
446    pub fn add(&mut self, group: AtomicGroup) -> &mut Self {
447        self.groups.push(group);
448        self
449    }
450
451    pub(crate) fn optimize(
452        &mut self,
453        options: &TransactionGroupOptions,
454        luts: &AddressLookupTables,
455        allow_payer_change: bool,
456    ) -> &mut Self {
457        if options.optimize(&mut self.groups, luts, allow_payer_change) {
458            self.groups = self
459                .groups
460                .drain(..)
461                .filter(|group| !group.is_empty())
462                .collect();
463        }
464        self
465    }
466
467    pub(crate) fn single(&self) -> Option<&AtomicGroup> {
468        if self.groups.len() == 1 {
469            Some(&self.groups[0])
470        } else {
471            None
472        }
473    }
474
475    pub(crate) fn single_mut(&mut self) -> Option<&mut AtomicGroup> {
476        if self.groups.len() == 1 {
477            Some(&mut self.groups[0])
478        } else {
479            None
480        }
481    }
482
483    pub(crate) fn into_single(mut self) -> Option<AtomicGroup> {
484        if self.groups.len() == 1 {
485            Some(self.groups.remove(0))
486        } else {
487            None
488        }
489    }
490
491    /// Returns the total number of transactions.
492    pub fn len(&self) -> usize {
493        self.groups.len()
494    }
495
496    /// Returns whether the group is empty.
497    pub fn is_empty(&self) -> bool {
498        self.groups.is_empty()
499    }
500
501    /// Estiamtes the execution fee of the result transactions
502    pub fn estimate_execution_fee(
503        &self,
504        compute_unit_price_micro_lamports: Option<u64>,
505        compute_unit_min_priority_lamports: Option<u64>,
506    ) -> u64 {
507        self.estimate_execution_fee_with_extra_units(
508            compute_unit_price_micro_lamports,
509            compute_unit_min_priority_lamports,
510            0,
511        )
512    }
513
514    /// Estiamtes the execution fee of the result transactions with extra units.
515    pub fn estimate_execution_fee_with_extra_units(
516        &self,
517        compute_unit_price_micro_lamports: Option<u64>,
518        compute_unit_min_priority_lamports: Option<u64>,
519        extra_compute_units: u32,
520    ) -> u64 {
521        self.groups
522            .iter()
523            .map(|ag| {
524                ag.estimate_execution_fee_with_extra_units(
525                    compute_unit_price_micro_lamports,
526                    compute_unit_min_priority_lamports,
527                    extra_compute_units,
528                )
529            })
530            .sum()
531    }
532}
533
534impl From<AtomicGroup> for ParallelGroup {
535    fn from(value: AtomicGroup) -> Self {
536        let mut this = Self::default();
537        this.add(value);
538        this
539    }
540}
541
542impl FromIterator<AtomicGroup> for ParallelGroup {
543    fn from_iter<T: IntoIterator<Item = AtomicGroup>>(iter: T) -> Self {
544        Self::with_options(iter, Default::default())
545    }
546}
547
548impl Deref for ParallelGroup {
549    type Target = [AtomicGroup];
550
551    fn deref(&self) -> &Self::Target {
552        self.groups.deref()
553    }
554}