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
26pub trait IntoAtomicGroup {
28 type Hint;
30
31 fn into_atomic_group(self, hint: &Self::Hint) -> crate::Result<AtomicGroup>;
33
34 #[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#[derive(Debug, Clone, Default)]
55pub struct GetInstructionsOptions {
56 pub compute_budget: ComputeBudgetOptions,
58 pub memo: Option<String>,
60}
61
62#[derive(Debug, Clone, Default)]
64pub struct ComputeBudgetOptions {
65 pub without_compute_budget: bool,
67 pub compute_unit_price_micro_lamports: Option<u64>,
69 pub compute_unit_min_priority_lamports: Option<u64>,
71}
72
73#[derive(Debug, Clone)]
75pub struct AtomicGroupOptions {
76 pub is_mergeable: bool,
78}
79
80impl Default for AtomicGroupOptions {
81 fn default() -> Self {
82 Self { is_mergeable: true }
83 }
84}
85
86#[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 pub fn is_mergeable(&self) -> bool {
100 self.options().is_mergeable
101 }
102
103 pub fn options(&self) -> &AtomicGroupOptions {
105 &self.options
106 }
107
108 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 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 pub fn new(payer: &Pubkey) -> Self {
134 Self::with_instructions(payer, None)
135 }
136
137 pub fn add(&mut self, instruction: Instruction) -> &mut Self {
139 self.instructions.push(instruction);
140 self
141 }
142
143 pub fn add_signer(&mut self, signer: &Pubkey) -> &mut Self {
145 self.signers.insert(*signer, NullSigner::new(signer));
146 self
147 }
148
149 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 pub fn compute_budget(&self) -> &ComputeBudget {
158 &self.compute_budget
159 }
160
161 pub fn compute_budget_mut(&mut self) -> &mut ComputeBudget {
163 &mut self.compute_budget
164 }
165
166 pub fn payer(&self) -> &Pubkey {
168 &self.payer
169 }
170
171 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 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 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 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 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 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 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 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#[derive(Debug, Clone)]
366pub struct ParallelGroupOptions {
367 pub is_mergeable: bool,
369}
370
371impl Default for ParallelGroupOptions {
372 fn default() -> Self {
373 Self { is_mergeable: true }
374 }
375}
376
377#[derive(Debug, Clone, Default)]
379pub struct ParallelGroup {
380 groups: SmallVec<[AtomicGroup; PARALLEL_SIZE]>,
381 options: ParallelGroupOptions,
382}
383
384impl ParallelGroup {
385 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 pub fn options(&self) -> &ParallelGroupOptions {
398 &self.options
399 }
400
401 pub fn is_mergeable(&self) -> bool {
403 self.options().is_mergeable
404 }
405
406 pub fn set_is_mergeable(&mut self, is_mergeable: bool) -> &mut Self {
408 self.options.is_mergeable = is_mergeable;
409 self
410 }
411
412 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 pub fn len(&self) -> usize {
460 self.groups.len()
461 }
462
463 pub fn is_empty(&self) -> bool {
465 self.groups.is_empty()
466 }
467
468 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}