feat: optimize undelegate instruction using Pinocchio for lower CU consumption

Cargo.toml

@@ -29,6 +29,7 @@ name = "dlp"
 no-entrypoint = []
 default = ["solana-security-txt"]
 unit_test_config = []
+log-cost = []
 
 [dependencies]
 borsh = { version = "1.5.3", features = [ "derive" ] }

src/cu.rs

@@ -0,0 +1,41 @@
+use pinocchio::syscalls::sol_remaining_compute_units;
+use pinocchio_log::log;
+
+pub struct BenchmarkComputeUnit {
+    name: &'static str,
+    remaining_at_start: u64,
+}
+
+impl BenchmarkComputeUnit {
+    pub fn start(name: &'static str) -> BenchmarkComputeUnit {
+        log!("BENCHMARK BEGIN: [{}]", name);
+        Self {
+            name,
+            remaining_at_start: Self::remaining_cu(),
+        }
+    }
+
+    fn remaining_cu() -> u64 {
+        unsafe { sol_remaining_compute_units() }
+    }
+}
+
+impl Drop for BenchmarkComputeUnit {
+    fn drop(&mut self) {
+        let consumed = self.remaining_at_start - Self::remaining_cu();
+        log!(
+            "BENCHMARK END: [{}] consumed {} of {} compute units.",
+            self.name,
+            consumed,
+            self.remaining_at_start
+        )
+    }
+}
+
+#[macro_export]
+macro_rules! compute {
+    ($msg:expr=> $($tt:tt)*) => {
+        let _log = $crate::cu::BenchmarkComputeUnit::start($msg);
+        $($tt)*
+    };
+}

src/lib.rs

@@ -22,6 +22,9 @@ pub mod pda;
 mod processor;
 pub mod state;
 
+#[cfg(feature = "log-cost")]
+mod cu;
+
 declare_id!("DELeGGvXpWV2fqJUhqcF5ZSYMS4JTLjteaAMARRSaeSh");
 
 pub mod fast {
@@ -106,6 +109,9 @@ pub fn fast_process_instruction(
         discriminator::DlpDiscriminator::Delegate => Some(processor::fast::process_delegate(
             program_id, accounts, data,
         )),
+        discriminator::DlpDiscriminator::Undelegate => Some(processor::fast::process_undelegate(
+            program_id, accounts, data,
+        )),
         _ => None,
     }
 }

src/processor/fast/mod.rs

@@ -1,7 +1,9 @@
 mod delegate;
+mod undelegate;
 mod utils;
 
 pub use delegate::*;
