Skip to main content

ab_system_contract_simple_wallet_base/
payload.rs

1//! This module contains generic utilities for serializing and deserializing method calls to/from
2//! payload bytes.
3//!
4//! It can be reused to implement a different wallet implementation as well as read and verify the
5//! contents of the transaction (for example, to display it on the screen of the hardware wallet).
6//!
7//! Builder interface requires heap allocations and can be enabled with `payload-builder` feature,
8//! while the rest works in `no_std` environment without a global allocator.
9
10#[cfg(feature = "payload-builder")]
11pub mod builder;
12
13use crate::EXTERNAL_ARGS_BUFFER_SIZE;
14use ab_contracts_common::MAX_TOTAL_METHOD_ARGS;
15use ab_contracts_common::env::{MethodContext, PreparedMethod};
16use ab_contracts_common::method::MethodFingerprint;
17use ab_core_primitives::address::Address;
18use ab_io_type::MAX_ALIGNMENT;
19use ab_io_type::trivial_type::TrivialType;
20use core::ffi::c_void;
21use core::marker::PhantomData;
22use core::mem::{MaybeUninit, offset_of};
23use core::num::{NonZeroU8, NonZeroUsize};
24use core::ops::{Deref, DerefMut};
25use core::ptr::NonNull;
26use core::{ptr, slice};
27
28#[derive(Copy, Clone)]
29#[repr(C)]
30struct FfiDataSizeCapacityRo {
31    data_ptr: NonNull<u8>,
32    size: u32,
33    capacity: u32,
34}
35
36#[derive(Copy, Clone)]
37#[repr(C)]
38struct FfiDataSizeCapacityRw {
39    data_ptr: *mut u8,
40    size: u32,
41    capacity: u32,
42}
43
44#[derive(Debug, Copy, Clone, Eq, PartialEq, TrivialType)]
45#[repr(u8)]
46pub enum TransactionMethodContext {
47    /// Call contract under [`Address::NULL`] context (corresponds to [`MethodContext::Reset`])
48    Null,
49    /// Call contract under context of the wallet (corresponds to [`MethodContext::Replace`])
50    Wallet,
51}
52
53impl const TryFrom<u8> for TransactionMethodContext {
54    type Error = ();
55
56    #[inline(always)]
57    fn try_from(value: u8) -> Result<Self, Self::Error> {
58        Ok(match value {
59            0 => Self::Null,
60            1 => Self::Wallet,
61            _ => return Err(()),
62        })
63    }
64}
65
66#[derive(Debug, Copy, Eq, PartialEq, Clone)]
67pub enum TransactionInputType {
68    Value { alignment_power: u8 },
69    OutputIndex { output_index: u8 },
70}
71
72#[derive(Debug, Copy, Eq, PartialEq, Clone)]
73pub enum TransactionSlotType {
74    Address,
75    OutputIndex { output_index: u8 },
76}
77
78/// The type of transaction slot could be either explicit slot address or output index.
79///
80/// Specifically, if the previous method has `#[output]` or return value, those values are collected
81/// and pushed into a virtual "stack". Then, if [`Self::slot_type()`] returns
82/// [`TransactionSlotType::OutputIndex`], then the corresponding input will use the value at
83/// `output_index` of this stack instead of what was specified in `external_args`. This allows
84/// composing calls to multiple methods into more sophisticated workflows without writing special
85/// contracts for this.
86#[derive(Debug, Copy, Clone)]
87pub struct TransactionSlot(TransactionSlotType);
88
89impl TransactionSlot {
90    /// Explicit slot address
91    #[inline(always)]
92    pub const fn new_address() -> Self {
93        Self(TransactionSlotType::Address)
94    }
95
96    /// Output index value.
97    ///
98    /// Valid index values are 0..=127.
99    #[inline(always)]
100    pub const fn new_output_index(output_index: u8) -> Option<Self> {
101        if output_index > 0b0111_1111 {
102            return None;
103        }
104
105        Some(Self(TransactionSlotType::OutputIndex { output_index }))
106    }
107
108    /// Create an instance from `u8`
109    #[inline(always)]
110    pub const fn from_u8(n: u8) -> Self {
111        // The first bit is set to 1 for explicit slot address and 0 for output index
112        if n & 0b1000_0000 == 0 {
113            Self(TransactionSlotType::OutputIndex { output_index: n })
114        } else {
115            Self(TransactionSlotType::Address)
116        }
117    }
118
119    /// Convert instance into `u8`
120    #[inline(always)]
121    pub const fn into_u8(self) -> u8 {
122        // The first bit is set to 1 for explicit slot address and 0 for output index
123        match self.0 {
124            TransactionSlotType::Address => 0b1000_0000,
125            TransactionSlotType::OutputIndex { output_index } => output_index,
126        }
127    }
128
129    /// Returns `Some(output_index)` or `None` if explicit slot address
130    #[inline(always)]
131    pub const fn slot_type(self) -> TransactionSlotType {
132        self.0
133    }
134}
135
136/// The type of transaction input could be either explicit value or output index.
137///
138/// Specifically, if the previous method has `#[output]` or return value, those values are collected
139/// and pushed into a virtual "stack". Then, if [`Self::input_type()`] returns
140/// [`TransactionInputType::OutputIndex`], then the corresponding input will use the value at
141/// `output_index` of this stack instead of what was specified in `external_args`. This allows
142/// composing calls to multiple methods into more sophisticated workflows without writing special
143/// contracts for this.
144#[derive(Debug, Copy, Clone)]
145pub struct TransactionInput(TransactionInputType);
146
147impl TransactionInput {
148    /// Regular input value with specified alignment.
149    ///
150    /// Valid alignment values are: 1, 2, 4, 8, 16.
151    #[inline(always)]
152    pub const fn new_value(alignment: NonZeroU8) -> Option<Self> {
153        match alignment.get() {
154            1 | 2 | 4 | 8 | 16 => Some(Self(TransactionInputType::Value {
155                alignment_power: alignment.ilog2() as u8,
156            })),
157            _ => None,
158        }
159    }
160
161    /// Output index value.
162    ///
163    /// Valid index values are 0..=127.
164    #[inline(always)]
165    pub const fn new_output_index(output_index: u8) -> Option<Self> {
166        if output_index > 0b0111_1111 {
167            return None;
168        }
169
170        Some(Self(TransactionInputType::OutputIndex { output_index }))
171    }
172
173    /// Create an instance from `u8`
174    #[inline(always)]
175    pub const fn from_u8(n: u8) -> Self {
176        // The first bit is set to 1 for value and 0 for output index
177        if n & 0b1000_0000 == 0 {
178            Self(TransactionInputType::OutputIndex { output_index: n })
179        } else {
180            Self(TransactionInputType::Value {
181                alignment_power: n & 0b0111_1111,
182            })
183        }
184    }
185
186    /// Convert instance into `u8`
187    #[inline(always)]
188    pub const fn into_u8(self) -> u8 {
189        // The first bit is set to 1 for value and 0 for output index
190        match self.0 {
191            TransactionInputType::Value { alignment_power } => 0b1000_0000 | alignment_power,
192            TransactionInputType::OutputIndex { output_index } => output_index,
193        }
194    }
195
196    /// Returns `Some(output_index)` or `None` if regular input value
197    #[inline(always)]
198    pub const fn input_type(self) -> TransactionInputType {
199        self.0
200    }
201}
202
203/// Errors for [`TransactionPayloadDecoder`]
204#[derive(Debug, thiserror::Error)]
205pub enum TransactionPayloadDecoderError {
206    /// Payload too small
207    #[error("Payload too small")]
208    PayloadTooSmall,
209    /// Too many arguments
210    #[error("Too many arguments")]
211    TooManyArguments(u8),
212    /// `ExternalArgs` buffer too small
213    #[error("`ExternalArgs` buffer too small")]
214    ExternalArgsBufferTooSmall,
215    /// Output index not found
216    #[error("Output index not found: {0}")]
217    OutputIndexNotFound(u8),
218    /// Invalid output index size for slot
219    #[error("Invalid output index {output_index} size for slot: {size}")]
220    InvalidSlotOutputIndexSize { output_index: u8, size: u32 },
221    /// Invalid output index alignment for slot
222    #[error("Invalid output index {output_index} alignment for slot: {alignment}")]
223    InvalidSlotOutputIndexAlign { output_index: u8, alignment: u32 },
224    /// Alignment power is too large
225    #[error("Alignment power is too large: {0}")]
226    AlignmentPowerTooLarge(u8),
227    /// Output buffer too small
228    #[error("Output buffer too small")]
229    OutputBufferTooSmall,
230    /// Output buffer offsets too small
231    #[error("Output buffer offsets too small")]
232    OutputBufferOffsetsTooSmall,
233}
234
235#[derive(Debug, Copy, Clone)]
236pub struct OutputBufferDetails {
237    /// Offset of output bytes inside `output_buffer`
238    output_offset: u32,
239    /// Size of the output buffer `output_offset` points to.
240    ///
241    /// NOTE: It temporarily stores the offset (in bytes) into `external_args_buffer` while
242    /// decoding a method. Before decoding the next method, the previous `external_args_buffer`
243    /// is read and an updated size is read from it to correct the value.
244    size_or_external_args_offset: u32,
245}
246
247#[derive(Debug, Copy, Clone)]
248struct OutputBufferOffsetsCursor {
249    /// Cursor pointing to the first free entry before the last method decoding, which allows
250    /// updating output sizes using [`OutputBufferDetails::size_or_external_args_offset`] field
251    before_last: usize,
252    /// Current cursor pointing to the next free entry in `output_buffer_details`
253    current: usize,
254}
255
256/// Decoder for transaction payload created using `TransactionPayloadBuilder`.
257#[derive(Debug)]
258pub struct TransactionPayloadDecoder<'a> {
259    payload: &'a [u8],
260    external_args_buffer: &'a mut [*mut c_void; EXTERNAL_ARGS_BUFFER_SIZE],
261    // TODO: Cast `output_buffer` into `&'a mut [MaybeUninit<u8>]` and remove
262    //  `output_buffer_cursor`
263    output_buffer: &'a mut [MaybeUninit<u128>],
264    output_buffer_cursor: usize,
265    output_buffer_details: &'a mut [MaybeUninit<OutputBufferDetails>],
266    output_buffer_offsets_cursor: OutputBufferOffsetsCursor,
267    map_context: fn(TransactionMethodContext) -> MethodContext,
268}
269
270impl<'a> TransactionPayloadDecoder<'a> {
271    /// Create a new instance.
272    ///
273    /// The size of `external_args_buffer` defines the max number of bytes allocated for
274    /// `ExternalArgs`, which impacts the number of arguments that can be represented by
275    /// `ExternalArgs`. The size is specified in pointers with `#[slot]` argument using one
276    /// pointer, `#[input]` two pointers, and `#[output]` three pointers each.
277    ///
278    /// The size of `output_buffer` defines how big the total size of `#[output]` and return values
279    /// could be in all methods of the payload together.
280    ///
281    /// The size of `output_buffer_details` defines how many `#[output]` arguments and return values
282    /// could exist in all methods of the payload together.
283    #[inline]
284    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
285    pub fn new(
286        payload: &'a [u128],
287        external_args_buffer: &'a mut [*mut c_void; EXTERNAL_ARGS_BUFFER_SIZE],
288        output_buffer: &'a mut [MaybeUninit<u128>],
289        output_buffer_details: &'a mut [MaybeUninit<OutputBufferDetails>],
290        map_context: fn(TransactionMethodContext) -> MethodContext,
291    ) -> Self {
292        debug_assert_eq!(align_of_val(payload), usize::from(MAX_ALIGNMENT));
293        debug_assert_eq!(align_of_val(output_buffer), usize::from(MAX_ALIGNMENT));
294
295        // SAFETY: Memory is valid and bound by an argument's lifetime
296        let payload =
297            unsafe { slice::from_raw_parts(payload.as_ptr().cast::<u8>(), size_of_val(payload)) };
298
299        Self {
300            payload,
301            external_args_buffer,
302            output_buffer,
303            output_buffer_cursor: 0,
304            output_buffer_details,
305            output_buffer_offsets_cursor: OutputBufferOffsetsCursor {
306                before_last: 0,
307                current: 0,
308            },
309            map_context,
310        }
311    }
312}
313
314impl<'a> TransactionPayloadDecoder<'a> {
315    /// Decode the next method (if any) in the payload
316    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
317    pub fn decode_next_method(
318        &mut self,
319    ) -> Result<Option<PreparedMethod<'_>>, TransactionPayloadDecoderError> {
320        TransactionPayloadDecoderInternal::<true>(self).decode_next_method()
321    }
322
323    /// Decode the next method (if any) in the payload without checking size.
324    ///
325    /// # Safety
326    /// Must be used with trusted input created using `TransactionPayloadBuilder` or pre-verified
327    /// using [`Self::decode_next_method()`] earlier.
328    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
329    pub unsafe fn decode_next_method_unchecked(&mut self) -> Option<PreparedMethod<'_>> {
330        TransactionPayloadDecoderInternal::<false>(self)
331            .decode_next_method()
332            .expect("No decoding errors are possible with trusted input; qed")
333    }
334}
335
336/// Write to an external arguments pointer and move it forward.
337///
338/// # Safety
339/// `external_args` must have enough capacity for the written value, and the current offset must
340/// have the correct alignment for the type being written.
341#[inline(always)]
342unsafe fn write_external_args<T>(external_args: &mut NonNull<c_void>, value: T) {
343    // SAFETY: guaranteed by this function signature
344    unsafe {
345        external_args.cast::<T>().write(value);
346        *external_args = external_args.byte_add(size_of::<T>());
347    }
348}
349
350/// # Safety
351/// When `VERIFY == false` input must be trusted and created using `TransactionPayloadBuilder` or
352/// pre-verified using `VERIFY == true` earlier.
353struct TransactionPayloadDecoderInternal<'tmp, 'decoder, const VERIFY: bool>(
354    &'tmp mut TransactionPayloadDecoder<'decoder>,
355);
356
357impl<'tmp, 'decoder, const VERIFY: bool> Deref
358    for TransactionPayloadDecoderInternal<'tmp, 'decoder, VERIFY>
359{
360    type Target = TransactionPayloadDecoder<'decoder>;
361
362    #[inline(always)]
363    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
364    fn deref(&self) -> &Self::Target {
365        self.0
366    }
367}
368
369impl<'tmp, 'decoder, const VERIFY: bool> DerefMut
370    for TransactionPayloadDecoderInternal<'tmp, 'decoder, VERIFY>
371{
372    #[inline(always)]
373    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
374    fn deref_mut(&mut self) -> &mut Self::Target {
375        self.0
376    }
377}
378
379impl<'tmp, 'decoder, const VERIFY: bool> TransactionPayloadDecoderInternal<'tmp, 'decoder, VERIFY> {
380    #[inline(always)]
381    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
382    fn decode_next_method(
383        mut self,
384    ) -> Result<Option<PreparedMethod<'decoder>>, TransactionPayloadDecoderError> {
385        if self.payload.len() <= usize::from(MAX_ALIGNMENT) {
386            return Ok(None);
387        }
388
389        self.update_output_buffer_details();
390
391        let contract = self.get_trivial_type::<Address>()?;
392        let method_fingerprint = self.get_trivial_type::<MethodFingerprint>()?;
393        let method_context =
394            (self.map_context)(*self.get_trivial_type::<TransactionMethodContext>()?);
395
396        let mut transaction_slots_inputs =
397            [MaybeUninit::<u8>::uninit(); MAX_TOTAL_METHOD_ARGS as usize];
398
399        let num_slot_arguments = self.read_u8()?;
400        for transaction_slot in transaction_slots_inputs
401            .iter_mut()
402            .take(usize::from(num_slot_arguments))
403        {
404            transaction_slot.write(self.read_u8()?);
405        }
406
407        let num_input_arguments = self.read_u8()?;
408        for transaction_input in transaction_slots_inputs
409            .iter_mut()
410            .skip(usize::from(num_slot_arguments))
411            .take(usize::from(num_input_arguments))
412        {
413            transaction_input.write(self.read_u8()?);
414        }
415
416        let num_output_arguments = self.read_u8()?;
417
418        // SAFETY: Just initialized elements above
419        let (transaction_slots, transaction_inputs) = unsafe {
420            let (transaction_slots, transaction_inputs) =
421                transaction_slots_inputs.split_at_unchecked(usize::from(num_slot_arguments));
422            let transaction_inputs =
423                transaction_inputs.get_unchecked(..usize::from(num_input_arguments));
424
425            (
426                transaction_slots.assume_init_ref(),
427                transaction_inputs.assume_init_ref(),
428            )
429        };
430
431        // This can be off by 1 due to `self` not included in `ExternalArgs`, but it is good enough
432        // for this context
433        let number_of_arguments = num_slot_arguments
434            .saturating_add(num_input_arguments)
435            .saturating_add(num_output_arguments);
436        if VERIFY && number_of_arguments > MAX_TOTAL_METHOD_ARGS {
437            return Err(TransactionPayloadDecoderError::TooManyArguments(
438                number_of_arguments,
439            ));
440        }
441
442        let external_args = NonNull::new(self.external_args_buffer.as_mut_ptr())
443            .expect("Not null; qed")
444            .cast::<c_void>();
445        {
446            let external_args_cursor = &mut external_args.clone();
447
448            for &transaction_slot in transaction_slots {
449                let address = match TransactionSlot::from_u8(transaction_slot).slot_type() {
450                    TransactionSlotType::Address => self.get_trivial_type::<Address>()?,
451                    TransactionSlotType::OutputIndex { output_index } => {
452                        let (bytes, size) = self.get_from_output_buffer(output_index)?;
453
454                        if VERIFY && size != Address::SIZE {
455                            return Err(
456                                TransactionPayloadDecoderError::InvalidSlotOutputIndexSize {
457                                    output_index,
458                                    size,
459                                },
460                            );
461                        }
462
463                        // SAFETY: All bytes are valid as long as alignment and size match
464                        let maybe_address = unsafe { Address::from_bytes(bytes) };
465
466                        if VERIFY {
467                            let Some(address) = maybe_address else {
468                                let error =
469                                    TransactionPayloadDecoderError::InvalidSlotOutputIndexAlign {
470                                        output_index,
471                                        alignment: align_of_val(bytes) as u32,
472                                    };
473                                return Err(error);
474                            };
475
476                            address
477                        } else {
478                            // SAFETY: The unverified version, see struct description
479                            unsafe { maybe_address.unwrap_unchecked() }
480                        }
481                    }
482                };
483
484                // SAFETY: Size of `self.external_args_buffer` is statically sized for worst-case,
485                // address is correctly aligned
486                unsafe {
487                    write_external_args(external_args_cursor, ptr::from_ref(address));
488                }
489            }
490
491            for &transaction_input in transaction_inputs {
492                let (bytes, size) = match TransactionInput::from_u8(transaction_input).input_type()
493                {
494                    TransactionInputType::Value { alignment_power } => {
495                        // Optimized version of the following:
496                        // let alignment = 2usize.pow(u32::from(alignment_power));
497                        let alignment = if VERIFY {
498                            1_usize.checked_shl(u32::from(alignment_power)).ok_or(
499                                TransactionPayloadDecoderError::AlignmentPowerTooLarge(
500                                    alignment_power,
501                                ),
502                            )?
503                        } else {
504                            // SAFETY: The unverified version, see struct description
505                            unsafe { 1_usize.unchecked_shl(u32::from(alignment_power)) }
506                        };
507
508                        let size = *self.get_trivial_type::<u32>()?;
509                        let bytes = self.get_bytes(
510                            size,
511                            NonZeroUsize::new(alignment).expect("Not zero; qed"),
512                        )?;
513
514                        (bytes, size)
515                    }
516                    TransactionInputType::OutputIndex { output_index } => {
517                        self.get_from_output_buffer(output_index)?
518                    }
519                };
520
521                // SAFETY: Size of `self.external_args_buffer` is statically sized for worst-case,
522                // buffer is correctly aligned
523                unsafe {
524                    write_external_args(
525                        external_args_cursor,
526                        FfiDataSizeCapacityRo {
527                            data_ptr: NonNull::from_ref(bytes).as_non_null_ptr(),
528                            size,
529                            capacity: size,
530                        },
531                    );
532                }
533            }
534
535            for _ in 0..num_output_arguments {
536                let recommended_capacity = *self.get_trivial_type::<u32>()?;
537                let alignment_power = *self.get_trivial_type::<u8>()?;
538                // Optimized version of the following:
539                // let alignment = 2usize.pow(u32::from(alignment_power));
540                let alignment = if VERIFY {
541                    1_usize.checked_shl(u32::from(alignment_power)).ok_or(
542                        TransactionPayloadDecoderError::AlignmentPowerTooLarge(alignment_power),
543                    )?
544                } else {
545                    // SAFETY: The unverified version, see struct description
546                    unsafe { 1_usize.unchecked_shl(u32::from(alignment_power)) }
547                };
548
549                // SAFETY: `external_args_cursor` is created from `external_args` and is within the
550                // same allocation
551                let external_args_size_offset = unsafe {
552                    external_args_cursor
553                        .byte_add(offset_of!(FfiDataSizeCapacityRw, size))
554                        .byte_offset_from_unsigned(external_args)
555                };
556
557                let data = self.allocate_output_buffer(
558                    recommended_capacity,
559                    NonZeroUsize::new(alignment).expect("Not zero; qed"),
560                    external_args_size_offset as u32,
561                )?;
562
563                // SAFETY: Size of `self.external_args_buffer` is statically sized for worst-case,
564                // buffer is correctly aligned
565                unsafe {
566                    write_external_args(
567                        external_args_cursor,
568                        FfiDataSizeCapacityRw {
569                            data_ptr: data.as_ptr(),
570                            size: 0,
571                            capacity: recommended_capacity,
572                        },
573                    );
574                }
575            }
576        }
577
578        Ok(Some(PreparedMethod {
579            contract: *contract,
580            fingerprint: *method_fingerprint,
581            external_args,
582            method_context,
583            phantom: PhantomData,
584        }))
585    }
586
587    /// Get a reference to a [`TrivialType`] value inside the payload
588    #[inline(always)]
589    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
590    fn get_trivial_type<T>(&mut self) -> Result<&'decoder T, TransactionPayloadDecoderError>
591    where
592        T: TrivialType,
593    {
594        self.ensure_alignment(NonZeroUsize::new(align_of::<T>()).expect("Not zero; qed"))?;
595
596        let bytes;
597        if VERIFY {
598            (bytes, self.payload) = self
599                .payload
600                .split_at_checked(size_of::<T>())
601                .ok_or(TransactionPayloadDecoderError::PayloadTooSmall)?;
602        } else {
603            // SAFETY: The unverified version, see struct description
604            (bytes, self.payload) = unsafe { self.payload.split_at_unchecked(size_of::<T>()) };
605        }
606
607        // SAFETY: Correctly aligned bytes of the correct size
608        let value_ref = unsafe { bytes.as_ptr().cast::<T>().as_ref().expect("Not null; qed") };
609
610        Ok(value_ref)
611    }
612
613    /// Get a reference to opaque bytes inside the payload
614    #[inline(always)]
615    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
616    fn get_bytes(
617        &mut self,
618        size: u32,
619        alignment: NonZeroUsize,
620    ) -> Result<&'decoder [u8], TransactionPayloadDecoderError> {
621        self.ensure_alignment(alignment)?;
622
623        let bytes;
624        if VERIFY {
625            (bytes, self.payload) = self
626                .payload
627                .split_at_checked(size as usize)
628                .ok_or(TransactionPayloadDecoderError::PayloadTooSmall)?;
629        } else {
630            // SAFETY: The unverified version, see struct description
631            (bytes, self.payload) = unsafe { self.payload.split_at_unchecked(size as usize) };
632        }
633
634        Ok(bytes)
635    }
636
637    #[inline(always)]
638    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
639    fn read_u8(&mut self) -> Result<u8, TransactionPayloadDecoderError> {
640        let value;
641        if VERIFY {
642            (value, self.payload) = self
643                .payload
644                .split_at_checked(1)
645                .ok_or(TransactionPayloadDecoderError::PayloadTooSmall)?;
646        } else {
647            // SAFETY: The unverified version, see struct description
648            (value, self.payload) = unsafe { self.payload.split_at_unchecked(1) };
649        }
650
651        Ok(value[0])
652    }
653
654    #[inline(always)]
655    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
656    fn ensure_alignment(
657        &mut self,
658        alignment: NonZeroUsize,
659    ) -> Result<(), TransactionPayloadDecoderError> {
660        let alignment = alignment.get();
661        debug_assert!(alignment <= usize::from(MAX_ALIGNMENT));
662
663        // Optimized version of the following that expects `alignment` to be a power of 2:
664        // let unaligned_by = self.payload.as_ptr().addr() % alignment;
665        let unaligned_by = self.payload.as_ptr().addr() & (alignment - 1);
666        if unaligned_by > 0 {
667            // SAFETY: Subtracted value is always smaller than alignment
668            let padding_bytes = unsafe { alignment.unchecked_sub(unaligned_by) };
669            if VERIFY {
670                self.payload = self
671                    .payload
672                    .split_off(padding_bytes..)
673                    .ok_or(TransactionPayloadDecoderError::PayloadTooSmall)?;
674            } else {
675                // SAFETY: Subtracted value is always smaller than alignment
676                self.payload = unsafe { self.payload.get_unchecked(padding_bytes..) };
677            }
678        }
679        Ok(())
680    }
681
682    /// Returns a tuple of `(size_ptr, output_ptr)` of a newly allocated output buffer
683    #[inline(always)]
684    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
685    fn allocate_output_buffer(
686        &mut self,
687        capacity: u32,
688        output_alignment: NonZeroUsize,
689        external_args_size_offset: u32,
690    ) -> Result<NonNull<u8>, TransactionPayloadDecoderError> {
691        if VERIFY && self.output_buffer_details.len() == self.output_buffer_offsets_cursor.current {
692            return Err(TransactionPayloadDecoderError::OutputBufferOffsetsTooSmall);
693        }
694
695        let (output_offset, output_ptr) = self
696            .allocate_output_buffer_ptr(output_alignment, capacity as usize)
697            .ok_or(TransactionPayloadDecoderError::OutputBufferTooSmall)?;
698
699        // SAFETY: Checked above that `output_buffer_details` is not full
700        let output_buffer_details = unsafe {
701            // Borrow checker doesn't understand that these variables are disjoint
702            let output_buffer_offsets_cursor = self.output_buffer_offsets_cursor.current;
703            self.output_buffer_details
704                .get_unchecked_mut(output_buffer_offsets_cursor)
705        };
706        output_buffer_details.write(OutputBufferDetails {
707            output_offset: output_offset as u32,
708            size_or_external_args_offset: external_args_size_offset,
709        });
710        self.output_buffer_offsets_cursor.current += 1;
711
712        Ok(output_ptr)
713    }
714
715    /// Returns `None` if output buffer is not large enough
716    #[inline(always)]
717    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
718    fn allocate_output_buffer_ptr<T>(
719        &mut self,
720        alignment: NonZeroUsize,
721        size: usize,
722    ) -> Option<(usize, NonNull<T>)> {
723        let alignment = alignment.get();
724        debug_assert!(alignment <= usize::from(MAX_ALIGNMENT));
725
726        // Optimized version of the following that expects `alignment` to be a power of 2:
727        // let unaligned_by = self.output_buffer_cursor % alignment;
728        let unaligned_by = self.output_buffer_cursor & (alignment - 1);
729        // SAFETY: Subtracted value is always smaller than alignment
730        let padding_bytes = unsafe { alignment.unchecked_sub(unaligned_by) };
731
732        let new_output_buffer_cursor = if VERIFY {
733            let new_output_buffer_cursor = self
734                .output_buffer_cursor
735                .checked_add(padding_bytes)?
736                .checked_add(size)?;
737
738            if new_output_buffer_cursor > size_of_val(self.output_buffer) {
739                return None;
740            }
741
742            new_output_buffer_cursor
743        } else {
744            // SAFETY: The unverified version, see struct description
745            unsafe {
746                self.output_buffer_cursor
747                    .unchecked_add(padding_bytes)
748                    .unchecked_add(size)
749            }
750        };
751
752        // SAFETY: Bounds and alignment checks are done above
753        let (offset, buffer_ptr) = unsafe {
754            let offset = self.output_buffer_cursor.unchecked_add(padding_bytes);
755            let buffer_ptr = NonNull::new_unchecked(
756                self.output_buffer.as_mut_ptr().byte_add(offset).cast::<T>(),
757            );
758
759            (offset, buffer_ptr)
760        };
761        self.output_buffer_cursor = new_output_buffer_cursor;
762
763        Some((offset, buffer_ptr))
764    }
765
766    /// Update all [`OutputBufferDetails::size_or_external_args_offset`] to store sizes rather than
767    /// offsets into `external_args_buffer` and advances [`OutputBufferOffsetsCursor::before_last`]
768    #[inline(always)]
769    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
770    fn update_output_buffer_details(&mut self) {
771        let TransactionPayloadDecoder {
772            external_args_buffer,
773            output_buffer_details,
774            output_buffer_offsets_cursor,
775            ..
776        } = &mut **self;
777
778        // SAFETY: Elements in the selected range were initialized
779        for output_buffer_details in unsafe {
780            output_buffer_details
781                .get_unchecked_mut(
782                    output_buffer_offsets_cursor.before_last..output_buffer_offsets_cursor.current,
783                )
784                .assume_init_mut()
785        } {
786            // SAFETY: Protected invariant in `decode_next_method`
787            output_buffer_details.size_or_external_args_offset = unsafe {
788                external_args_buffer
789                    .as_ptr()
790                    .cast::<u32>()
791                    .byte_add(output_buffer_details.size_or_external_args_offset as usize)
792                    .read()
793            };
794        }
795
796        output_buffer_offsets_cursor.before_last = output_buffer_offsets_cursor.current;
797    }
798
799    #[inline(always)]
800    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
801    fn get_from_output_buffer(
802        &self,
803        output_index: u8,
804    ) -> Result<(&[u8], u32), TransactionPayloadDecoderError> {
805        let OutputBufferDetails {
806            output_offset,
807            size_or_external_args_offset: size,
808        } = if VERIFY {
809            if usize::from(output_index) < self.output_buffer_offsets_cursor.current {
810                // SAFETY: Checked that index is initialized
811                unsafe {
812                    self.output_buffer_details
813                        .get_unchecked(usize::from(output_index))
814                        .assume_init()
815                }
816            } else {
817                return Err(TransactionPayloadDecoderError::OutputIndexNotFound(
818                    output_index,
819                ));
820            }
821        } else {
822            // SAFETY: The unverified version, see struct description
823            unsafe {
824                self.output_buffer_details
825                    .get_unchecked(usize::from(output_index))
826                    .assume_init()
827            }
828        };
829
830        // SAFETY: Offset was created as the result of writing value at the correct
831        // offset into `output_buffer_details` earlier
832        let bytes = unsafe {
833            let bytes_ptr = self
834                .output_buffer
835                .as_ptr()
836                .cast::<u8>()
837                .add(output_offset as usize);
838
839            slice::from_raw_parts(bytes_ptr, size as usize)
840        };
841
842        Ok((bytes, size))
843    }
844}