refactor: use requires macros instead of if-else for validations (#129)

snawaz · web-flow · commit d58eb954302e · 2025-12-31T13:23:47.000+01:00
diff --git a/src/diff/algorithm.rs b/src/diff/algorithm.rs
@@ -10,6 +10,8 @@ use super::{
     SIZE_OF_SINGLE_OFFSET_PAIR,
 };
 
+use crate::{require_eq, require_le};
+
 ///
 /// Compute diff between original and changed.
 ///
@@ -258,18 +260,18 @@ pub fn merge_diff_copy<'a>(
     original: &[u8],
     diffset: &DiffSet<'_>,
 ) -> Result<&'a mut [u8], ProgramError> {
-    if destination.len() != diffset.changed_len() {
-        return Err(DlpError::MergeDiffError.into());
-    }
+    require_eq!(
+        destination.len(),
+        diffset.changed_len(),
+        DlpError::MergeDiffError
+    );
 
     let mut write_index = 0;
     for item in diffset.iter() {
         let (diff_segment, OffsetInData { start, end }) = item?;
 
         if write_index < start {
-            if start > original.len() {
-                return Err(DlpError::InvalidDiff.into());
-            }
+            require_le!(start, original.len(), DlpError::InvalidDiff);
             // copy the unchanged bytes
             destination[write_index..start].copy_from_slice(&original[write_index..start]);
         }
diff --git a/src/error.rs b/src/error.rs
@@ -7,86 +7,130 @@ use thiserror::Error;
 pub enum DlpError {
     #[error("Invalid Authority")]
     InvalidAuthority = 0,
+
     #[error("Account cannot be undelegated, is_undelegatable is false")]
     NotUndelegatable = 1,
+
     #[error("Unauthorized Operation")]
     Unauthorized = 2,
+
     #[error("Invalid Authority for the current target program")]
     InvalidAuthorityForProgram = 3,
+
     #[error("Delegated account does not match the expected account")]
     InvalidDelegatedAccount = 4,
+
     #[error("Delegated account is not in a valid state")]
     InvalidDelegatedState = 5,
+
     #[error("Reimbursement account does not match the expected account")]
     InvalidReimbursementAccount = 6,
+
     #[error("Invalid account data after CPI")]
     InvalidAccountDataAfterCPI = 7,
+
     #[error("Invalid validator balance after CPI")]
     InvalidValidatorBalanceAfterCPI = 8,
+
     #[error("Invalid reimbursement address for delegation rent")]
     InvalidReimbursementAddressForDelegationRent = 9,
+
     #[error("Authority is invalid for the delegated account program owner")]
     InvalidWhitelistProgramConfig = 10,
+
     #[error("Account already undelegated")]
     AlreadyUndelegated = 11,
+
     #[error("Commit is out of order")]
     NonceOutOfOrder = 12,
+
     #[error("Computation overflow detected")]
     Overflow = 13,
+
     #[error("Too many seeds")]
     TooManySeeds = 14,
+
     #[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,
+
     #[error("Commit state PDA invalid seeds")]
     CommitStateInvalidSeeds = 18,
+
     #[error("Commit state PDA invalid account owner")]
     CommitStateInvalidAccountOwner = 19,
+
     #[error("Commit state PDA is already initialized")]
     CommitStateAlreadyInitialized = 20,
+
     #[error("Commit state PDA immutable")]
     CommitStateImmutable = 21,
+
     #[error("Commit record PDA invalid seeds")]
     CommitRecordInvalidSeeds = 22,
+
     #[error("Commit record PDA invalid account owner")]
     CommitRecordInvalidAccountOwner = 23,
+
     #[error("Commit record PDA is already initialized")]
     CommitRecordAlreadyInitialized = 24,
+
     #[error("Commit record PDA immutable")]
     CommitRecordImmutable = 25,
+
     #[error("Delegation record PDA invalid seeds")]
     DelegationRecordInvalidSeeds = 26,
+
     #[error("Delegation record PDA invalid account owner")]
     DelegationRecordInvalidAccountOwner = 27,
+
     #[error("Delegation record PDA is already initialized")]
     DelegationRecordAlreadyInitialized = 28,
+
     #[error("Delegation record PDA immutable")]
     DelegationRecordImmutable = 29,
+
     #[error("Delegation metadata PDA invalid seeds")]
     DelegationMetadataInvalidSeeds = 30,
+
     #[error("Delegation metadata PDA invalid account owner")]
     DelegationMetadataInvalidAccountOwner = 31,
+
     #[error("Delegation metadata PDA is already initialized")]
     DelegationMetadataAlreadyInitialized = 32,
+
     #[error("Delegation metadata PDA immutable")]
     DelegationMetadataImmutable = 33,
+
     #[error("Undelegate buffer PDA invalid seeds")]
     UndelegateBufferInvalidSeeds = 34,
+
     #[error("Undelegate buffer PDA invalid account owner")]
     UndelegateBufferInvalidAccountOwner = 35,
+
     #[error("Undelegate buffer PDA is already initialized")]
     UndelegateBufferAlreadyInitialized = 36,
+
     #[error("Undelegate buffer PDA immutable")]
     UndelegateBufferImmutable = 37,
+
     #[error("Invalid data length for deserialization")]
     InvalidDataLength = 38,
+
     #[error("Invalid discriminator for delegation record")]
     InvalidDiscriminator = 39,
+
     #[error("Invalid delegation record deserialization")]
     InvalidDelegationRecordData = 40,
+
+    #[error("Too many account keys passed to the instruction")]
+    TooManyAccountKeys = 41,
+
     #[error("An infallible error is encountered possibly due to logic error")]
     InfallibleError = 100,
 }
diff --git a/src/processor/fast/commit_diff.rs b/src/processor/fast/commit_diff.rs
@@ -10,6 +10,8 @@ use crate::DiffSet;
 
 use super::NewState;
 
+use crate::{require, require_n_accounts};
+
 /// Commit diff to a delegated PDA
 ///
 /// Accounts:
@@ -47,15 +49,22 @@ pub fn process_commit_diff(
     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);
-    };
+    let [
+        validator, // force multi-line
+        delegated_account,
+        commit_state_account,
+        commit_record_account,
+        delegation_record_account,
+        delegation_metadata_account,
+        validator_fees_vault,
+        program_config_account,
+        _system_program,
+    ] = require_n_accounts!(accounts, 9);
 
