feat: Add CommitDiff Instruction for Efficient Account Updates (#110)

* feat: Add CommitDiff Instruction for Efficient Account Updates

* Make CI happy

Author: snawaz

Cargo.lock

@@ -48,6 +48,7 @@ pinocchio-pubkey = "0.3.0"
 pinocchio-system = "0.3.0"
 rkyv = "0.7.45"
 static_assertions = "1.1.0"
+strum = { version = "0.27.2", features = ["derive"] }
 
 [dev-dependencies]
 base64 = "0.22.1"

Cargo.toml

@@ -1,3 +1,5 @@
+use std::mem::size_of;
+
 use borsh::{BorshDeserialize, BorshSerialize};
 
 #[derive(Default, Debug, BorshSerialize, BorshDeserialize)]
@@ -23,3 +25,37 @@ pub struct CommitStateFromBufferArgs {
     /// Whether the account can be undelegated after the commit completes
     pub allow_undelegation: bool,
 }
+
+#[derive(Default, Debug, BorshSerialize)]
+pub struct CommitDiffArgs {
+    /// The account diff
+    /// SAFETY: this must be the FIRST field in the struct because the serialized format
+    /// is manually split: the diff (with Borsh Vec prefix) followed by the fixed-size
+    /// fields. The processor uses `data.split_at(data.len() - SIZE_COMMIT_DIFF_ARGS_WITHOUT_DIFF)`
+    /// to separate them during deserialization.
+    pub diff: Vec<u8>,
+
+    /// "Nonce" of an account. Updates are submitted historically and nonce incremented by 1
+    /// Deprecated: The ephemeral slot at which the account data is committed
+    pub nonce: u64,
+
+    /// The lamports that the account holds in the ephemeral validator
+    pub lamports: u64,
+
+    /// Whether the account can be undelegated after the commit completes
+    pub allow_undelegation: bool,
+}
+
+#[derive(Default, Debug, BorshDeserialize)]
+pub struct CommitDiffArgsWithoutDiff {
+    /// "Nonce" of an account. Updates are submitted historically and nonce incremented by 1
+    /// Deprecated: The ephemeral slot at which the account data is committed
+    pub nonce: u64,
+    /// The lamports that the account holds in the ephemeral validator
+    pub lamports: u64,
+    /// Whether the account can be undelegated after the commit completes
+    pub allow_undelegation: bool,
+}
+
+pub const SIZE_COMMIT_DIFF_ARGS_WITHOUT_DIFF: usize =
+    size_of::<u64>() + size_of::<u64>() + size_of::<bool>();

src/args/commit_state.rs

@@ -103,6 +103,13 @@ impl<'a> DiffSet<'a> {
         Ok(this)
     }
 
+    pub fn try_new_from_borsh_vec(vec_buffer: &'a [u8]) -> Result<Self, ProgramError> {
+        if vec_buffer.len() < 4 {
+            return Err(ProgramError::InvalidInstructionData);
+        }
+        Self::try_new(&vec_buffer[4..])
+    }
+
     pub fn raw_diff(&self) -> &'a [u8] {
         // SAFETY: it does not do any "computation" as such. It merely reverses try_new
         // and get the immutable slice back.

src/diff/types.rs

@@ -1,7 +1,8 @@
 use num_enum::TryFromPrimitive;
+use strum::IntoStaticStr;
 
 #[repr(u8)]
-#[derive(Clone, Copy, Debug, Eq, PartialEq, TryFromPrimitive)]
+#[derive(Clone, Copy, Debug, Eq, PartialEq, TryFromPrimitive, IntoStaticStr)]
 #[rustfmt::skip]
 pub enum DlpDiscriminator {
     /// See [crate::processor::process_delegate] for docs.
@@ -34,11 +35,17 @@ pub enum DlpDiscriminator {
     CloseValidatorFeesVault = 14,
     /// See [crate::processor::process_call_handler] for docs.
     CallHandler = 15,
+    /// See [crate::processor::process_commit_diff] for docs.
+    CommitDiff = 16,
 }
 
 impl DlpDiscriminator {
     pub fn to_vec(self) -> Vec<u8> {
         let num = self as u64;
         num.to_le_bytes().to_vec()
     }
+
+    pub fn name(&self) -> &'static str {
+        self.into()
+    }
 }

src/discriminator.rs

@@ -0,0 +1,45 @@
+use borsh::to_vec;
+use solana_program::instruction::Instruction;
+use solana_program::system_program;
+use solana_program::{instruction::AccountMeta, pubkey::Pubkey};
+
+use crate::args::CommitDiffArgs;
+use crate::discriminator::DlpDiscriminator;
+use crate::pda::{
+    commit_record_pda_from_delegated_account, commit_state_pda_from_delegated_account,
+    delegation_metadata_pda_from_delegated_account, delegation_record_pda_from_delegated_account,
+    program_config_from_program_id, validator_fees_vault_pda_from_validator,
+};
+
+/// Builds a commit state instruction.
+/// See [crate::processor::fast::process_commit_diff] for docs.
+pub fn commit_diff(
+    validator: Pubkey,
+    delegated_account: Pubkey,
+    delegated_account_owner: Pubkey,
+    commit_args: CommitDiffArgs,
+) -> Instruction {
+    let commit_args = to_vec(&commit_args).unwrap();
+    let delegation_record_pda = delegation_record_pda_from_delegated_account(&delegated_account);
+    let commit_state_pda = commit_state_pda_from_delegated_account(&delegated_account);
+    let commit_record_pda = commit_record_pda_from_delegated_account(&delegated_account);
+    let validator_fees_vault_pda = validator_fees_vault_pda_from_validator(&validator);
+    let delegation_metadata_pda =
+        delegation_metadata_pda_from_delegated_account(&delegated_account);
+    let program_config_pda = program_config_from_program_id(&delegated_account_owner);
+    Instruction {
+        program_id: crate::id(),
+        accounts: vec![
+            AccountMeta::new_readonly(validator, true),
+            AccountMeta::new_readonly(delegated_account, false),
+            AccountMeta::new(commit_state_pda, false),
+            AccountMeta::new(commit_record_pda, false),
+            AccountMeta::new_readonly(delegation_record_pda, false),
+            AccountMeta::new(delegation_metadata_pda, false),
+            AccountMeta::new_readonly(validator_fees_vault_pda, false),
+            AccountMeta::new_readonly(program_config_pda, false),
+            AccountMeta::new_readonly(system_program::id(), false),
+        ],
+        data: [DlpDiscriminator::CommitDiff.to_vec(), commit_args].concat(),
+    }
+}

src/instruction_builder/commit_diff.rs

@@ -1,6 +1,7 @@
 mod call_handler;
 mod close_ephemeral_balance;
 mod close_validator_fees_vault;
+mod commit_diff;
 mod commit_state;
 mod commit_state_from_buffer;
 mod delegate;
@@ -17,6 +18,7 @@ mod whitelist_validator_for_program;
 pub use call_handler::*;
 pub use close_ephemeral_balance::*;
 pub use close_validator_fees_vault::*;
+pub use commit_diff::*;
 pub use commit_state::*;
 pub use commit_state_from_buffer::*;
 pub use delegate::*;

src/instruction_builder/mod.rs

@@ -82,6 +82,9 @@ pub fn fast_process_instruction(
         DlpDiscriminator::CommitStateFromBuffer => Some(
             processor::fast::process_commit_state_from_buffer(program_id, accounts, data),
         ),
+        DlpDiscriminator::CommitDiff => Some(processor::fast::process_commit_diff(
+            program_id, accounts, data,
+        )),
         DlpDiscriminator::Finalize => Some(processor::fast::process_finalize(
             program_id, accounts, data,
         )),

src/lib.rs

@@ -0,0 +1,95 @@
+use borsh::BorshDeserialize;
+use pinocchio::{
+    account_info::AccountInfo, program_error::ProgramError, pubkey::Pubkey, ProgramResult,
+};
+use pinocchio_log::log;
+
+use crate::args::{CommitDiffArgsWithoutDiff, SIZE_COMMIT_DIFF_ARGS_WITHOUT_DIFF};
+use crate::processor::fast::{process_commit_state_internal, CommitStateInternalArgs};
+use crate::{apply_diff_copy, DiffSet};
+
+/// Commit diff to a delegated PDA
+///
+/// Accounts:
+///
+/// 0: `[signer]`   the validator requesting the commit
+/// 1: `[]`         the delegated account
+/// 2: `[writable]` the PDA storing the new state
+/// 3: `[writable]` the PDA storing the commit record
+/// 4: `[]`         the delegation record
+/// 5: `[writable]` the delegation metadata
+/// 6: `[]`         the validator fees vault
+/// 7: `[]`         the program config account
+/// 8: `[]`         the system program
+///
+/// Requirements:
+///
+/// - The following accounts must be initialized:
+///   - delegation record
+///   - delegation metadata
+///   - validator fees vault
+///   - program config
+/// - The following accounts must be uninitialized:
+///   - commit state
+///   - commit record
+/// - delegated account holds at least the lamports indicated in the delegation record
+/// - account was not committed at a later slot
+///
+/// Steps:
+/// 1. Check that the pda is delegated
+/// 2. Init a new PDA to store the new state
+/// 3. Copy the new state to the new PDA
+/// 4. Init a new PDA to store the record of the new state commitment
+pub fn process_commit_diff(
+    _program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    data: &[u8],
+) -> ProgramResult {
+    let [validator, delegated_account, commit_state_account, commit_record_account, delegation_record_account, delegation_metadata_account, validator_fees_vault, program_config_account, _system_program] =
+        accounts
+    else {
+        return Err(ProgramError::NotEnoughAccountKeys);
+    };
+
+    if data.len() < SIZE_COMMIT_DIFF_ARGS_WITHOUT_DIFF {
+        return Err(ProgramError::InvalidInstructionData);
+    }
+
+    let (diff, data) = data.split_at(data.len() - SIZE_COMMIT_DIFF_ARGS_WITHOUT_DIFF);
+
+    let args =
+        CommitDiffArgsWithoutDiff::try_from_slice(data).map_err(|_| ProgramError::BorshIoError)?;
+
+    let diffset = DiffSet::try_new_from_borsh_vec(diff)?;
+
+    if diffset.segments_count() == 0 {
+        log!("WARN: noop; empty diff sent");
+    }
+
+    let commit_record_lamports = args.lamports;
+    let commit_record_nonce = args.nonce;
+    let allow_undelegation = args.allow_undelegation;
+
+    // TODO (snawaz): the following approach to apply diff works, but it's not efficient.
+    // It is also problematic for larger account as it allocates memory on the heap.
+    // It will be fixed in a separate PR.
+    let original = unsafe { delegated_account.borrow_data_unchecked() };
+    let changed = apply_diff_copy(original, &diffset)?;
+
+    let commit_args = CommitStateInternalArgs {
+        commit_state_bytes: &changed,
+        commit_record_lamports,
+        commit_record_nonce,
+        allow_undelegation,
+        validator,
+        delegated_account,
+        commit_state_account,
+        commit_record_account,
+        delegation_record_account,
+        delegation_metadata_account,
+        validator_fees_vault,
+        program_config_account,
+    };
+
+    process_commit_state_internal(commit_args)
+}

src/processor/fast/commit_diff.rs

@@ -1,10 +1,12 @@
+mod commit_diff;
 mod commit_state;
 mod commit_state_from_buffer;
 mod delegate;
 mod finalize;
 mod undelegate;
 mod utils;
 
+pub use commit_diff::*;
 pub use commit_state::*;
 pub use commit_state_from_buffer::*;
 pub use delegate::*;