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