pubkey: Optional bump parameter (#196)

* Remove separate bump parameter

* Add optional bump

* Tweak docs

* Tweaks

* More tweaks

* Update sdk/pubkey/src/lib.rs

Co-authored-by: Illia Bobyr <[email protected]>

---------

Co-authored-by: Illia Bobyr <[email protected]>

Author: febo

sdk/pubkey/src/lib.rs

@@ -14,68 +14,77 @@ use pinocchio::pubkey::{Pubkey, MAX_SEEDS, PDA_MARKER};
 #[cfg(target_os = "solana")]
 use pinocchio::syscalls::sol_sha256;
 
-/// Derive a [program address][pda] from the given seeds, bump and program id.
+/// Derive a [program address][pda] from the given seeds, optional bump and
+/// program id.
 ///
 /// [pda]: https://solana.com/docs/core/pda
 ///
-/// This function avoids the cost of the `create_program_address` syscall,
-/// which is `1500` compute units, by directly computing the derived address
-/// calculating the hash of the seeds, bump, and program id using the
-/// `sol_sha256` syscall.
+/// In general, the derivation uses an optional bump (byte) value to ensure a
+/// valid PDA (off-curve) is generated. Even when a program stores a bump to
+/// derive a program address, it is necessary to use the
+/// [`pinocchio::pubkey::create_program_address`] to validate the derivation. In
+/// most cases, the program has the correct seeds for the derivation, so it would
+/// be sufficient to just perform the derivation and compare it against the
+/// expected resulting address.
 ///
-/// Even when a program stores a bump to derive a program address, it is necessary
-/// to use the [`pinocchio::pubkey::create_program_address`] to validate the
-/// derivation. In most cases, the program has the correct seeds for the derivation,
-/// so it would be sufficient to just perform the derivation and compare it against
-/// the expected resulting address.
+/// This function avoids the cost of the `create_program_address` syscall
+/// (`1500` compute units) by directly computing the derived address
+/// calculating the hash of the seeds, bump and program id using the
+/// `sol_sha256` syscall.
 ///
 /// # Important
 ///
-/// This function differs from [`pinocchio::pubkey::create_program_address`] in that it
-/// does not perform a validation to ensure that the derived address is a valid
+/// This function differs from [`pinocchio::pubkey::create_program_address`] in that
+/// it does not perform a validation to ensure that the derived address is a valid
 /// (off-curve) program derived address. It is intended for use in cases where the
 /// seeds, bump, and program id are known to be valid, and the caller wants to derive
 /// the address without incurring the cost of the `create_program_address` syscall.
-pub fn derive_address<const N: usize>(seeds: &[&[u8]; N], bump: u8, program_id: &Pubkey) -> Pubkey {
+pub fn derive_address<const N: usize>(
+    seeds: &[&[u8]; N],
+    bump: Option<u8>,
+    program_id: &Pubkey,
+) -> Pubkey {
     const {
-        assert!(
-            N <= MAX_SEEDS,
-            "number of seeds must be less than MAX_SEEDS"
-        );
+        assert!(N < MAX_SEEDS, "number of seeds must be less than MAX_SEEDS");
     }
 
     const UNINIT: MaybeUninit<&[u8]> = MaybeUninit::<&[u8]>::uninit();
-    let mut data = [UNINIT; MAX_SEEDS + 3];
+    let mut data = [UNINIT; MAX_SEEDS + 2];
     let mut i = 0;
 
     while i < N {
-        // SAFETY: `data` is guanranteed to have enough space for `N` seeds,
+        // SAFETY: `data` is guaranteed to have enough space for `N` seeds,
         // so `i` will always be within bounds.
         unsafe {
             data.get_unchecked_mut(i).write(seeds.get_unchecked(i));
         }
         i += 1;
     }
 
-    let bump = [bump];
+    // TODO: replace this with `as_slice` when the MSRV is upgraded
+    // to `1.84.0+`.
+    let bump_seed = [bump.unwrap_or_default()];
 
-    // SAFETY: `data` is guaranteed to have enough space for `MAX_SEEDS + 3`
+    // SAFETY: `data` is guaranteed to have enough space for `MAX_SEEDS + 2`
     // elements, and `MAX_SEEDS` is as large as `N`.
     unsafe {
-        data.get_unchecked_mut(i).write(bump.as_ref());
-        data.get_unchecked_mut(i + 1).write(program_id.as_ref());
-        data.get_unchecked_mut(i + 2).write(PDA_MARKER.as_ref());
+        if bump.is_some() {
+            data.get_unchecked_mut(i).write(&bump_seed);
+            i += 1;
+        }
+        data.get_unchecked_mut(i).write(program_id.as_ref());
+        data.get_unchecked_mut(i + 1).write(PDA_MARKER.as_ref());
     }
 
     #[cfg(target_os = "solana")]
     {
         let mut pda = MaybeUninit::<[u8; 32]>::uninit();
 
-        // SAFETY: `data` has `i + 3` elements initialized.
+        // SAFETY: `data` has `i + 2` elements initialized.
         unsafe {
             sol_sha256(
                 data.as_ptr() as *const u8,
-                (N + 3) as u64,
+                (i + 2) as u64,
                 pda.as_mut_ptr() as *mut u8,
             );
         }
@@ -88,20 +97,24 @@ pub fn derive_address<const N: usize>(seeds: &[&[u8]; N], bump: u8, program_id:
     unreachable!("deriving a pda is only available on target `solana`");
 }
 
-/// Derive a [program address][pda] from the given seeds, bump and program id.
+/// Derive a [program address][pda] from the given seeds, optional bump and
+/// program id.
 ///
 /// [pda]: https://solana.com/docs/core/pda
 ///
-/// This function avoids the cost of the `create_program_address` syscall,
-/// which is `1500` compute units, by directly computing the derived address
-/// using the SHA-256 hash of the seeds, bump, and program id.
+/// In general, the derivation uses an optional bump (byte) value to ensure a
+/// valid PDA (off-curve) is generated.
 ///
-/// This function is intended for use in `const` contexts.
+/// This function is intended for use in `const` contexts - i.e., the seeds and
+/// bump are known at compile time and the program id is also a constant. It avoids
+/// the cost of the `create_program_address` syscall (`1500` compute units) by
+/// directly computing the derived address using the SHA-256 hash of the seeds,
+/// bump and program id.
 ///
 /// # Important
 ///
-/// This function differs from [`pinocchio::pubkey::create_program_address`] in that it
-/// does not perform a validation to ensure that the derived address is a valid
+/// This function differs from [`pinocchio::pubkey::create_program_address`] in that
+/// it does not perform a validation to ensure that the derived address is a valid
 /// (off-curve) program derived address. It is intended for use in cases where the
 /// seeds, bump, and program id are known to be valid, and the caller wants to derive
 /// the address without incurring the cost of the `create_program_address` syscall.
@@ -110,14 +123,11 @@ pub fn derive_address<const N: usize>(seeds: &[&[u8]; N], bump: u8, program_id:
 #[cfg(feature = "const")]
 pub const fn derive_address_const<const N: usize>(
     seeds: &[&[u8]; N],
-    bump: u8,
+    bump: Option<u8>,
     program_id: &Pubkey,
 ) -> Pubkey {
     const {
-        assert!(
-            N <= MAX_SEEDS,
-            "number of seeds must be less than MAX_SEEDS"
-        );
+        assert!(N < MAX_SEEDS, "number of seeds must be less than MAX_SEEDS");
     }
 
     let mut hasher = Sha256::new();
@@ -128,13 +138,17 @@ pub const fn derive_address_const<const N: usize>(
         i += 1;
     }
 
-    let bump = [bump];
-
-    hasher
-        .update(&bump)
-        .update(program_id)
-        .update(PDA_MARKER)
-        .finalize()
+    // TODO: replace this with `is_some` when the MSRV is upgraded
+    // to `1.84.0+`.
+    if let Some(bump) = bump {
+        hasher
+            .update(&[bump])
+            .update(program_id)
+            .update(PDA_MARKER)
+            .finalize()
+    } else {
+        hasher.update(program_id).update(PDA_MARKER).finalize()
+    }
 }
 
 /// Convenience macro to define a static `Pubkey` value.