Skip to content

pinocchio: Add pubkey_eq helper #248

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Sep 11, 2025
Merged

pinocchio: Add pubkey_eq helper #248

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
b1daf04
Add pubkey_eq helper
febo Sep 10, 2025
e80ef09
Fix typo
febo Sep 10, 2025
9e857a9
Update pubkey comparison
febo Sep 11, 2025
efcab54
Add proptest
febo Sep 11, 2025
97f1d87
Add unlikely
febo Sep 11, 2025
5a3b61e
Replace proptest
febo Sep 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions sdk/pinocchio/src/account_info.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
5 changes: 3 additions & 2 deletions sdk/pinocchio/src/cpi.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand Down Expand Up @@ -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);
}

Expand Down
42 changes: 42 additions & 0 deletions sdk/pinocchio/src/pubkey.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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.
Comment on lines +40 to +42
Copy link
Contributor

@ilya-bobyr ilya-bobyr Sep 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If this is just a more efficient implementation of ==, it seems better to redefine PartialEq:

// Comparing `Pubkey`s 8 bytes at a time is faster, than comparing byte at a
// time.  One would expect LLVM to do this optimization for us, but it does not
// at the moment (as of platform tools v1.43).  Comparing keys is a pretty
// common operation, so we want to make it faster if we can.
impl PartialEq for Pubkey {
    fn eq(&self, other: &Self) -> bool {
        let p1 = self.as_ptr() as *const u64;
        let p2 = other.as_ptr() as *const u64;

        // SAFETY: `p1` and `p2` both point to `[u8; 32]` bytes of data, which
        // is the same size as `[u64; 4]`.
        unsafe {
            read_unaligned(p1) == read_unaligned(p2)
                && read_unaligned(p1.add(1)) == read_unaligned(p2.add(1))
                && read_unaligned(p1.add(2)) == read_unaligned(p2.add(2))
                && read_unaligned(p1.add(3)) == read_unaligned(p2.add(3))
        }
    }
}

Bikeshedding

Maybe it is possible to get the same effect if you just help the compiler a bit,
by casting [u8; 32] into [u64; 4]?

// Comparing `Pubkey`s 8 bytes at a time is faster than comparing byte at a
// time.  One would expect LLVM to do this optimization for us, but it does not
// at the moment (as of platform tools v1.43).  Comparing keys is a pretty
// common operation, so we want to make it faster if we can.
impl PartialEq for Pubkey {
    fn eq(&self, other: &Self) -> bool {
        let p1 = self.as_ptr() as *const [u64; 4];
        let p2 = other.as_ptr() as *const [u64; 4];

        // SAFETY: `p1` and `p2` both point to `[u8; 32]` bytes of data, which
        // is the same size as `[u64; 4]`.
        unsafe { read_unaligned(p1) == read_unaligned(p2) }
    }
}

I would also add static_assertions, to be able to encode the SAFETEY comment
in code:

use static_assertions::assert_eq_size;

// Comparing `Pubkey`s 8 bytes at a time is faster than comparing byte at a
// time.  One would expect LLVM to do this optimization for us, but it does not
// at the moment (as of platform tools v1.43).  Comparing keys is a pretty
// common operation, so we want to make it faster if we can.
impl PartialEq for Pubkey {
    fn eq(&self, other: &Self) -> bool {
        let p1 = self.as_ptr() as *const [u64; 4];
        let p2 = other.as_ptr() as *const [u64; 4];

        assert_eq_size!(*self, [u64; 4]);
        assert_eq_size!(*other, [u64; 4]);
        // SAFETY: `self` and `[u64; 4]` are of the same size.
        unsafe { read_unaligned(p1) == read_unaligned(p2) }
    }
}

(It is a compile time assertion with no runtime code.)

Copy link
Contributor

@rustopian rustopian Sep 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wanted this, too – at first. It would be great to get the savings everywhere devs use ==.

But since our pinocchio::pubkey::Pubkey is just an alias for [u8;32], we can't impl PartialEq on it. Creating a new type is maybe not worth it and probably a separate discussion.

Also, sounds like Lucas is looking into this, so maybe this discussion will be irrelevant soon.

Copy link
Collaborator Author

@febo febo Sep 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If this is just a more efficient implementation of ==, it seems better to redefine PartialEq:

We have done this for the solana-address here, since Address is a type.

Bikeshedding

Maybe it is possible to get the same effect if you just help the compiler a bit, by casting [u8; 32] into [u64; 4]?

Unfortunately no, it is even worse. 😅 I think the compiler tries to do something different when it sees an array, which in the end becomes "bloated" and uses more CU. You can use this simple bench to see the effect.

As @rustopian mentioned, Lucas will have a look at this to see if we can improve the code generation so the compiler does the "right" thing.

Copy link
Contributor

@ilya-bobyr ilya-bobyr Sep 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wanted this, too – at first. It would be great to get the savings everywhere devs use ==.

But since our pinocchio::pubkey::Pubkey is just an alias for [u8;32], we can't impl PartialEq on it. Creating a new type is maybe not worth it and probably a separate discussion.

This is out of scope for this PR obviously, but, considering Pubkey is such a fundamental abstraction, I think it should be a type.

Copy link
Collaborator Author

@febo febo Sep 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It will be very soon: #118 😊

#[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
Expand Down Expand Up @@ -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);
}
}
}
7 changes: 4 additions & 3 deletions sdk/pinocchio/src/sysvars/clock.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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 {
Expand All @@ -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(
Expand Down
7 changes: 4 additions & 3 deletions sdk/pinocchio/src/sysvars/rent.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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 {
Expand All @@ -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(
Expand Down
5 changes: 3 additions & 2 deletions sdk/pinocchio/src/sysvars/slot_hashes/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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);
}

Expand Down
4 changes: 2 additions & 2 deletions sdk/pubkey/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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."]
Expand Down