gmsol_solana_utils/
instruction_group.rs

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