feat: Optimize CommitDiff using efficient merge_diff_copy that avoids heap allocation (#113)

Author: snawaz

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<Vec<u8>
     })
 }
 
+/// 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);
     }
 }

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<DlpError> for ProgramError {

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,

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
 

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,