p-token: Add custom entrypoint

p-token/src/entrypoint.rs

@@ -1,25 +1,222 @@
 use {
     crate::processor::*,
+    core::{
+        mem::{size_of, transmute, MaybeUninit},
+        slice::from_raw_parts,
+    },
     pinocchio::{
         account_info::AccountInfo,
-        no_allocator, nostd_panic_handler, program_entrypoint,
+        entrypoint::deserialize,
+        log::sol_log,
+        no_allocator, nostd_panic_handler,
         program_error::{ProgramError, ToStr},
-        pubkey::Pubkey,
-        ProgramResult,
+        ProgramResult, MAX_TX_ACCOUNTS, SUCCESS,
     },
-    pinocchio_token_interface::error::TokenError,
+    pinocchio_token_interface::{error::TokenError, likely},
 };
 
-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 = 0x7a78;
+
+    /// 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(input, &mut accounts);
+
+    match process_instruction(
+        from_raw_parts(accounts.as_ptr() as _, count),
+        instruction_data,
+    ) {
+        Ok(()) => SUCCESS,
+        Err(error) => error.into(),
+    }
+}
+
 /// Log an error.
 #[cold]
 fn log_error(error: &ProgramError) {
-    pinocchio::log::sol_log(error.to_str::<TokenError>());
+    sol_log(error.to_str::<TokenError>());
 }
 
 /// Process an instruction.
@@ -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());
     };
@@ -138,12 +331,12 @@ pub(crate) fn inner_process_instruction(
 
             process_burn_checked(accounts, instruction_data)
         }
-        // 16 - InitializeAccount2
-        16 => {
+        // 17 - SyncNative
+        17 => {
             #[cfg(feature = "logging")]
-            pinocchio::msg!("Instruction: InitializeAccount2");
+            pinocchio::msg!("Instruction: SyncNative");
 
-            process_initialize_account2(accounts, instruction_data)
+            process_sync_native(accounts)
         }
         // 18 - InitializeAccount3
         18 => {
@@ -159,6 +352,13 @@ pub(crate) fn inner_process_instruction(
 
             process_initialize_mint2(accounts, instruction_data)
         }
+        // 22 - InitializeImmutableOwner
+        22 => {
+            #[cfg(feature = "logging")]
+            pinocchio::msg!("Instruction: InitializeImmutableOwner");
+
+            process_initialize_immutable_owner(accounts)
+        }
         d => inner_process_remaining_instruction(accounts, instruction_data, d),
     }
 }
@@ -231,12 +431,12 @@ fn inner_process_remaining_instruction(
 
             process_mint_to_checked(accounts, instruction_data)
         }
-        // 17 - SyncNative
-        17 => {
+        // 16 - InitializeAccount2
+        16 => {
             #[cfg(feature = "logging")]
-            pinocchio::msg!("Instruction: SyncNative");
+            pinocchio::msg!("Instruction: InitializeAccount2");
 
-            process_sync_native(accounts)
+            process_initialize_account2(accounts, instruction_data)
         }
         // 19 - InitializeMultisig2
         19 => {
@@ -252,13 +452,6 @@ fn inner_process_remaining_instruction(
 
             process_get_account_data_size(accounts)
         }
-        // 22 - InitializeImmutableOwner
-        22 => {
-            #[cfg(feature = "logging")]
-            pinocchio::msg!("Instruction: InitializeImmutableOwner");
-
-            process_initialize_immutable_owner(accounts)
-        }
         // 23 - AmountToUiAmount
         23 => {
             #[cfg(feature = "logging")]