-    if data.len() < SIZE_COMMIT_DIFF_ARGS_WITHOUT_DIFF {
-        return Err(ProgramError::InvalidInstructionData);
-    }
+    require!(
+        data.len() >= SIZE_COMMIT_DIFF_ARGS_WITHOUT_DIFF,
+        ProgramError::InvalidInstructionData
+    );
 
     let (diff, data) = data.split_at(data.len() - SIZE_COMMIT_DIFF_ARGS_WITHOUT_DIFF);
 
diff --git a/src/processor/fast/utils/requires.rs b/src/processor/fast/utils/requires.rs
@@ -27,6 +27,139 @@ mod pubkey {
     }
 }
 
+// require true
+#[macro_export]
+macro_rules! require {
+    ($cond:expr, $error:expr) => {{
+        if !$cond {
+            let expr = stringify!($cond);
+            pinocchio_log::log!("require!({}) failed.", expr);
+            return Err($error.into());
+        }
+    }};
+}
+
+// require key1 == key2
+#[macro_export]
+macro_rules! require_eq_keys {
+    ( $key1:expr, $key2:expr, $error:expr) => {{
+        if !pinocchio::pubkey::pubkey_eq($key1, $key2) {
+            pinocchio_log::log!(
+                "require_eq_keys!({}, {}) failed: ",
+                stringify!($key1),
+                stringify!($key2)
+            );
+            pinocchio::pubkey::log($key1);
+            pinocchio::pubkey::log($key2);
+            return Err($error.into());
+        }
+    }};
+}
+
+// require a == b
+#[macro_export]
+macro_rules! require_eq {
+    ( $val1:expr, $val2:expr, $error:expr) => {{
+        if !($val1 == $val2) {
+            pinocchio_log::log!(
+                "require_eq!({}, {}) failed: {} == {}",
+                stringify!($val1),
+                stringify!($val2),
+                $val1,
+                $val2
+            );
+            return Err($error.into());
+        }
+    }};
+}
+
+// require a <= b
+#[macro_export]
+macro_rules! require_le {
+    ( $val1:expr, $val2:expr, $error:expr) => {{
+        if !($val1 <= $val2) {
+            pinocchio_log::log!(
+                "require_le!({}, {}) failed: {} <= {}",
+                stringify!($val1),
+                stringify!($val2),
+                $val1,
+                $val2
+            );
+            return Err($error.into());
+        }
+    }};
+}
+
+// require a < b
+#[macro_export]
+macro_rules! require_lt {
+    ( $val1:expr, $val2:expr, $error:expr) => {{
+        if !($val1 < $val2) {
+            pinocchio_log::log!(
+                "require_lt!({}, {}) failed: {} < {}",
+                stringify!($val1),
+                stringify!($val2),
+                $val1,
+                $val2
+            );
+            return Err($error.into());
+        }
+    }};
+}
+
+// require a >= b
+#[macro_export]
+macro_rules! require_ge {
+    ( $val1:expr, $val2:expr, $error:expr) => {{
+        if !($val1 >= $val2) {
+            pinocchio_log::log!(
+                "require_ge!({}, {}) failed: {} >= {}",
+                stringify!($val1),
+                stringify!($val2),
+                $val1,
+                $val2
+            );
+            return Err($error.into());
+        }
+    }};
+}
+
+#[macro_export]
+macro_rules! require_n_accounts {
+    ( $accounts:expr, $n:literal) => {{
+        match $accounts.len().cmp(&$n) {
+            core::cmp::Ordering::Less => {
+                pinocchio_log::log!(
+                    "Need {} accounts, but got less ({}) accounts",
+                    $n,
+                    $accounts.len()
+                );
+                return Err(pinocchio::program_error::ProgramError::NotEnoughAccountKeys);
+            }
+            core::cmp::Ordering::Equal => TryInto::<&[_; $n]>::try_into($accounts)
+                .map_err(|_| $crate::error::DlpError::InfallibleError)?,
+            core::cmp::Ordering::Greater => {
+                pinocchio_log::log!(
+                    "Need {} accounts, but got more ({}) accounts",
+                    $n,
+                    $accounts.len()
+                );
+                return Err($crate::error::DlpError::TooManyAccountKeys.into());
+            }
+        }
+    }};
+}
+
+#[macro_export]
+macro_rules! require_some {
+    ($option:expr, $error:expr) => {{
+        match $option {
+            Some(val) => val,
+            None => return Err($error.into()),
+        }
+    }};
+}
+
 /// Errors if:
 /// - Account is not owned by expected program.
 #[inline(always)]
