1use {
2 crate::{
3 compute_budget::DEFAULT_HEAP_COST,
4 prioritization_fee::{PrioritizationFeeDetails, PrioritizationFeeType},
5 },
6 miraland_sdk::{
7 borsh1::try_from_slice_unchecked,
8 compute_budget::{self, ComputeBudgetInstruction},
9 entrypoint::HEAP_LENGTH as MIN_HEAP_FRAME_BYTES,
10 fee::FeeBudgetLimits,
11 instruction::{CompiledInstruction, InstructionError},
12 pubkey::Pubkey,
13 transaction::TransactionError,
14 },
15};
16
17const MAX_HEAP_FRAME_BYTES: u32 = 256 * 1024;
18pub const DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT: u32 = 200_000;
19pub const MAX_COMPUTE_UNIT_LIMIT: u32 = 1_400_000;
20
21pub const MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES: u32 = 64 * 1024 * 1024;
24
25#[derive(Clone, Debug, PartialEq, Eq)]
26pub struct ComputeBudgetLimits {
27 pub updated_heap_bytes: u32,
28 pub compute_unit_limit: u32,
29 pub compute_unit_price: u64,
30 pub loaded_accounts_bytes: u32,
31}
32
33impl Default for ComputeBudgetLimits {
34 fn default() -> Self {
35 ComputeBudgetLimits {
36 updated_heap_bytes: u32::try_from(MIN_HEAP_FRAME_BYTES).unwrap(),
37 compute_unit_limit: MAX_COMPUTE_UNIT_LIMIT,
38 compute_unit_price: 0,
39 loaded_accounts_bytes: MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES,
40 }
41 }
42}
43
44impl From<ComputeBudgetLimits> for FeeBudgetLimits {
45 fn from(val: ComputeBudgetLimits) -> Self {
46 let prioritization_fee_details = PrioritizationFeeDetails::new(
47 PrioritizationFeeType::ComputeUnitPrice(val.compute_unit_price),
48 u64::from(val.compute_unit_limit),
49 );
50 let prioritization_fee = prioritization_fee_details.get_fee();
51
52 FeeBudgetLimits {
53 loaded_accounts_data_size_limit: usize::try_from(val.loaded_accounts_bytes).unwrap(),
56 heap_cost: DEFAULT_HEAP_COST,
57 compute_unit_limit: u64::from(val.compute_unit_limit),
58 prioritization_fee,
59 }
60 }
61}
62
63pub fn process_compute_budget_instructions<'a>(
69 instructions: impl Iterator<Item = (&'a Pubkey, &'a CompiledInstruction)>,
70) -> Result<ComputeBudgetLimits, TransactionError> {
71 let mut num_non_compute_budget_instructions: u32 = 0;
72 let mut updated_compute_unit_limit = None;
73 let mut updated_compute_unit_price = None;
74 let mut requested_heap_size = None;
75 let mut updated_loaded_accounts_data_size_limit = None;
76
77 for (i, (program_id, instruction)) in instructions.enumerate() {
78 if compute_budget::check_id(program_id) {
79 let invalid_instruction_data_error = TransactionError::InstructionError(
80 i as u8,
81 InstructionError::InvalidInstructionData,
82 );
83 let duplicate_instruction_error = TransactionError::DuplicateInstruction(i as u8);
84
85 match try_from_slice_unchecked(&instruction.data) {
86 Ok(ComputeBudgetInstruction::RequestHeapFrame(bytes)) => {
87 if requested_heap_size.is_some() {
88 return Err(duplicate_instruction_error);
89 }
90 if sanitize_requested_heap_size(bytes) {
91 requested_heap_size = Some(bytes);
92 } else {
93 return Err(invalid_instruction_data_error);
94 }
95 }
96 Ok(ComputeBudgetInstruction::SetComputeUnitLimit(compute_unit_limit)) => {
97 if updated_compute_unit_limit.is_some() {
98 return Err(duplicate_instruction_error);
99 }
100 updated_compute_unit_limit = Some(compute_unit_limit);
101 }
102 Ok(ComputeBudgetInstruction::SetComputeUnitPrice(micro_lamports)) => {
103 if updated_compute_unit_price.is_some() {
104 return Err(duplicate_instruction_error);
105 }
106 updated_compute_unit_price = Some(micro_lamports);
107 }
108 Ok(ComputeBudgetInstruction::SetLoadedAccountsDataSizeLimit(bytes)) => {
109 if updated_loaded_accounts_data_size_limit.is_some() {
110 return Err(duplicate_instruction_error);
111 }
112 updated_loaded_accounts_data_size_limit = Some(bytes);
113 }
114 _ => return Err(invalid_instruction_data_error),
115 }
116 } else {
117 num_non_compute_budget_instructions =
119 num_non_compute_budget_instructions.saturating_add(1);
120 }
121 }
122
123 let updated_heap_bytes = requested_heap_size
125 .unwrap_or(u32::try_from(MIN_HEAP_FRAME_BYTES).unwrap()) .min(MAX_HEAP_FRAME_BYTES);
127
128 let compute_unit_limit = updated_compute_unit_limit
129 .unwrap_or_else(|| {
130 num_non_compute_budget_instructions
131 .saturating_mul(DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT)
132 })
133 .min(MAX_COMPUTE_UNIT_LIMIT);
134
135 let compute_unit_price = updated_compute_unit_price.unwrap_or(0);
136
137 let loaded_accounts_bytes = updated_loaded_accounts_data_size_limit
138 .unwrap_or(MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES)
139 .min(MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES);
140
141 Ok(ComputeBudgetLimits {
142 updated_heap_bytes,
143 compute_unit_limit,
144 compute_unit_price,
145 loaded_accounts_bytes,
146 })
147}
148
149fn sanitize_requested_heap_size(bytes: u32) -> bool {
150 (u32::try_from(MIN_HEAP_FRAME_BYTES).unwrap()..=MAX_HEAP_FRAME_BYTES).contains(&bytes)
151 && bytes % 1024 == 0
152}
153
154#[cfg(test)]
155mod tests {
156 use {
157 super::*,
158 miraland_sdk::{
159 hash::Hash,
160 instruction::Instruction,
161 message::Message,
162 pubkey::Pubkey,
163 signature::Keypair,
164 signer::Signer,
165 system_instruction::{self},
166 transaction::{SanitizedTransaction, Transaction},
167 },
168 };
169
170 macro_rules! test {
171 ( $instructions: expr, $expected_result: expr) => {
172 let payer_keypair = Keypair::new();
173 let tx = SanitizedTransaction::from_transaction_for_tests(Transaction::new(
174 &[&payer_keypair],
175 Message::new($instructions, Some(&payer_keypair.pubkey())),
176 Hash::default(),
177 ));
178 let result =
179 process_compute_budget_instructions(tx.message().program_instructions_iter());
180 assert_eq!($expected_result, result);
181 };
182 }
183
184 #[test]
185 fn test_process_instructions() {
186 test!(
188 &[],
189 Ok(ComputeBudgetLimits {
190 compute_unit_limit: 0,
191 ..ComputeBudgetLimits::default()
192 })
193 );
194 test!(
195 &[
196 ComputeBudgetInstruction::set_compute_unit_limit(1),
197 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
198 ],
199 Ok(ComputeBudgetLimits {
200 compute_unit_limit: 1,
201 ..ComputeBudgetLimits::default()
202 })
203 );
204 test!(
205 &[
206 ComputeBudgetInstruction::set_compute_unit_limit(MAX_COMPUTE_UNIT_LIMIT + 1),
207 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
208 ],
209 Ok(ComputeBudgetLimits {
210 compute_unit_limit: MAX_COMPUTE_UNIT_LIMIT,
211 ..ComputeBudgetLimits::default()
212 })
213 );
214 test!(
215 &[
216 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
217 ComputeBudgetInstruction::set_compute_unit_limit(MAX_COMPUTE_UNIT_LIMIT),
218 ],
219 Ok(ComputeBudgetLimits {
220 compute_unit_limit: MAX_COMPUTE_UNIT_LIMIT,
221 ..ComputeBudgetLimits::default()
222 })
223 );
224 test!(
225 &[
226 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
227 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
228 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
229 ComputeBudgetInstruction::set_compute_unit_limit(1),
230 ],
231 Ok(ComputeBudgetLimits {
232 compute_unit_limit: 1,
233 ..ComputeBudgetLimits::default()
234 })
235 );
236 test!(
237 &[
238 ComputeBudgetInstruction::set_compute_unit_limit(1),
239 ComputeBudgetInstruction::set_compute_unit_price(42)
240 ],
241 Ok(ComputeBudgetLimits {
242 compute_unit_limit: 1,
243 compute_unit_price: 42,
244 ..ComputeBudgetLimits::default()
245 })
246 );
247
248 test!(
250 &[],
251 Ok(ComputeBudgetLimits {
252 compute_unit_limit: 0,
253 ..ComputeBudgetLimits::default()
254 })
255 );
256 test!(
257 &[
258 ComputeBudgetInstruction::request_heap_frame(40 * 1024),
259 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
260 ],
261 Ok(ComputeBudgetLimits {
262 compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
263 updated_heap_bytes: 40 * 1024,
264 ..ComputeBudgetLimits::default()
265 })
266 );
267 test!(
268 &[
269 ComputeBudgetInstruction::request_heap_frame(40 * 1024 + 1),
270 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
271 ],
272 Err(TransactionError::InstructionError(
273 0,
274 InstructionError::InvalidInstructionData,
275 ))
276 );
277 test!(
278 &[
279 ComputeBudgetInstruction::request_heap_frame(31 * 1024),
280 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
281 ],
282 Err(TransactionError::InstructionError(
283 0,
284 InstructionError::InvalidInstructionData,
285 ))
286 );
287 test!(
288 &[
289 ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES + 1),
290 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
291 ],
292 Err(TransactionError::InstructionError(
293 0,
294 InstructionError::InvalidInstructionData,
295 ))
296 );
297 test!(
298 &[
299 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
300 ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES),
301 ],
302 Ok(ComputeBudgetLimits {
303 compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
304 updated_heap_bytes: MAX_HEAP_FRAME_BYTES,
305 ..ComputeBudgetLimits::default()
306 })
307 );
308 test!(
309 &[
310 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
311 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
312 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
313 ComputeBudgetInstruction::request_heap_frame(1),
314 ],
315 Err(TransactionError::InstructionError(
316 3,
317 InstructionError::InvalidInstructionData,
318 ))
319 );
320 test!(
321 &[
322 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
323 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
324 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
325 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
326 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
327 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
328 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
329 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
330 ],
331 Ok(ComputeBudgetLimits {
332 compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT * 7,
333 ..ComputeBudgetLimits::default()
334 })
335 );
336
337 test!(
339 &[
340 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
341 ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES),
342 ComputeBudgetInstruction::set_compute_unit_limit(MAX_COMPUTE_UNIT_LIMIT),
343 ComputeBudgetInstruction::set_compute_unit_price(u64::MAX),
344 ],
345 Ok(ComputeBudgetLimits {
346 compute_unit_price: u64::MAX,
347 compute_unit_limit: MAX_COMPUTE_UNIT_LIMIT,
348 updated_heap_bytes: MAX_HEAP_FRAME_BYTES,
349 ..ComputeBudgetLimits::default()
350 })
351 );
352 test!(
353 &[
354 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
355 ComputeBudgetInstruction::set_compute_unit_limit(1),
356 ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES),
357 ComputeBudgetInstruction::set_compute_unit_price(u64::MAX),
358 ],
359 Ok(ComputeBudgetLimits {
360 compute_unit_price: u64::MAX,
361 compute_unit_limit: 1,
362 updated_heap_bytes: MAX_HEAP_FRAME_BYTES,
363 ..ComputeBudgetLimits::default()
364 })
365 );
366
367 test!(
369 &[
370 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
371 ComputeBudgetInstruction::set_compute_unit_limit(MAX_COMPUTE_UNIT_LIMIT),
372 ComputeBudgetInstruction::set_compute_unit_limit(MAX_COMPUTE_UNIT_LIMIT - 1),
373 ],
374 Err(TransactionError::DuplicateInstruction(2))
375 );
376
377 test!(
378 &[
379 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
380 ComputeBudgetInstruction::request_heap_frame(MIN_HEAP_FRAME_BYTES as u32),
381 ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES),
382 ],
383 Err(TransactionError::DuplicateInstruction(2))
384 );
385 test!(
386 &[
387 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
388 ComputeBudgetInstruction::set_compute_unit_price(0),
389 ComputeBudgetInstruction::set_compute_unit_price(u64::MAX),
390 ],
391 Err(TransactionError::DuplicateInstruction(2))
392 );
393 }
394
395 #[test]
396 fn test_process_loaded_accounts_data_size_limit_instruction() {
397 test!(
398 &[],
399 Ok(ComputeBudgetLimits {
400 compute_unit_limit: 0,
401 ..ComputeBudgetLimits::default()
402 })
403 );
404
405 let data_size = 1;
408 let expected_result = Ok(ComputeBudgetLimits {
409 compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
410 loaded_accounts_bytes: data_size,
411 ..ComputeBudgetLimits::default()
412 });
413
414 test!(
415 &[
416 ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size),
417 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
418 ],
419 expected_result
420 );
421
422 let data_size = MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES + 1;
425 let expected_result = Ok(ComputeBudgetLimits {
426 compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
427 loaded_accounts_bytes: MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES,
428 ..ComputeBudgetLimits::default()
429 });
430
431 test!(
432 &[
433 ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size),
434 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
435 ],
436 expected_result
437 );
438
439 let expected_result = Ok(ComputeBudgetLimits {
442 compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
443 loaded_accounts_bytes: MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES,
444 ..ComputeBudgetLimits::default()
445 });
446
447 test!(
448 &[Instruction::new_with_bincode(
449 Pubkey::new_unique(),
450 &0_u8,
451 vec![]
452 ),],
453 expected_result
454 );
455
456 let data_size = MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES;
459 let expected_result = Err(TransactionError::DuplicateInstruction(2));
460
461 test!(
462 &[
463 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
464 ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size),
465 ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size),
466 ],
467 expected_result
468 );
469 }
470
471 #[test]
472 fn test_process_mixed_instructions_without_compute_budget() {
473 let payer_keypair = Keypair::new();
474
475 let transaction =
476 SanitizedTransaction::from_transaction_for_tests(Transaction::new_signed_with_payer(
477 &[
478 Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
479 system_instruction::transfer(&payer_keypair.pubkey(), &Pubkey::new_unique(), 2),
480 ],
481 Some(&payer_keypair.pubkey()),
482 &[&payer_keypair],
483 Hash::default(),
484 ));
485
486 let result =
487 process_compute_budget_instructions(transaction.message().program_instructions_iter());
488
489 assert_eq!(
493 result,
494 Ok(ComputeBudgetLimits {
495 compute_unit_limit: 2 * DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
496 ..ComputeBudgetLimits::default()
497 })
498 );
499 }
500}