pinocchio: Add pubkey_eq helper (#248)

* Add pubkey_eq helper

* Fix typo

* Update pubkey comparison

* Add proptest

* Add unlikely

* Replace proptest

Author: febo

sdk/pinocchio/src/account_info.rs

@@ -10,7 +10,11 @@ use core::{
 #[cfg(target_os = "solana")]
 use crate::syscalls::sol_memset_;
 
-use crate::{program_error::ProgramError, pubkey::Pubkey, ProgramResult};
+use crate::{
+    program_error::ProgramError,
+    pubkey::{pubkey_eq, Pubkey},
+    ProgramResult,
+};
 
 /// Maximum number of bytes a program may add to an account during a
 /// single top-level instruction.
@@ -186,7 +190,7 @@ impl AccountInfo {
     /// Checks if the account is owned by the given program.
     #[inline(always)]
     pub fn is_owned_by(&self, program: &Pubkey) -> bool {
-        self.owner() == program
+        pubkey_eq(self.owner(), program)
     }
 
     /// Changes the owner of the account.

sdk/pinocchio/src/cpi.rs

@@ -4,9 +4,10 @@ use core::{mem::MaybeUninit, ops::Deref, slice::from_raw_parts};
 
 use crate::{
     account_info::{AccountInfo, BorrowState},
+    hint::unlikely,
     instruction::{Account, Instruction, Signer},
     program_error::ProgramError,
-    pubkey::Pubkey,
+    pubkey::{pubkey_eq, Pubkey},
     ProgramResult,
 };
 
@@ -299,7 +300,7 @@ unsafe fn inner_invoke_signed_with_bounds<const MAX_ACCOUNTS: usize>(
             // In order to check whether the borrow state is compatible
             // with the invocation, we need to check that we have the
             // correct account info and meta pair.
-            if account_info.key() != account_meta.pubkey {
+            if unlikely(!pubkey_eq(account_info.key(), account_meta.pubkey)) {
                 return Err(ProgramError::InvalidArgument);
             }
 

sdk/pinocchio/src/pubkey.rs

@@ -1,5 +1,7 @@
 //! Public key type and functions.
 
+use core::ptr::read_unaligned;
+
 use crate::program_error::ProgramError;
 
 /// Number of bytes in a pubkey.
@@ -33,6 +35,24 @@ pub fn log(pubkey: &Pubkey) {
     core::hint::black_box(pubkey);
 }
 
+/// Compare two `Pubkey`s for equality.
+///
+/// The implementation of this function is currently more efficient
+/// than `p1 == p2` since it compares 8 bytes at a time instead of
+/// byte-by-byte.
+#[inline(always)]
+pub const fn pubkey_eq(p1: &Pubkey, p2: &Pubkey) -> bool {
+    let p1_ptr = p1.as_ptr() as *const u64;
+    let p2_ptr = p2.as_ptr() as *const u64;
+
+    unsafe {
+        read_unaligned(p1_ptr) == read_unaligned(p2_ptr)
+            && read_unaligned(p1_ptr.add(1)) == read_unaligned(p2_ptr.add(1))
+            && read_unaligned(p1_ptr.add(2)) == read_unaligned(p2_ptr.add(2))
+            && read_unaligned(p1_ptr.add(3)) == read_unaligned(p2_ptr.add(3))
+    }
+}
+
 /// Find a valid [program derived address][pda] and its corresponding bump seed.
 ///
 /// [pda]: https://solana.com/docs/core/cpi#program-derived-addresses
@@ -277,3 +297,25 @@ pub fn create_with_seed(
         panic!("create_with_seed is only available on target `solana`")
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::pubkey::{pubkey_eq, Pubkey, PUBKEY_BYTES};
+
+    #[test]
+    fn test_pubkey_eq_matches_default_eq() {
+        for i in 0..u8::MAX {
+            let p1: Pubkey = [i; PUBKEY_BYTES];
+            let p2: Pubkey = [i; PUBKEY_BYTES];
+
+            assert_eq!(pubkey_eq(&p1, &p2), p1 == p2);
+        }
+
+        for i in 0..u8::MAX {
+            let p1: Pubkey = [i; PUBKEY_BYTES];
+            let p2: Pubkey = [u8::MAX - i; PUBKEY_BYTES];
+
+            assert_eq!(!pubkey_eq(&p1, &p2), p1 != p2);
+        }
+    }
+}

sdk/pinocchio/src/sysvars/clock.rs

@@ -3,9 +3,10 @@
 use super::Sysvar;
 use crate::{
     account_info::{AccountInfo, Ref},
+    hint::unlikely,
     impl_sysvar_get,
     program_error::ProgramError,
-    pubkey::Pubkey,
+    pubkey::{pubkey_eq, Pubkey},
 };
 
 /// The ID of the clock sysvar.
@@ -86,7 +87,7 @@ impl Clock {
     /// This method performs a check on the account info key.
     #[inline]
     pub fn from_account_info(account_info: &AccountInfo) -> Result<Ref<Clock>, ProgramError> {
-        if account_info.key() != &CLOCK_ID {
+        if unlikely(!pubkey_eq(account_info.key(), &CLOCK_ID)) {
             return Err(ProgramError::InvalidArgument);
         }
         Ok(Ref::map(account_info.try_borrow_data()?, |data| unsafe {
@@ -107,7 +108,7 @@ impl Clock {
     pub unsafe fn from_account_info_unchecked(
         account_info: &AccountInfo,
     ) -> Result<&Self, ProgramError> {
-        if account_info.key() != &CLOCK_ID {
+        if unlikely(!pubkey_eq(account_info.key(), &CLOCK_ID)) {
             return Err(ProgramError::InvalidArgument);
         }
         Ok(Self::from_bytes_unchecked(

sdk/pinocchio/src/sysvars/rent.rs

@@ -5,9 +5,10 @@
 use super::Sysvar;
 use crate::{
     account_info::{AccountInfo, Ref},
+    hint::unlikely,
     impl_sysvar_get,
     program_error::ProgramError,
-    pubkey::Pubkey,
+    pubkey::{pubkey_eq, Pubkey},
 };
 
 /// The ID of the rent sysvar.
@@ -73,7 +74,7 @@ impl Rent {
     /// This method performs a check on the account info key.
     #[inline]
     pub fn from_account_info(account_info: &AccountInfo) -> Result<Ref<Rent>, ProgramError> {
-        if account_info.key() != &RENT_ID {
+        if unlikely(!pubkey_eq(account_info.key(), &RENT_ID)) {
             return Err(ProgramError::InvalidArgument);
         }
         Ok(Ref::map(account_info.try_borrow_data()?, |data| unsafe {
@@ -94,7 +95,7 @@ impl Rent {
     pub unsafe fn from_account_info_unchecked(
         account_info: &AccountInfo,
     ) -> Result<&Self, ProgramError> {
-        if account_info.key() != &RENT_ID {
+        if unlikely(!pubkey_eq(account_info.key(), &RENT_ID)) {
             return Err(ProgramError::InvalidArgument);
         }
         Ok(Self::from_bytes_unchecked(

sdk/pinocchio/src/sysvars/slot_hashes/mod.rs

@@ -15,8 +15,9 @@ mod test_utils;
 
 use crate::{
     account_info::{AccountInfo, Ref},
+    hint::unlikely,
     program_error::ProgramError,
-    pubkey::Pubkey,
+    pubkey::{pubkey_eq, Pubkey},
     sysvars::clock::Slot,
 };
 use core::{mem, ops::Deref, slice::from_raw_parts};
@@ -277,7 +278,7 @@ impl<'a> SlotHashes<Ref<'a, [u8]>> {
     /// - `ProgramError::AccountBorrowFailed` if the account data is already mutably borrowed
     #[inline(always)]
     pub fn from_account_info(account_info: &'a AccountInfo) -> Result<Self, ProgramError> {
-        if account_info.key() != &SLOTHASHES_ID {
+        if unlikely(!pubkey_eq(account_info.key(), &SLOTHASHES_ID)) {
             return Err(ProgramError::InvalidArgument);
         }
 

sdk/pubkey/src/lib.rs

@@ -4,7 +4,7 @@
 #[doc(hidden)]
 // Re-export dependencies used in macros.
 pub mod reexport {
-    pub use pinocchio::pubkey::Pubkey;
+    pub use pinocchio::pubkey::{pubkey_eq, Pubkey};
 }
 
 use core::mem::MaybeUninit;
@@ -176,7 +176,7 @@ macro_rules! declare_id {
         #[doc = "Returns `true` if given pubkey is the program ID."]
         #[inline]
         pub fn check_id(id: &$crate::reexport::Pubkey) -> bool {
-            id == &ID
+            $crate::reexport::pubkey_eq(id, &ID)
         }
 
         #[doc = "Returns the program ID."]