1#[cfg(feature = "payload-builder")]
11pub mod builder;
12
13use ab_contracts_common::env::{MethodContext, PreparedMethod};
14use ab_contracts_common::method::MethodFingerprint;
15use ab_contracts_common::{Address, MAX_TOTAL_METHOD_ARGS};
16use ab_contracts_io_type::MAX_ALIGNMENT;
17use ab_contracts_io_type::trivial_type::TrivialType;
18use core::ffi::c_void;
19use core::marker::PhantomData;
20use core::mem::MaybeUninit;
21use core::num::{NonZeroU8, NonZeroUsize};
22use core::ops::{Deref, DerefMut};
23use core::ptr::NonNull;
24use core::slice;
25
26#[derive(Debug, Copy, Clone, Eq, PartialEq, TrivialType)]
27#[repr(u8)]
28pub enum TransactionMethodContext {
29 Null,
31 Wallet,
33}
34
35impl TransactionMethodContext {
36 #[inline(always)]
39 pub const fn try_from_u8(n: u8) -> Option<Self> {
40 Some(match n {
41 0 => Self::Null,
42 1 => Self::Wallet,
43 _ => {
44 return None;
45 }
46 })
47 }
48}
49
50#[derive(Debug, Copy, Eq, PartialEq, Clone)]
51pub enum TransactionInputType {
52 Value { alignment_power: u8 },
53 OutputIndex { output_index: u8 },
54}
55
56#[derive(Debug, Copy, Eq, PartialEq, Clone)]
57pub enum TransactionSlotType {
58 Address,
59 OutputIndex { output_index: u8 },
60}
61
62#[derive(Debug, Copy, Clone)]
71pub struct TransactionSlot(TransactionSlotType);
72
73impl TransactionSlot {
74 #[inline(always)]
76 pub const fn new_address() -> Self {
77 Self(TransactionSlotType::Address)
78 }
79
80 #[inline(always)]
84 pub const fn new_output_index(output_index: u8) -> Option<Self> {
85 if output_index > 0b0111_1111 {
86 return None;
87 }
88
89 Some(Self(TransactionSlotType::OutputIndex { output_index }))
90 }
91
92 #[inline(always)]
94 pub const fn from_u8(n: u8) -> Self {
95 if n & 0b1000_0000 == 0 {
97 Self(TransactionSlotType::OutputIndex { output_index: n })
98 } else {
99 Self(TransactionSlotType::Address)
100 }
101 }
102
103 #[inline(always)]
105 pub const fn into_u8(self) -> u8 {
106 match self.0 {
108 TransactionSlotType::Address => 0b1000_0000,
109 TransactionSlotType::OutputIndex { output_index } => output_index,
110 }
111 }
112
113 #[inline(always)]
115 pub const fn slot_type(self) -> TransactionSlotType {
116 self.0
117 }
118}
119
120#[derive(Debug, Copy, Clone)]
129pub struct TransactionInput(TransactionInputType);
130
131impl TransactionInput {
132 #[inline(always)]
136 pub const fn new_value(alignment: NonZeroU8) -> Option<Self> {
137 match alignment.get() {
138 1 | 2 | 4 | 8 | 16 => Some(Self(TransactionInputType::Value {
139 alignment_power: alignment.ilog2() as u8,
140 })),
141 _ => None,
142 }
143 }
144
145 #[inline(always)]
149 pub const fn new_output_index(output_index: u8) -> Option<Self> {
150 if output_index > 0b0111_1111 {
151 return None;
152 }
153
154 Some(Self(TransactionInputType::OutputIndex { output_index }))
155 }
156
157 #[inline(always)]
159 pub const fn from_u8(n: u8) -> Self {
160 if n & 0b1000_0000 == 0 {
162 Self(TransactionInputType::OutputIndex { output_index: n })
163 } else {
164 Self(TransactionInputType::Value {
165 alignment_power: n & 0b0111_1111,
166 })
167 }
168 }
169
170 #[inline(always)]
172 pub const fn into_u8(self) -> u8 {
173 match self.0 {
175 TransactionInputType::Value { alignment_power } => 0b1000_0000 | alignment_power,
176 TransactionInputType::OutputIndex { output_index } => output_index,
177 }
178 }
179
180 #[inline(always)]
182 pub const fn input_type(self) -> TransactionInputType {
183 self.0
184 }
185}
186
187#[derive(Debug, thiserror::Error)]
189pub enum TransactionPayloadDecoderError {
190 #[error("Payload too small")]
192 PayloadTooSmall,
193 #[error("Too many arguments")]
195 TooManyArguments(u8),
196 #[error("`ExternalArgs` buffer too small")]
198 ExternalArgsBufferTooSmall,
199 #[error("Output index not found: {0}")]
201 OutputIndexNotFound(u8),
202 #[error("Invalid output index {output_index} size for slot: {size}")]
204 InvalidSlotOutputIndexSize { output_index: u8, size: u32 },
205 #[error("Invalid output index {output_index} alignment for slot: {alignment}")]
207 InvalidSlotOutputIndexAlign { output_index: u8, alignment: u32 },
208 #[error("Alignment power is too large: {0}")]
210 AlignmentPowerTooLarge(u8),
211 #[error("Output buffer too small")]
213 OutputBufferTooSmall,
214 #[error("Output buffer offsets too small")]
216 OutputBufferOffsetsTooSmall,
217}
218
219#[derive(Debug)]
221pub struct TransactionPayloadDecoder<'a> {
222 payload: &'a [u8],
223 external_args_buffer: &'a mut [*mut c_void],
224 output_buffer: &'a mut [MaybeUninit<u128>],
225 output_buffer_cursor: usize,
226 output_buffer_offsets: &'a mut [MaybeUninit<(u32, u32)>],
227 output_buffer_offsets_cursor: usize,
228 map_context: fn(TransactionMethodContext) -> MethodContext,
229}
230
231impl<'a> TransactionPayloadDecoder<'a> {
232 #[inline]
245 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
246 pub fn new(
247 payload: &'a [u128],
248 external_args_buffer: &'a mut [*mut c_void],
249 output_buffer: &'a mut [MaybeUninit<u128>],
250 output_buffer_offsets: &'a mut [MaybeUninit<(u32, u32)>],
251 map_context: fn(TransactionMethodContext) -> MethodContext,
252 ) -> Self {
253 debug_assert_eq!(align_of_val(payload), usize::from(MAX_ALIGNMENT));
254 debug_assert_eq!(align_of_val(output_buffer), usize::from(MAX_ALIGNMENT));
255
256 let payload =
258 unsafe { slice::from_raw_parts(payload.as_ptr().cast::<u8>(), size_of_val(payload)) };
259
260 Self {
261 payload,
262 external_args_buffer,
263 output_buffer,
264 output_buffer_cursor: 0,
265 output_buffer_offsets,
266 output_buffer_offsets_cursor: 0,
267 map_context,
268 }
269 }
270}
271
272impl<'a> TransactionPayloadDecoder<'a> {
273 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
275 pub fn decode_next_method(
276 &mut self,
277 ) -> Result<Option<PreparedMethod<'_>>, TransactionPayloadDecoderError> {
278 TransactionPayloadDecoderInternal::<true>(self).decode_next_method()
279 }
280
281 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
287 pub unsafe fn decode_next_method_unchecked(&mut self) -> Option<PreparedMethod<'_>> {
288 TransactionPayloadDecoderInternal::<false>(self)
289 .decode_next_method()
290 .expect("No decoding errors are possible with trusted input; qed")
291 }
292}
293
294struct TransactionPayloadDecoderInternal<'tmp, 'decoder, const VERIFY: bool>(
298 &'tmp mut TransactionPayloadDecoder<'decoder>,
299);
300
301impl<'tmp, 'decoder, const VERIFY: bool> Deref
302 for TransactionPayloadDecoderInternal<'tmp, 'decoder, VERIFY>
303{
304 type Target = TransactionPayloadDecoder<'decoder>;
305
306 #[inline(always)]
307 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
308 fn deref(&self) -> &Self::Target {
309 self.0
310 }
311}
312
313impl<'tmp, 'decoder, const VERIFY: bool> DerefMut
314 for TransactionPayloadDecoderInternal<'tmp, 'decoder, VERIFY>
315{
316 #[inline(always)]
317 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
318 fn deref_mut(&mut self) -> &mut Self::Target {
319 self.0
320 }
321}
322
323impl<'tmp, 'decoder, const VERIFY: bool> TransactionPayloadDecoderInternal<'tmp, 'decoder, VERIFY> {
324 #[inline(always)]
325 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
326 fn decode_next_method(
327 mut self,
328 ) -> Result<Option<PreparedMethod<'decoder>>, TransactionPayloadDecoderError> {
329 if self.payload.len() <= usize::from(MAX_ALIGNMENT) {
330 return Ok(None);
331 }
332
333 let contract = self.get_trivial_type::<Address>()?;
334 let method_fingerprint = self.get_trivial_type::<MethodFingerprint>()?;
335 let method_context =
336 (self.map_context)(*self.get_trivial_type::<TransactionMethodContext>()?);
337
338 let mut transaction_slots_inputs =
339 [MaybeUninit::<u8>::uninit(); MAX_TOTAL_METHOD_ARGS as usize];
340
341 let num_slot_arguments = self.read_u8()?;
342 for transaction_slot in transaction_slots_inputs
343 .iter_mut()
344 .take(usize::from(num_slot_arguments))
345 {
346 transaction_slot.write(self.read_u8()?);
347 }
348
349 let num_input_arguments = self.read_u8()?;
350 for transaction_input in transaction_slots_inputs
351 .iter_mut()
352 .skip(usize::from(num_slot_arguments))
353 .take(usize::from(num_input_arguments))
354 {
355 transaction_input.write(self.read_u8()?);
356 }
357
358 let num_output_arguments = self.read_u8()?;
359
360 let (transaction_slots, transaction_inputs) = unsafe {
362 let (transaction_slots, transaction_inputs) =
363 transaction_slots_inputs.split_at_unchecked(usize::from(num_slot_arguments));
364 let (transaction_inputs, _) =
365 transaction_inputs.split_at_unchecked(usize::from(num_input_arguments));
366
367 (
368 transaction_slots.assume_init_ref(),
369 transaction_inputs.assume_init_ref(),
370 )
371 };
372
373 let number_of_arguments = num_slot_arguments
376 .saturating_add(num_input_arguments)
377 .saturating_add(num_output_arguments);
378 if VERIFY && number_of_arguments > MAX_TOTAL_METHOD_ARGS {
379 return Err(TransactionPayloadDecoderError::TooManyArguments(
380 number_of_arguments,
381 ));
382 }
383
384 let expected_external_args_buffer_size =
387 usize::from(num_slot_arguments + num_input_arguments * 2 + num_output_arguments * 3);
388 if VERIFY && expected_external_args_buffer_size > self.external_args_buffer.len() {
389 return Err(TransactionPayloadDecoderError::ExternalArgsBufferTooSmall);
390 }
391
392 let external_args =
393 NonNull::new(self.external_args_buffer.as_mut_ptr()).expect("Not null; qed");
394 {
395 let mut external_args_cursor = external_args;
396
397 for &transaction_slot in transaction_slots {
398 let address = match TransactionSlot::from_u8(transaction_slot).slot_type() {
399 TransactionSlotType::Address => self.get_trivial_type::<Address>()?,
400 TransactionSlotType::OutputIndex { output_index } => {
401 let (bytes, &size) = self.get_from_output_buffer(output_index)?;
402
403 if VERIFY && size != Address::SIZE {
404 return Err(
405 TransactionPayloadDecoderError::InvalidSlotOutputIndexSize {
406 output_index,
407 size,
408 },
409 );
410 }
411
412 let maybe_address = unsafe { Address::from_bytes(bytes) };
414
415 if VERIFY {
416 let Some(address) = maybe_address else {
417 let error =
418 TransactionPayloadDecoderError::InvalidSlotOutputIndexAlign {
419 output_index,
420 alignment: align_of_val(bytes) as u32,
421 };
422 return Err(error);
423 };
424
425 address
426 } else {
427 unsafe { maybe_address.unwrap_unchecked() }
429 }
430 }
431 };
432
433 unsafe {
436 external_args_cursor.cast::<*const Address>().write(address);
437 external_args_cursor = external_args_cursor.offset(1);
438 }
439 }
440
441 for &transaction_input in transaction_inputs {
442 let (bytes, size) = match TransactionInput::from_u8(transaction_input).input_type()
443 {
444 TransactionInputType::Value { alignment_power } => {
445 let alignment = if VERIFY {
448 1_usize.checked_shl(u32::from(alignment_power)).ok_or(
449 TransactionPayloadDecoderError::AlignmentPowerTooLarge(
450 alignment_power,
451 ),
452 )?
453 } else {
454 unsafe { 1_usize.unchecked_shl(u32::from(alignment_power)) }
456 };
457
458 let size = self.get_trivial_type::<u32>()?;
459 let bytes = self.get_bytes(
460 *size,
461 NonZeroUsize::new(alignment).expect("Not zero; qed"),
462 )?;
463
464 (bytes, size)
465 }
466 TransactionInputType::OutputIndex { output_index } => {
467 self.get_from_output_buffer(output_index)?
468 }
469 };
470
471 unsafe {
474 external_args_cursor
475 .cast::<*const u8>()
476 .write(bytes.as_ptr());
477 external_args_cursor = external_args_cursor.offset(1);
478
479 external_args_cursor.cast::<*const u32>().write(size);
480 external_args_cursor = external_args_cursor.offset(1);
481 }
482 }
483
484 for _ in 0..num_output_arguments {
485 let recommended_capacity = self.get_trivial_type::<u32>()?;
486 let alignment_power = *self.get_trivial_type::<u8>()?;
487 let alignment = if VERIFY {
490 1_usize.checked_shl(u32::from(alignment_power)).ok_or(
491 TransactionPayloadDecoderError::AlignmentPowerTooLarge(alignment_power),
492 )?
493 } else {
494 unsafe { 1_usize.unchecked_shl(u32::from(alignment_power)) }
496 };
497
498 let (size, data) = self.allocate_output_buffer(
499 *recommended_capacity,
500 NonZeroUsize::new(alignment).expect("Not zero; qed"),
501 )?;
502
503 unsafe {
506 external_args_cursor.cast::<*mut u8>().write(data.as_ptr());
507 external_args_cursor = external_args_cursor.offset(1);
508
509 external_args_cursor.cast::<*mut u32>().write(size.as_ptr());
510 external_args_cursor = external_args_cursor.offset(1);
511
512 external_args_cursor
513 .cast::<*const u32>()
514 .write(recommended_capacity);
515 external_args_cursor = external_args_cursor.offset(1);
516 }
517 }
518 }
519
520 Ok(Some(PreparedMethod {
521 contract: *contract,
522 fingerprint: *method_fingerprint,
523 external_args: external_args.cast::<NonNull<c_void>>(),
524 method_context,
525 phantom: PhantomData,
526 }))
527 }
528
529 #[inline(always)]
530 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
531 fn get_trivial_type<T>(&mut self) -> Result<&'decoder T, TransactionPayloadDecoderError>
532 where
533 T: TrivialType,
534 {
535 self.ensure_alignment(NonZeroUsize::new(align_of::<T>()).expect("Not zero; qed"));
536
537 let bytes;
538 if VERIFY {
539 (bytes, self.payload) = self
540 .payload
541 .split_at_checked(size_of::<T>())
542 .ok_or(TransactionPayloadDecoderError::PayloadTooSmall)?;
543 } else {
544 (bytes, self.payload) = unsafe { self.payload.split_at_unchecked(size_of::<T>()) };
546 }
547
548 let value_ref = unsafe { bytes.as_ptr().cast::<T>().as_ref().expect("Not null; qed") };
550
551 Ok(value_ref)
552 }
553
554 #[inline(always)]
555 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
556 fn get_bytes(
557 &mut self,
558 size: u32,
559 alignment: NonZeroUsize,
560 ) -> Result<&'decoder [u8], TransactionPayloadDecoderError> {
561 self.ensure_alignment(alignment);
562
563 let bytes;
564 if VERIFY {
565 (bytes, self.payload) = self
566 .payload
567 .split_at_checked(size as usize)
568 .ok_or(TransactionPayloadDecoderError::PayloadTooSmall)?;
569 } else {
570 (bytes, self.payload) = unsafe { self.payload.split_at_unchecked(size as usize) };
572 }
573
574 Ok(bytes)
575 }
576
577 #[inline(always)]
578 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
579 fn read_u8(&mut self) -> Result<u8, TransactionPayloadDecoderError> {
580 let value;
581 if VERIFY {
582 (value, self.payload) = self
583 .payload
584 .split_at_checked(1)
585 .ok_or(TransactionPayloadDecoderError::PayloadTooSmall)?;
586 } else {
587 (value, self.payload) = unsafe { self.payload.split_at_unchecked(1) };
589 }
590
591 Ok(value[0])
592 }
593
594 #[inline(always)]
595 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
596 fn ensure_alignment(&mut self, alignment: NonZeroUsize) {
597 debug_assert!(alignment.get() <= usize::from(MAX_ALIGNMENT));
598
599 let unaligned_by = {
602 let mask = alignment
603 .get()
604 .checked_sub(1)
605 .expect("Left side is not zero; qed");
606 self.payload.len() & mask
607 };
608 self.payload = &self.payload[unaligned_by..];
609 }
610
611 #[inline(always)]
612 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
613 fn allocate_output_buffer(
614 &mut self,
615 capacity: u32,
616 output_alignment: NonZeroUsize,
617 ) -> Result<(NonNull<u32>, NonNull<u8>), TransactionPayloadDecoderError> {
618 if VERIFY && self.output_buffer_offsets.len() == self.output_buffer_offsets_cursor {
619 return Err(TransactionPayloadDecoderError::OutputBufferOffsetsTooSmall);
620 }
621
622 let (size_offset, size_ptr) = self
623 .allocate_output_buffer_ptr(
624 NonZeroUsize::new(align_of::<u32>()).expect("Not zero; qed"),
625 size_of::<u32>(),
626 )
627 .ok_or(TransactionPayloadDecoderError::OutputBufferTooSmall)?;
628 let (output_offset, output_ptr) = self
629 .allocate_output_buffer_ptr(output_alignment, capacity as usize)
630 .ok_or(TransactionPayloadDecoderError::OutputBufferTooSmall)?;
631
632 let output_buffer_offsets = unsafe {
634 let output_buffer_offsets_cursor = self.output_buffer_offsets_cursor;
636 self.output_buffer_offsets
637 .get_unchecked_mut(output_buffer_offsets_cursor)
638 };
639 output_buffer_offsets.write((size_offset as u32, output_offset as u32));
640 self.output_buffer_offsets_cursor += 1;
641
642 Ok((size_ptr, output_ptr))
643 }
644
645 #[inline(always)]
647 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
648 fn allocate_output_buffer_ptr<T>(
649 &mut self,
650 alignment: NonZeroUsize,
651 size: usize,
652 ) -> Option<(usize, NonNull<T>)> {
653 debug_assert!(alignment.get() <= usize::from(MAX_ALIGNMENT));
654
655 let unaligned_by = {
658 let mask = alignment
659 .get()
660 .checked_sub(1)
661 .expect("Left side is not zero; qed");
662 self.output_buffer_cursor & mask
663 };
664
665 let new_output_buffer_cursor = if VERIFY {
666 let new_output_buffer_cursor = self
667 .output_buffer_cursor
668 .checked_add(unaligned_by)?
669 .checked_add(size)?;
670
671 if new_output_buffer_cursor > size_of_val(self.output_buffer) {
672 return None;
673 }
674
675 new_output_buffer_cursor
676 } else {
677 unsafe {
679 self.output_buffer_cursor
680 .unchecked_add(unaligned_by)
681 .unchecked_add(size)
682 }
683 };
684
685 let (offset, buffer_ptr) = unsafe {
687 let offset = self.output_buffer_cursor.unchecked_add(unaligned_by);
688 let buffer_ptr = NonNull::new_unchecked(
689 self.output_buffer.as_mut_ptr().byte_add(offset).cast::<T>(),
690 );
691
692 (offset, buffer_ptr)
693 };
694 self.output_buffer_cursor = new_output_buffer_cursor;
695
696 Some((offset, buffer_ptr))
697 }
698
699 #[inline(always)]
700 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
701 fn get_from_output_buffer(
702 &self,
703 output_index: u8,
704 ) -> Result<(&[u8], &u32), TransactionPayloadDecoderError> {
705 let (size_offset, output_offset) = if VERIFY {
706 if usize::from(output_index) < self.output_buffer_offsets_cursor {
707 unsafe {
709 self.output_buffer_offsets
710 .get_unchecked(usize::from(output_index))
711 .assume_init()
712 }
713 } else {
714 return Err(TransactionPayloadDecoderError::OutputIndexNotFound(
715 output_index,
716 ));
717 }
718 } else {
719 unsafe {
721 self.output_buffer_offsets
722 .get_unchecked(usize::from(output_index))
723 .assume_init()
724 }
725 };
726
727 let size = unsafe {
730 self.output_buffer
731 .as_ptr()
732 .byte_add(size_offset as usize)
733 .cast::<u32>()
734 .as_ref_unchecked()
735 };
736 let bytes = unsafe {
739 let bytes_ptr = self
740 .output_buffer
741 .as_ptr()
742 .cast::<u8>()
743 .add(output_offset as usize);
744
745 slice::from_raw_parts(bytes_ptr, *size as usize)
746 };
747
748 Ok((bytes, size))
749 }
750}