+pub use undelegate::*;
 
 pub fn to_pinocchio_program_error(
     error: solana_program::program_error::ProgramError,

src/processor/fast/undelegate.rs

@@ -0,0 +1,311 @@
+use pinocchio::{
+    account_info::AccountInfo,
+    cpi::invoke_signed,
+    instruction::{AccountMeta, Instruction, Signer},
+    program_error::ProgramError,
+    pubkey::{pubkey_eq, Pubkey},
+    sysvars::{rent::Rent, Sysvar},
+    ProgramResult,
+};
+use pinocchio::{pubkey, seeds};
+use pinocchio_log::log;
+use pinocchio_system::instructions as system;
+
+use crate::consts::{EXTERNAL_UNDELEGATE_DISCRIMINATOR, RENT_FEES_PERCENTAGE};
+use crate::error::DlpError;
+use crate::pda;
+use crate::processor::fast::utils::{
+    pda::{close_pda, close_pda_with_fees, create_pda},
+    requires::require_uninitialized_pda,
+};
+use crate::state::{DelegationMetadata, DelegationRecord};
+
+#[cfg(feature = "log-cost")]
+use crate::compute;
+
+use super::{
+    to_pinocchio_program_error,
+    utils::requires::{
+        require_initialized_delegation_metadata, require_initialized_delegation_record,
+        require_initialized_protocol_fees_vault, require_initialized_validator_fees_vault,
+        require_owned_pda, require_program, require_signer,
+    },
+};
+
+pub fn process_undelegate(
+    _program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    _data: &[u8],
+) -> ProgramResult {
+    let [validator, delegated_account, owner_program, undelegate_buffer_account, commit_state_account, commit_record_account, delegation_record_account, delegation_metadata_account, rent_reimbursement, fees_vault, validator_fees_vault, system_program] =
+        accounts
+    else {
+        return Err(ProgramError::NotEnoughAccountKeys);
+    };
+
+    // Check accounts
+    require_signer(validator, "validator")?;
+    require_owned_pda(delegated_account, &crate::fast::ID, "delegated account")?;
+    require_initialized_delegation_record(delegated_account, delegation_record_account, true)?;
+    require_initialized_delegation_metadata(delegated_account, delegation_metadata_account, true)?;
+    require_initialized_protocol_fees_vault(fees_vault, true)?;
+    require_initialized_validator_fees_vault(validator, validator_fees_vault, true)?;
+    require_program(system_program, &pinocchio_system::ID, "system program")?;
+
+    // Make sure there is no pending commits to be finalized before this call
+    require_uninitialized_pda(
+        commit_state_account,
+        &[pda::COMMIT_STATE_TAG, delegated_account.key()],
+        &crate::fast::ID,
+        false,
+        "commit state",
+    )?;
+    require_uninitialized_pda(
+        commit_record_account,
+        &[pda::COMMIT_RECORD_TAG, delegated_account.key()],
+        &crate::fast::ID,
+        false,
+        "commit record",
+    )?;
+
+    // Load delegation record
+    let delegation_record_data = delegation_record_account.try_borrow_data()?;
+    let delegation_record =
+        DelegationRecord::try_from_bytes_with_discriminator(&delegation_record_data)
+            .map_err(to_pinocchio_program_error)?;
+
+    // Check passed owner and owner stored in the delegation record match
+    if !pubkey_eq(delegation_record.owner.as_array(), owner_program.key()) {
+        log!("Expected delegation record owner to be : ");
+        pubkey::log(delegation_record.owner.as_array());
+        log!("but got : ");
+        pubkey::log(owner_program.key());
+        return Err(ProgramError::InvalidAccountOwner);
+    }
+
+    // Load delegated account metadata
+    let delegation_metadata_data = delegation_metadata_account.try_borrow_data()?;
+    let delegation_metadata =
+        DelegationMetadata::try_from_bytes_with_discriminator(&delegation_metadata_data)
+            .map_err(to_pinocchio_program_error)?;
+
+    // Check if the delegated account is undelegatable
+    if !delegation_metadata.is_undelegatable {
+        log!("delegation metadata indicates the account is not undelegatable : ");
+        pubkey::log(delegation_metadata_account.key());
+        return Err(DlpError::NotUndelegatable.into());
+    }
+
+    // Check if the rent payer is correct
+    if !pubkey_eq(
+        delegation_metadata.rent_payer.as_array(),
+        rent_reimbursement.key(),
+    ) {
+        log!("Expected rent payer to be : ");
+        pubkey::log(delegation_metadata.rent_payer.as_array());
+        log!("but got : ");
+        pubkey::log(rent_reimbursement.key());
+        return Err(DlpError::InvalidReimbursementAddressForDelegationRent.into());
+    }
+
+    // Dropping delegation references
+    drop(delegation_record_data);
+    drop(delegation_metadata_data);
+
+    // If there is no program to call CPI to, we can just assign the owner back and we're done
+    if delegated_account.data_is_empty() {
+        // TODO - we could also do this fast-path if the data was non-empty but zeroed-out
+        unsafe {
+            delegated_account.assign(owner_program.key());
+        }
+        process_delegation_cleanup(
+            delegation_record_account,
+            delegation_metadata_account,
+            rent_reimbursement,
+            fees_vault,
+            validator_fees_vault,
+        )?;
+        return Ok(());
+    }
+
+    // Initialize the undelegation buffer PDA
+
+    let undelegate_buffer_bump: u8 = require_uninitialized_pda(
+        undelegate_buffer_account,
+        &[pda::UNDELEGATE_BUFFER_TAG, delegated_account.key()],
+        &crate::fast::ID,
+        true,
+        "undelegate buffer",
+    )?;
+
+    create_pda(
+        undelegate_buffer_account,
+        &crate::fast::ID,
+        delegated_account.data_len(),
+        &[Signer::from(&seeds!(
+            pda::UNDELEGATE_BUFFER_TAG,
+            delegated_account.key(),
+            &[undelegate_buffer_bump]
+        ))],
+        validator,
+    )?;
+
+    // Copy data in the undelegation buffer PDA
+    (*undelegate_buffer_account.try_borrow_mut_data()?)
+        .copy_from_slice(&delegated_account.try_borrow_data()?);
+
+    // Call a CPI to the owner program to give it back the new state
+    process_undelegation_with_cpi(
+        validator,
+        delegated_account,
+        owner_program,
+        undelegate_buffer_account,
+        &[Signer::from(&seeds!(
+            pda::UNDELEGATE_BUFFER_TAG,
+            delegated_account.key(),
+            &[undelegate_buffer_bump]
+        ))],
+        delegation_metadata,
+        system_program,
+    )?;
+
+    // Done, close undelegation buffer
+    close_pda(undelegate_buffer_account, validator)?;
+
+    // Closing delegation accounts
+    process_delegation_cleanup(
+        delegation_record_account,
+        delegation_metadata_account,
+        rent_reimbursement,
+        fees_vault,
+        validator_fees_vault,
+    )?;
+    Ok(())
+}
+
+/// 1. Close the delegated account
+/// 2. CPI to the owner program
+/// 3. Check state
+/// 4. Settle lamports balance
+#[allow(clippy::too_many_arguments)]
+fn process_undelegation_with_cpi(
+    validator: &AccountInfo,
+    delegated_account: &AccountInfo,
+    owner_program: &AccountInfo,
+    undelegate_buffer_account: &AccountInfo,
+    undelegate_buffer_signer_seeds: &[Signer],
+    delegation_metadata: DelegationMetadata,
+    system_program: &AccountInfo,
+) -> ProgramResult {
+    let delegated_account_lamports_before_close = delegated_account.lamports();
+    close_pda(delegated_account, validator)?;
+
+    // Invoke the owner program's post-undelegation IX, to give the state back to the original program
+    let validator_lamports_before_cpi = validator.lamports();
+
+    cpi_external_undelegate(
+        validator,
+        delegated_account,
+        undelegate_buffer_account,
+        undelegate_buffer_signer_seeds,
+        system_program,
+        owner_program.key(),
+        delegation_metadata,
+    )?;
+
+    let validator_lamports_after_cpi = validator.lamports();
+
+    // Check that the validator lamports are exactly as expected
+    let delegated_account_min_rent = Rent::get()?.minimum_balance(delegated_account.data_len());
+    if validator_lamports_before_cpi
+        != validator_lamports_after_cpi
+            .checked_add(delegated_account_min_rent)
+            .ok_or(DlpError::Overflow)?
+    {
+        return Err(DlpError::InvalidValidatorBalanceAfterCPI.into());
+    }
+
+    // Check that the owner program properly moved the state back into the original account during CPI
+    if delegated_account.try_borrow_data()?.as_ref()
+        != undelegate_buffer_account.try_borrow_data()?.as_ref()
+    {
+        return Err(DlpError::InvalidAccountDataAfterCPI.into());
+    }
+
+    // Return the extra lamports to the delegated account
+    let delegated_account_extra_lamports = delegated_account_lamports_before_close
+        .checked_sub(delegated_account_min_rent)
+        .ok_or(DlpError::Overflow)?;
+
+    system::Transfer {
+        from: validator,
+        to: delegated_account,
+        lamports: delegated_account_extra_lamports,
+    }
+    .invoke()?;
+    Ok(())
+}
+
+/// CPI to the original owner program to re-open the PDA with the new state
+fn cpi_external_undelegate(
+    payer: &AccountInfo,
+    delegated_account: &AccountInfo,
+    undelegate_buffer_account: &AccountInfo,
+    undelegate_buffer_signer_seeds: &[Signer],
+    system_program: &AccountInfo,
+    owner_program_id: &Pubkey,
+    delegation_metadata: DelegationMetadata,
+) -> ProgramResult {
+    let data = {
+        // GAIN: 299  (42075 => 41776)
+        let mut data = Vec::with_capacity(32);
+        data.extend_from_slice(&EXTERNAL_UNDELEGATE_DISCRIMINATOR);
+        borsh::to_writer(&mut data, &delegation_metadata.seeds)
+            .map_err(|_| ProgramError::BorshIoError)?;
+        data
+    };
+
+    let external_undelegate_instruction = Instruction {
+        program_id: owner_program_id,
+        data: &data,
+        accounts: &[
+            AccountMeta::new(delegated_account.key(), true, false),
+            AccountMeta::new(undelegate_buffer_account.key(), true, true),
+            AccountMeta::new(payer.key(), true, true),
+            AccountMeta::new(system_program.key(), false, false),
+        ],
+    };
+
+    invoke_signed(
+        &external_undelegate_instruction,
+        &[
+            delegated_account,
+            undelegate_buffer_account,
+            payer,
+            system_program,
+        ],
+        undelegate_buffer_signer_seeds,
+    )
+}
+
+fn process_delegation_cleanup(
+    delegation_record_account: &AccountInfo,
+    delegation_metadata_account: &AccountInfo,
+    rent_reimbursement: &AccountInfo,
+    fees_vault: &AccountInfo,
+    validator_fees_vault: &AccountInfo,
+) -> ProgramResult {
+    close_pda_with_fees(
+        delegation_record_account,
+        rent_reimbursement,
+        &[validator_fees_vault, fees_vault],
+        RENT_FEES_PERCENTAGE,
+    )?;
+    close_pda_with_fees(
+        delegation_metadata_account,
+        rent_reimbursement,
+        &[validator_fees_vault, fees_vault],
+        RENT_FEES_PERCENTAGE,
+    )?;
+    Ok(())
+}