From 383919beaa526a02ade48c45ac5badf6fd9926ee Mon Sep 17 00:00:00 2001 From: Sarfaraz Nawaz Date: Thu, 30 Oct 2025 15:37:08 +0530 Subject: [PATCH] feat: Optimize CommitDiff using efficient merge_diff_copy that avoids heap allocation --- src/diff/algorithm.rs | 54 +++++++++++++++++-- src/error.rs | 4 +- src/processor/fast/commit_diff.rs | 12 ++--- src/processor/fast/commit_state.rs | 32 ++++++++--- .../fast/commit_state_from_buffer.rs | 5 +- 5 files changed, 87 insertions(+), 20 deletions(-) diff --git a/src/diff/algorithm.rs b/src/diff/algorithm.rs index 43d41399..1d7af040 100644 --- a/src/diff/algorithm.rs +++ b/src/diff/algorithm.rs @@ -3,8 +3,11 @@ use std::cmp::{min, Ordering}; use pinocchio::program_error::ProgramError; use rkyv::util::AlignedVec; +use crate::error::DlpError; + use super::{ - DiffSet, SizeChanged, SIZE_OF_CHANGED_LEN, SIZE_OF_NUM_OFFSET_PAIRS, SIZE_OF_SINGLE_OFFSET_PAIR, + DiffSet, OffsetInData, SizeChanged, SIZE_OF_CHANGED_LEN, SIZE_OF_NUM_OFFSET_PAIRS, + SIZE_OF_SINGLE_OFFSET_PAIR, }; /// @@ -229,6 +232,35 @@ pub fn apply_diff_copy(original: &[u8], diffset: &DiffSet<'_>) -> Result }) } +/// This function constructs destination by merging original with diff such that destination +/// becomes the changed version of the original. +/// +/// Precondition: +/// - destination.len() == original.len() +pub fn merge_diff_copy( + destination: &mut [u8], + original: &[u8], + diffset: &DiffSet<'_>, +) -> Result<(), ProgramError> { + if destination.len() != original.len() { + return Err(DlpError::MergeDiffError.into()); + } + let mut write_index = 0; + for item in diffset.iter() { + let (diff_segment, OffsetInData { start, end }) = item?; + if write_index < start { + // copy the unchanged bytes + destination[write_index..start].copy_from_slice(&original[write_index..start]); + } + destination[start..end].copy_from_slice(diff_segment); + write_index = end; + } + if write_index < original.len() { + destination[write_index..].copy_from_slice(&original[write_index..]); + } + Ok(()) +} + // private function that does the actual work. fn apply_diff_impl(original: &mut [u8], diffset: &DiffSet<'_>) -> Result<(), ProgramError> { for item in diffset.iter() { @@ -247,7 +279,7 @@ mod tests { Rng, RngCore, SeedableRng, }; - use crate::{apply_diff_copy, apply_diff_in_place, compute_diff, DiffSet}; + use crate::{apply_diff_copy, apply_diff_in_place, compute_diff, merge_diff_copy, DiffSet}; #[test] fn test_no_change() { @@ -311,6 +343,14 @@ mod tests { let expected_changed = apply_diff_copy(&original, &actual_diffset).unwrap(); assert_eq!(changed.as_slice(), expected_changed.as_slice()); + + let expected_changed = { + let mut destination = vec![255; original.len()]; + merge_diff_copy(&mut destination, &original, &actual_diffset).unwrap(); + destination + }; + + assert_eq!(changed.as_slice(), expected_changed.as_slice()); } #[test] @@ -394,11 +434,19 @@ mod tests { // apply diff back to verify correctness let expected_changed = { - let mut copy = original; + let mut copy = original.clone(); apply_diff_in_place(&mut copy, &actual_diffset).unwrap(); copy }; assert_eq!(changed, expected_changed); + + let expected_changed = { + let mut destination = vec![255; original.len()]; + merge_diff_copy(&mut destination, &original, &actual_diffset).unwrap(); + destination + }; + + assert_eq!(changed, expected_changed); } } diff --git a/src/error.rs b/src/error.rs index 96e55dfd..1709eb61 100644 --- a/src/error.rs +++ b/src/error.rs @@ -35,10 +35,12 @@ pub enum DlpError { Overflow = 13, #[error("Too many seeds")] TooManySeeds = 14, - #[error("Invalid diff passed to DiffSet::try_new")] + #[error("Invalid length of diff passed to DiffSet::try_new")] InvalidDiff = 15, #[error("Diff is not properly aligned")] InvalidDiffAlignment = 16, + #[error("MergeDiff precondition did not meet")] + MergeDiffError = 17, } impl From for ProgramError { diff --git a/src/processor/fast/commit_diff.rs b/src/processor/fast/commit_diff.rs index 448f36fd..4bd2522e 100644 --- a/src/processor/fast/commit_diff.rs +++ b/src/processor/fast/commit_diff.rs @@ -6,7 +6,9 @@ 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}; +use crate::DiffSet; + +use super::NewState; /// Commit diff to a delegated PDA /// @@ -70,14 +72,8 @@ pub fn process_commit_diff( 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_state_bytes: NewState::Diff(diffset), commit_record_lamports, commit_record_nonce, allow_undelegation, diff --git a/src/processor/fast/commit_state.rs b/src/processor/fast/commit_state.rs index 1fba49fa..130c0cb5 100644 --- a/src/processor/fast/commit_state.rs +++ b/src/processor/fast/commit_state.rs @@ -10,7 +10,6 @@ use pinocchio_system::instructions as system; use crate::args::CommitStateArgs; use crate::error::DlpError; -use crate::pda; use crate::processor::fast::utils::{ pda::create_pda, requires::{ @@ -20,6 +19,7 @@ use crate::processor::fast::utils::{ }, }; use crate::state::{CommitRecord, DelegationMetadata, DelegationRecord, ProgramConfig}; +use crate::{merge_diff_copy, pda, DiffSet}; use super::to_pinocchio_program_error; @@ -59,7 +59,6 @@ pub fn process_commit_state( ) -> ProgramResult { let args = CommitStateArgs::try_from_slice(data).map_err(|_| ProgramError::BorshIoError)?; - let commit_state_bytes: &[u8] = args.data.as_ref(); let commit_record_lamports = args.lamports; let commit_record_nonce = args.nonce; let allow_undelegation = args.allow_undelegation; @@ -71,7 +70,7 @@ pub fn process_commit_state( }; let commit_args = CommitStateInternalArgs { - commit_state_bytes, + commit_state_bytes: NewState::FullBytes(&args.data), commit_record_lamports, commit_record_nonce, allow_undelegation, @@ -88,9 +87,23 @@ pub fn process_commit_state( process_commit_state_internal(commit_args) } +pub(crate) enum NewState<'a> { + FullBytes(&'a [u8]), + Diff(DiffSet<'a>), +} + +impl NewState<'_> { + pub fn data_len(&self) -> usize { + match self { + NewState::FullBytes(bytes) => bytes.len(), + NewState::Diff(diff) => diff.changed_len(), + } + } +} + /// Arguments for the commit state internal function pub(crate) struct CommitStateInternalArgs<'a> { - pub(crate) commit_state_bytes: &'a [u8], + pub(crate) commit_state_bytes: NewState<'a>, pub(crate) commit_record_lamports: u64, pub(crate) commit_record_nonce: u64, pub(crate) allow_undelegation: bool, @@ -238,7 +251,7 @@ pub(crate) fn process_commit_state_internal( create_pda( args.commit_state_account, &crate::fast::ID, - args.commit_state_bytes.len(), + args.commit_state_bytes.data_len(), &[Signer::from(&seeds!( pda::COMMIT_STATE_TAG, args.delegated_account.key(), @@ -274,7 +287,14 @@ pub(crate) fn process_commit_state_internal( // Copy the new state to the initialized PDA let mut commit_state_data = args.commit_state_account.try_borrow_mut_data()?; - (*commit_state_data).copy_from_slice(args.commit_state_bytes); + + match args.commit_state_bytes { + NewState::FullBytes(bytes) => (*commit_state_data).copy_from_slice(bytes), + NewState::Diff(diff) => { + let original_data = args.delegated_account.try_borrow_data()?; + merge_diff_copy(&mut commit_state_data, &original_data, &diff)?; + } + } // TODO - Add additional validation for the commitment, e.g. sufficient validator stake diff --git a/src/processor/fast/commit_state_from_buffer.rs b/src/processor/fast/commit_state_from_buffer.rs index 97440e71..167b2227 100644 --- a/src/processor/fast/commit_state_from_buffer.rs +++ b/src/processor/fast/commit_state_from_buffer.rs @@ -7,6 +7,8 @@ use pinocchio::program_error::ProgramError; use pinocchio::pubkey::Pubkey; use pinocchio::ProgramResult; +use super::NewState; + pub fn process_commit_state_from_buffer( _program_id: &Pubkey, accounts: &[AccountInfo], @@ -26,10 +28,9 @@ pub fn process_commit_state_from_buffer( let allow_undelegation = args.allow_undelegation; let state = state_buffer_account.try_borrow_data()?; - let commit_state_bytes: &[u8] = &state; let commit_args = CommitStateInternalArgs { - commit_state_bytes, + commit_state_bytes: NewState::FullBytes(&state), commit_record_lamports, commit_record_nonce, allow_undelegation,