gmsol_solana_utils/
instruction_group.rs

1use std::{borrow::Cow, collections::BTreeMap, ops::Deref};
2
3use smallvec::SmallVec;
4use solana_sdk::{
5    hash::Hash,
6    instruction::Instruction,
7    message::{v0, VersionedMessage},
8    pubkey::Pubkey,
9    signature::NullSigner,
10    signer::Signer,
11    transaction::VersionedTransaction,
12};
13
14use crate::{
15    address_lookup_table::AddressLookupTables, compute_budget::ComputeBudget,
16    signer::BoxClonableSigner, transaction_group::TransactionGroupOptions,
17};
18
19const ATOMIC_SIZE: usize = 3;
20const PARALLEL_SIZE: usize = 2;
21
22/// A trait representing types that can be converted into [`AtomicGroup`]s.
23pub trait IntoAtomicGroup {
24    /// Hint.
25    type Hint;
26
27    /// Convert into [`AtomicGroup`]s.
28    fn into_atomic_group(self, hint: &Self::Hint) -> crate::Result<AtomicGroup>;
29}
30
31/// Options for getting instructions.
32#[derive(Debug, Clone, Default)]
33pub struct GetInstructionsOptions {
34    /// Without compute budget instruction.
35    pub without_compute_budget: bool,
36    /// Compute unit price in micro lamports.
37    pub compute_unit_price_micro_lamports: Option<u64>,
38    /// If set, a memo will be included in the final transaction.
39    pub memo: Option<String>,
40}
41
42/// A group of instructions that are expected to be executed in the same transaction.
43#[derive(Debug, Clone)]
44pub struct AtomicGroup {
45    payer: Pubkey,
46    signers: BTreeMap<Pubkey, NullSigner>,
47    owned_signers: BTreeMap<Pubkey, BoxClonableSigner<'static>>,
48    instructions: SmallVec<[Instruction; ATOMIC_SIZE]>,
49    compute_budget: ComputeBudget,
50}
51
52impl AtomicGroup {
53    /// Create from an iterator of instructions.
54    pub fn with_instructions(
55        payer: &Pubkey,
56        instructions: impl IntoIterator<Item = Instruction>,
57    ) -> Self {
58        Self {
59            payer: *payer,
60            signers: BTreeMap::from([(*payer, NullSigner::new(payer))]),
61            owned_signers: Default::default(),
62            instructions: SmallVec::from_iter(instructions),
63            compute_budget: Default::default(),
64        }
65    }
66
67    /// Create a new empty group.
68    pub fn new(payer: &Pubkey) -> Self {
69        Self::with_instructions(payer, None)
70    }
71
72    /// Add an instruction.
73    pub fn add(&mut self, instruction: Instruction) -> &mut Self {
74        self.instructions.push(instruction);
75        self
76    }
77
78    /// Add a signer.
79    pub fn add_signer(&mut self, signer: &Pubkey) -> &mut Self {
80        self.signers.insert(*signer, NullSigner::new(signer));
81        self
82    }
83
84    /// Add an owned signer.
85    pub fn add_owned_signer(&mut self, signer: impl Signer + Clone + 'static) -> &mut Self {
86        self.owned_signers
87            .insert(signer.pubkey(), BoxClonableSigner::new(signer));
88        self
89    }
90
91    /// Get compute budget.
92    pub fn compute_budget(&self) -> &ComputeBudget {
93        &self.compute_budget
94    }
95
96    /// Get mutable reference to the compute budget.
97    pub fn compute_budget_mut(&mut self) -> &mut ComputeBudget {
98        &mut self.compute_budget
99    }
100
101    /// Returns the pubkey of the payer.
102    pub fn payer(&self) -> &Pubkey {
103        &self.payer
104    }
105
106    /// Returns signers that need to be provided externally including the payer.
107    pub fn external_signers(&self) -> impl Iterator<Item = &Pubkey> + '_ {
108        self.signers.keys()
109    }
110
111    fn compute_budget_instructions(
112        &self,
113        compute_unit_price_micro_lamports: Option<u64>,
114    ) -> Vec<Instruction> {
115        self.compute_budget
116            .compute_budget_instructions(compute_unit_price_micro_lamports)
117    }
118
119    /// Returns instructions.
120    pub fn instructions_with_options(
121        &self,
122        options: GetInstructionsOptions,
123    ) -> impl Iterator<Item = Cow<'_, Instruction>> {
124        let compute_budget_instructions = if options.without_compute_budget {
125            Vec::default()
126        } else {
127            self.compute_budget_instructions(options.compute_unit_price_micro_lamports)
128        };
129        let memo_instruction = options
130            .memo
131            .as_ref()
132            .map(|s| spl_memo::build_memo(s.as_bytes(), &[&self.payer]));
133        compute_budget_instructions
134            .into_iter()
135            .chain(memo_instruction)
136            .map(Cow::Owned)
137            .chain(self.instructions.iter().map(Cow::Borrowed))
138    }
139
140    /// Estimates the transaciton size.
141    pub fn transaction_size(
142        &self,
143        is_versioned_transaction: bool,
144        luts: Option<&AddressLookupTables>,
145        options: GetInstructionsOptions,
146    ) -> usize {
147        let addresses = luts.as_ref().map(|luts| luts.addresses());
148        crate::utils::transaction_size(
149            &self.instructions_with_options(options).collect::<Vec<_>>(),
150            is_versioned_transaction,
151            addresses.as_ref(),
152            luts.as_ref().map(|luts| luts.len()).unwrap_or_default(),
153        )
154    }
155
156    /// Estimates the transaction size after merge.
157    pub fn transaction_size_after_merge(
158        &self,
159        other: &Self,
160        is_versioned_transaction: bool,
161        luts: Option<&AddressLookupTables>,
162        options: GetInstructionsOptions,
163    ) -> usize {
164        let addresses = luts.as_ref().map(|luts| luts.addresses());
165        crate::utils::transaction_size(
166            &self
167                .instructions_with_options(options)
168                .chain(other.instructions_with_options(GetInstructionsOptions {
169                    without_compute_budget: true,
170                    ..Default::default()
171                }))
172                .collect::<Vec<_>>(),
173            is_versioned_transaction,
174            addresses.as_ref(),
175            luts.as_ref().map(|luts| luts.len()).unwrap_or_default(),
176        )
177    }
178
179    /// Merge two [`AtomicGroup`]s.
180    ///
181    /// # Note
182    /// - Merging does not change the payer of the current [`AtomicGroup`].
183    pub fn merge(&mut self, mut other: Self) -> &mut Self {
184        self.signers.append(&mut other.signers);
185        self.owned_signers.append(&mut other.owned_signers);
186        self.instructions.extend(other.instructions);
187        self.compute_budget += other.compute_budget;
188        self
189    }
190
191    fn v0_message_with_blockhash_and_options(
192        &self,
193        recent_blockhash: Hash,
194        options: GetInstructionsOptions,
195        luts: Option<&AddressLookupTables>,
196    ) -> crate::Result<v0::Message> {
197        let instructions = self
198            .instructions_with_options(options)
199            .map(|ix| (*ix).clone())
200            .collect::<Vec<_>>();
201        let luts = luts
202            .map(|t| t.accounts().collect::<Vec<_>>())
203            .unwrap_or_default();
204        Ok(v0::Message::try_compile(
205            self.payer(),
206            &instructions,
207            &luts,
208            recent_blockhash,
209        )?)
210    }
211
212    /// Create versioned message with the given blockhash and options.
213    pub fn message_with_blockhash_and_options(
214        &self,
215        recent_blockhash: Hash,
216        options: GetInstructionsOptions,
217        luts: Option<&AddressLookupTables>,
218    ) -> crate::Result<VersionedMessage> {
219        Ok(VersionedMessage::V0(
220            self.v0_message_with_blockhash_and_options(recent_blockhash, options, luts)?,
221        ))
222    }
223
224    /// Create partially signed transaction with the given blockhash and options.
225    pub fn partially_signed_transaction_with_blockhash_and_options(
226        &self,
227        recent_blockhash: Hash,
228        options: GetInstructionsOptions,
229        luts: Option<&AddressLookupTables>,
230    ) -> crate::Result<VersionedTransaction> {
231        let message = self.message_with_blockhash_and_options(recent_blockhash, options, luts)?;
232        let signers = self
233            .signers
234            .values()
235            .map(|s| s as &dyn Signer)
236            .chain(self.owned_signers.values().map(|s| s as &dyn Signer))
237            .collect::<Vec<_>>();
238        Ok(VersionedTransaction::try_new(message, &signers)?)
239    }
240}
241
242impl Extend<Instruction> for AtomicGroup {
243    fn extend<T: IntoIterator<Item = Instruction>>(&mut self, iter: T) {
244        self.instructions.extend(iter);
245    }
246}
247
248impl Deref for AtomicGroup {
249    type Target = [Instruction];
250
251    fn deref(&self) -> &Self::Target {
252        self.instructions.deref()
253    }
254}
255
256/// A group of atomic instructions that can be executed in parallel.
257#[derive(Debug, Clone, Default)]
258pub struct ParallelGroup(SmallVec<[AtomicGroup; PARALLEL_SIZE]>);
259
260impl ParallelGroup {
261    /// Add an [`AtomicGroup`].
262    pub fn add(&mut self, group: AtomicGroup) -> &mut Self {
263        self.0.push(group);
264        self
265    }
266
267    pub(crate) fn optimize(
268        &mut self,
269        options: &TransactionGroupOptions,
270        luts: &AddressLookupTables,
271        allow_payer_change: bool,
272    ) -> &mut Self {
273        if options.optimize(&mut self.0, luts, allow_payer_change) {
274            self.0 = self.0.drain(..).filter(|group| !group.is_empty()).collect();
275        }
276        self
277    }
278
279    pub(crate) fn single(&self) -> Option<&AtomicGroup> {
280        if self.0.len() == 1 {
281            Some(&self.0[0])
282        } else {
283            None
284        }
285    }
286
287    pub(crate) fn single_mut(&mut self) -> Option<&mut AtomicGroup> {
288        if self.0.len() == 1 {
289            Some(&mut self.0[0])
290        } else {
291            None
292        }
293    }
294
295    pub(crate) fn into_single(mut self) -> Option<AtomicGroup> {
296        if self.0.len() == 1 {
297            Some(self.0.remove(0))
298        } else {
299            None
300        }
301    }
302}
303
304impl From<AtomicGroup> for ParallelGroup {
305    fn from(value: AtomicGroup) -> Self {
306        let mut this = Self::default();
307        this.add(value);
308        this
309    }
310}
311
312impl FromIterator<AtomicGroup> for ParallelGroup {
313    fn from_iter<T: IntoIterator<Item = AtomicGroup>>(iter: T) -> Self {
314        Self(FromIterator::from_iter(iter))
315    }
316}
317
318impl Deref for ParallelGroup {
319    type Target = [AtomicGroup];
320
321    fn deref(&self) -> &Self::Target {
322        self.0.deref()
323    }
324}