Reduce batch compute units

Author: febo

program/src/entrypoint.rs

@@ -1,4 +1,3 @@
-use batch::process_batch;
 use pinocchio::{
     account_info::AccountInfo, default_panic_handler, no_allocator, program_entrypoint,
     program_error::ProgramError, pubkey::Pubkey, ProgramResult,
@@ -12,6 +11,27 @@ no_allocator!();
 // Use the default panic handler.
 default_panic_handler!();
 
+#[inline(always)]
+pub fn process_instruction(
+    _program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    instruction_data: &[u8],
+) -> ProgramResult {
+    let [discriminator, instruction_data @ ..] = instruction_data else {
+        return Err(ProgramError::InvalidInstructionData);
+    };
+
+    if *discriminator == 255 {
+        // 255 - Batch
+        #[cfg(feature = "logging")]
+        pinocchio::msg!("Instruction: Batch");
+
+        return process_batch(accounts, instruction_data);
+    }
+
+    inner_process_instruction(accounts, instruction_data, *discriminator)
+}
+
 /// Process an instruction.
 ///
 /// The processor of the token program is divided into two parts to reduce the overhead
@@ -22,32 +42,35 @@ default_panic_handler!();
 ///
 /// Instructions on the first part of the processor:
 ///
-/// - `0`: `InitializeMint`
-/// - `3`:  `Transfer`
-/// - `7`:  `MintTo`
-/// - `9`:  `CloseAccount`
+/// -  `0`: `InitializeMint`
+/// -  `1`: `InitializeAccount`
+/// -  `3`: `Transfer`
+/// -  `7`: `MintTo`
+/// -  `9`: `CloseAccount`
+/// - `18`: `InitializeAccount2`
 /// - `18`: `InitializeAccount3`
 /// - `20`: `InitializeMint2`
-/// - `255`: `Batch`
 #[inline(always)]
-pub fn process_instruction(
-    _program_id: &Pubkey,
+pub fn inner_process_instruction(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
+    discriminator: u8,
 ) -> ProgramResult {
-    let [discriminator, instruction_data @ ..] = instruction_data else {
-        return Err(ProgramError::InvalidInstructionData);
-    };
-
-    match *discriminator {
+    match discriminator {
         // 0 - InitializeMint
         0 => {
             #[cfg(feature = "logging")]
             pinocchio::msg!("Instruction: InitializeMint");
 
             process_initialize_mint(accounts, instruction_data, true)
         }
+        // 1 - InitializeAccount
+        1 => {
+            #[cfg(feature = "logging")]
+            pinocchio::msg!("Instruction: InitializeAccount");
 
+            process_initialize_account(accounts)
+        }
         // 3 - Transfer
         3 => {
             #[cfg(feature = "logging")]
@@ -69,6 +92,13 @@ pub fn process_instruction(
 
             process_close_account(accounts)
         }
+        // 16 - InitializeAccount2
+        16 => {
+            #[cfg(feature = "logging")]
+            pinocchio::msg!("Instruction: InitializeAccount2");
+
+            process_initialize_account2(accounts, instruction_data)
+        }
         // 18 - InitializeAccount3
         18 => {
             #[cfg(feature = "logging")]
@@ -83,14 +113,7 @@ pub fn process_instruction(
 
             process_initialize_mint2(accounts, instruction_data)
         }
-        // 255 - Batch
-        255 => {
-            #[cfg(feature = "logging")]
-            pinocchio::msg!("Instruction: Batch");
-
-            process_batch(accounts, instruction_data)
-        }
-        _ => process_remaining_instruction(accounts, instruction_data, *discriminator),
+        _ => inner_process_remaining_instruction(accounts, instruction_data, discriminator),
     }
 }
 
@@ -99,19 +122,12 @@ pub fn process_instruction(
 /// This function is called by the `process_instruction` function if the discriminator
 /// does not match any of the common instructions. This function is used to reduce the
 /// overhead of having a large `match` statement in the `process_instruction` function.
-fn process_remaining_instruction(
+fn inner_process_remaining_instruction(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
     discriminator: u8,
 ) -> ProgramResult {
     match discriminator {
-        // 1 - InitializeAccount
-        1 => {
-            #[cfg(feature = "logging")]
-            pinocchio::msg!("Instruction: InitializeAccount");
-
-            process_initialize_account(accounts)
-        }
         // 2 - InitializeMultisig
         2 => {
             #[cfg(feature = "logging")]
@@ -189,13 +205,6 @@ fn process_remaining_instruction(
 
             process_burn_checked(accounts, instruction_data)
         }
-        // 16 - InitializeAccount2
-        16 => {
-            #[cfg(feature = "logging")]
-            pinocchio::msg!("Instruction: InitializeAccount2");
-
-            process_initialize_account2(accounts, instruction_data)
-        }
         // 17 - SyncNative
         17 => {
             #[cfg(feature = "logging")]

program/src/processor/batch.rs

@@ -1,6 +1,6 @@
 use pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult};
 
-use crate::entrypoint::process_instruction;
+use crate::entrypoint::inner_process_instruction;
 
 /// The size of the batch instruction header.
 ///
@@ -18,11 +18,12 @@ pub fn process_batch(mut accounts: &[AccountInfo], mut instruction_data: &[u8])
             return Err(ProgramError::InvalidInstructionData);
         }
 
-        // SAFETY: The instruction data is guaranteed to have at least two bytes.
+        // SAFETY: The instruction data is guaranteed to have at least two bytes (header)
+        // + one byte (discriminator).
         let expected_accounts = unsafe { *instruction_data.get_unchecked(0) as usize };
         let data_offset = IX_HEADER_SIZE + unsafe { *instruction_data.get_unchecked(1) as usize };
 
-        if instruction_data.len() < data_offset {
+        if instruction_data.len() < data_offset || data_offset == 0 {
             return Err(ProgramError::InvalidInstructionData);
         }
 
@@ -32,10 +33,12 @@ pub fn process_batch(mut accounts: &[AccountInfo], mut instruction_data: &[u8])
 
         // Process the instruction.
 
-        process_instruction(
-            &token_interface::program::ID,
-            &accounts[..expected_accounts],
-            &instruction_data[IX_HEADER_SIZE..data_offset],
+        // SAFETY: The instruction data and accounts lengths are already validated so all
+        // the slices are guaranteed to be valid.
+        inner_process_instruction(
+            unsafe { accounts.get_unchecked(..expected_accounts) },
+            unsafe { instruction_data.get_unchecked(IX_HEADER_SIZE + 1..data_offset) },
+            unsafe { *instruction_data.get_unchecked(IX_HEADER_SIZE) },
         )?;
 
         if data_offset == instruction_data.len() {

program/src/processor/mod.rs

@@ -50,6 +50,7 @@ pub mod shared;
 pub use amount_to_ui_amount::process_amount_to_ui_amount;
 pub use approve::process_approve;
 pub use approve_checked::process_approve_checked;
+pub use batch::process_batch;
 pub use burn::process_burn;
 pub use burn_checked::process_burn_checked;
 pub use close_account::process_close_account;