Add fast path entrypoint

Author: febo

p-token/src/entrypoint.rs

@@ -1,21 +1,218 @@
 use {
     crate::processor::*,
+    core::{
+        mem::{transmute, MaybeUninit},
+        slice::from_raw_parts,
+    },
     pinocchio::{
         account_info::AccountInfo,
-        no_allocator, nostd_panic_handler, program_entrypoint,
+        entrypoint::deserialize_into,
+        hint::likely,
+        no_allocator, nostd_panic_handler,
         program_error::{ProgramError, ToStr},
-        pubkey::Pubkey,
-        ProgramResult,
+        ProgramResult, MAX_TX_ACCOUNTS, SUCCESS,
     },
     pinocchio_token_interface::error::TokenError,
 };
 
-program_entrypoint!(process_instruction);
 // Do not allocate memory.
 no_allocator!();
 // Use the no_std panic handler.
 nostd_panic_handler!();
 
+/// Custom program entrypoint to give priority to `transfer` and
+/// `transfer_checked` instructions.
+///
+/// The entrypoint prioritizes the transfer instruction by validating
+/// account data lengths and instruction data. When it can reliably
+/// determine that the instruction is a transfer, it will invoke the
+/// processor directly.
+#[no_mangle]
+#[allow(clippy::arithmetic_side_effects)]
+pub unsafe extern "C" fn entrypoint(input: *mut u8) -> u64 {
+    // Constants that apply to both `transfer` and `transfer_checked`.
+
+    /// Offset for the first account.
+    const ACCOUNT1_HEADER_OFFSET: usize = 0x0008;
+
+    /// Offset for the first account data length. This is
+    /// expected to be a token account (165 bytes).
+    const ACCOUNT1_DATA_LEN: usize = 0x0058;
+
+    /// Offset for the second account.
+    const ACCOUNT2_HEADER_OFFSET: usize = 0x2910;
+
+    /// Offset for the second account data length. This is
+    /// expected to be a token account for `transfer` (165 bytes)
+    /// or a mint account for `transfer_checked` (82 bytes).
+    const ACCOUNT2_DATA_LEN: usize = 0x2960;
+
+    // Constants that apply to `transfer_checked` (instruction 12).
+
+    /// Offset for the third account.
+    const IX12_ACCOUNT3_HEADER_OFFSET: usize = 0x51c8;
+
+    /// Offset for the third account data length. This is
+    /// expected to be a token account (165 bytes).
+    const IX12_ACCOUNT3_DATA_LEN: usize = 0x5218;
+
+    /// Offset for the fourth account.
+    const IX12_ACCOUNT4_HEADER_OFFSET: usize = 0x7ad0;
+
+    /// Offset for the fourth account data length.
+    ///
+    /// This is expected to be an account with variable data
+    /// length.
+    const IX12_ACCOUNT4_DATA_LEN: usize = 0x7b20;
+
+    /// Expected offset for the instruction data in the case all
+    /// previous accounts have zero data.
+    ///
+    /// This value is adjusted before it is used.
+    const IX12_EXPECTED_INSTRUCTION_DATA_LEN_OFFSET: usize = 0xa330;
+
+    // Constants that apply to `transfer` (instruction 3).
+
+    /// Offset for the second account.
+    ///
+    /// Note that this assumes that both first and second accounts
+    /// have zero data, which is being validated before the offset
+    /// is used.
+    const IX3_ACCOUNT3_HEADER_OFFSET: usize = 0x5218;
+
+    /// Offset for the third account data length. This is
+    /// expected to be a mint account (82 bytes).
+    const IX3_ACCOUNT3_DATA_LEN: usize = 0x5268;
+
+    /// Expected offset for the instruction data in the case all
+    /// previous accounts have zero data.
+    ///
+    /// This value is adjusted before it is used.
+    const IX3_INSTRUCTION_DATA_LEN_OFFSET: usize = 0x7a28;
+
+    /// Align an address to the next multiple of 8.
+    #[inline(always)]
+    fn align(input: u64) -> u64 {
+        (input + 7) & (!7)
+    }
+
+    // Fast path for `transfer_checked`.
+    //
+    // It expects 4 accounts:
+    //   1. source: must be a token account (165 length)
+    //   2. mint: must be a mint account (82 length)
+    //   3. destination: must be a token account (165 length)
+    //   4. authority: can be any account (variable length)
+    //
+    // Instruction data is expected to be at least 9 bytes
+    // and discriminator equal to 12.
+    if *input == 4
+        && (*input.add(ACCOUNT1_DATA_LEN).cast::<u64>() == 165)
+        && (*input.add(ACCOUNT2_HEADER_OFFSET) == 255)
+        && (*input.add(ACCOUNT2_DATA_LEN).cast::<u64>() == 82)
+        && (*input.add(IX12_ACCOUNT3_HEADER_OFFSET) == 255)
+        && (*input.add(IX12_ACCOUNT3_DATA_LEN).cast::<u64>() == 165)
+        && (*input.add(IX12_ACCOUNT4_HEADER_OFFSET) == 255)
+    {
+        // The `authority` account can have variable data length.
+        let account_4_data_len_aligned =
+            align(*input.add(IX12_ACCOUNT4_DATA_LEN).cast::<u64>()) as usize;
+        let offset = IX12_EXPECTED_INSTRUCTION_DATA_LEN_OFFSET + account_4_data_len_aligned;
+
+        // Check that we have enough instruction data.
+        //
+        // Expected: instruction discriminator (u8) + amount (u64) + decimals (u8)
+        if input.add(offset).cast::<usize>().read() >= 10 {
+            let discriminator = input.add(offset + size_of::<u64>()).cast::<u8>().read();
+
+            // Check for transfer discriminator.
+            if likely(discriminator == 12) {
+                // instruction data length (u64) + discriminator (u8)
+                let instruction_data = unsafe { from_raw_parts(input.add(offset + 9), 9) };
+
+                let accounts = unsafe {
+                    [
+                        transmute::<*mut u8, AccountInfo>(input.add(ACCOUNT1_HEADER_OFFSET)),
+                        transmute::<*mut u8, AccountInfo>(input.add(ACCOUNT2_HEADER_OFFSET)),
+                        transmute::<*mut u8, AccountInfo>(input.add(IX12_ACCOUNT3_HEADER_OFFSET)),
+                        transmute::<*mut u8, AccountInfo>(input.add(IX12_ACCOUNT4_HEADER_OFFSET)),
+                    ]
+                };
+
+                return match process_transfer_checked(&accounts, instruction_data) {
+                    Ok(()) => SUCCESS,
+                    Err(error) => {
+                        log_error(&error);
+                        error.into()
+                    }
+                };
+            }
+        }
+    }
+    // Fast path for `transfer`.
+    //
+    // It expects 3 accounts:
+    //   1. source: must be a token account (165 length)
+    //   2. destination: must be a token account (165 length)
+    //   3. authority: can be any account (variable length)
+    //
+    // Instruction data is expected to be at least 8 bytes
+    // and discriminator equal to 3.
+    else if *input == 3
+        && (*input.add(ACCOUNT1_DATA_LEN).cast::<u64>() == 165)
+        && (*input.add(ACCOUNT2_HEADER_OFFSET) == 255)
+        && (*input.add(ACCOUNT2_DATA_LEN).cast::<u64>() == 165)
+        && (*input.add(IX3_ACCOUNT3_HEADER_OFFSET) == 255)
+    {
+        // The `authority` account can have variable data length.
+        let account_3_data_len_aligned =
+            align(*input.add(IX3_ACCOUNT3_DATA_LEN).cast::<u64>()) as usize;
+        let offset = IX3_INSTRUCTION_DATA_LEN_OFFSET + account_3_data_len_aligned;
+
+        // Check that we have enough instruction data.
+        if likely(input.add(offset).cast::<usize>().read() >= 9) {
+            let discriminator = input.add(offset + size_of::<u64>()).cast::<u8>().read();
+
+            // Check for transfer discriminator.
+            if likely(discriminator == 3) {
+                let instruction_data =
+                    unsafe { from_raw_parts(input.add(offset + 9), size_of::<u64>()) };
+
+                let accounts = unsafe {
+                    [
+                        transmute::<*mut u8, AccountInfo>(input.add(ACCOUNT1_HEADER_OFFSET)),
+                        transmute::<*mut u8, AccountInfo>(input.add(ACCOUNT2_HEADER_OFFSET)),
+                        transmute::<*mut u8, AccountInfo>(input.add(IX3_ACCOUNT3_HEADER_OFFSET)),
+                    ]
+                };
+
+                return match process_transfer(&accounts, instruction_data) {
+                    Ok(()) => SUCCESS,
+                    Err(error) => {
+                        log_error(&error);
+                        error.into()
+                    }
+                };
+            }
+        }
+    }
+
+    // Entrypoint for the remaining instructions.
+
+    const UNINIT: MaybeUninit<AccountInfo> = MaybeUninit::<AccountInfo>::uninit();
+    let mut accounts = [UNINIT; { MAX_TX_ACCOUNTS }];
+
+    let (count, instruction_data) = deserialize_into(input, &mut accounts);
+
+    match process_instruction(
+        from_raw_parts(accounts.as_ptr() as _, count as usize),
+        instruction_data,
+    ) {
+        Ok(()) => SUCCESS,
+        Err(error) => error.into(),
+    }
+}
+
 /// Log an error.
 #[cold]
 fn log_error(error: &ProgramError) {
@@ -30,11 +227,7 @@ fn log_error(error: &ProgramError) {
 /// instructions, since it is not sound to have a "batch" instruction inside
 /// another "batch" instruction.
 #[inline(always)]
-pub fn process_instruction(
-    _program_id: &Pubkey,
-    accounts: &[AccountInfo],
-    instruction_data: &[u8],
-) -> ProgramResult {
+pub fn process_instruction(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
     let [discriminator, remaining @ ..] = instruction_data else {
         return Err(TokenError::InvalidInstruction.into());
     };