harness: refactor account compilation step (#157)

buffalojoec · web-flow · commit 1c0e142e5f86 · 2025-11-18T21:03:09.000+08:00
diff --git a/harness/src/compile_accounts.rs b/harness/src/compile_accounts.rs
@@ -21,9 +21,9 @@ pub struct CompiledAccounts {
     pub transaction_accounts: Vec<TransactionAccount>,
 }
 
-pub fn compile_accounts(
+pub fn compile_accounts<'a>(
     instruction: &Instruction,
-    accounts: &[(Pubkey, Account)],
+    accounts: impl Iterator<Item = &'a (Pubkey, Account)>,
     loader_key: Pubkey,
 ) -> CompiledAccounts {
     let stub_out_program_account = move || {
diff --git a/harness/src/fuzz/firedancer.rs b/harness/src/fuzz/firedancer.rs
@@ -73,7 +73,7 @@ fn build_fixture_context(
         instruction_accounts,
         transaction_accounts,
         ..
-    } = compile_accounts(instruction, accounts, loader_key);
+    } = compile_accounts(instruction, accounts.iter(), loader_key);
 
     let accounts = transaction_accounts
         .into_iter()
diff --git a/harness/src/lib.rs b/harness/src/lib.rs
@@ -454,8 +454,6 @@ pub use mollusk_svm_result as result;
 use mollusk_svm_result::Compare;
 #[cfg(feature = "precompiles")]
 use solana_precompile_error::PrecompileError;
-#[cfg(feature = "invocation-inspect-callback")]
-use solana_transaction_context::InstructionAccount;
 use {
     crate::{
         account_store::AccountStore, compile_accounts::CompiledAccounts, epoch_stake::EpochStake,
@@ -464,7 +462,7 @@ use {
     agave_feature_set::FeatureSet,
     mollusk_svm_error::error::{MolluskError, MolluskPanic},
     mollusk_svm_result::{Check, CheckContext, Config, InstructionResult},
-    solana_account::Account,
+    solana_account::{Account, AccountSharedData},
     solana_compute_budget::compute_budget::ComputeBudget,
     solana_hash::Hash,
     solana_instruction::{AccountMeta, Instruction},
@@ -473,7 +471,7 @@ use {
     solana_svm_callback::InvokeContextCallback,
     solana_svm_log_collector::LogCollector,
     solana_svm_timings::ExecuteTimings,
-    solana_transaction_context::TransactionContext,
+    solana_transaction_context::{InstructionAccount, TransactionContext},
     std::{cell::RefCell, collections::HashSet, iter::once, rc::Rc},
 };
 
@@ -680,32 +678,28 @@ impl Mollusk {
         self.sysvars.warp_to_slot(slot)
     }
 
-    /// Process an instruction using the minified Solana Virtual Machine (SVM)
-    /// environment. Simply returns the result.
-    pub fn process_instruction(
+    fn get_loader_key(&self, program_id: &Pubkey) -> Pubkey {
+        if crate::program::precompile_keys::is_precompile(program_id) {
+            crate::program::loader_keys::NATIVE_LOADER
+        } else {
+            self.program_cache
+                .load_program(program_id)
+                .or_panic_with(MolluskError::ProgramNotCached(program_id))
+                .account_owner()
+        }
+    }
+
+    fn process_instruction_inner(
         &self,
         instruction: &Instruction,
         accounts: &[(Pubkey, Account)],
+        program_id_index: u16,
+        instruction_accounts: Vec<InstructionAccount>,
+        transaction_accounts: Vec<(Pubkey, AccountSharedData)>,
     ) -> InstructionResult {
         let mut compute_units_consumed = 0;
         let mut timings = ExecuteTimings::default();
 
-        let loader_key = if crate::program::precompile_keys::is_precompile(&instruction.program_id)
-        {
-            crate::program::loader_keys::NATIVE_LOADER
-        } else {
-            self.program_cache
-                .load_program(&instruction.program_id)
-                .or_panic_with(MolluskError::ProgramNotCached(&instruction.program_id))
-                .account_owner()
-        };
-
-        let CompiledAccounts {
-            program_id_index,
-            instruction_accounts,
-            transaction_accounts,
-        } = crate::compile_accounts::compile_accounts(instruction, accounts, loader_key);
-
         let mut transaction_context = TransactionContext::new(
             transaction_accounts,
             self.sysvars.rent.clone(),
@@ -805,6 +799,30 @@ impl Mollusk {
         }
     }
 
+    /// Process an instruction using the minified Solana Virtual Machine (SVM)
+    /// environment. Simply returns the result.
+    pub fn process_instruction(
+        &self,
+        instruction: &Instruction,
+        accounts: &[(Pubkey, Account)],
+    ) -> InstructionResult {
+        let loader_key = self.get_loader_key(&instruction.program_id);
+
+        let CompiledAccounts {
+            program_id_index,
+            instruction_accounts,
+            transaction_accounts,
+        } = crate::compile_accounts::compile_accounts(instruction, accounts.iter(), loader_key);
+
+        self.process_instruction_inner(
+            instruction,
+            accounts,
+            program_id_index,
+            instruction_accounts,
+            transaction_accounts,
+        )
+    }
+
     /// Process a chain of instructions using the minified Solana Virtual
     /// Machine (SVM) environment. The returned result is an
     /// `InstructionResult`, containing:
@@ -820,22 +838,36 @@ impl Mollusk {
         instructions: &[Instruction],
         accounts: &[(Pubkey, Account)],
     ) -> InstructionResult {
-        let mut result = InstructionResult {
+        let mut composite_result = InstructionResult {
             resulting_accounts: accounts.to_vec(),
             ..Default::default()
         };
 
         for instruction in instructions {
-            let this_result = self.process_instruction(instruction, &result.resulting_accounts);
+            let loader_key = self.get_loader_key(&instruction.program_id);
+
+            let CompiledAccounts {
+                program_id_index,
+                instruction_accounts,
+                transaction_accounts,
+            } = crate::compile_accounts::compile_accounts(instruction, accounts.iter(), loader_key);
+
+            let this_result = self.process_instruction_inner(
+                instruction,
+                accounts,
+                program_id_index,
+                instruction_accounts,
+                transaction_accounts,
+            );
 
-            result.absorb(this_result);
+            composite_result.absorb(this_result);
 
-            if result.program_result.is_err() {
+            if composite_result.program_result.is_err() {
                 break;
             }
         }
 
-        result
+        composite_result
     }
 
     /// Process an instruction using the minified Solana Virtual Machine (SVM)
@@ -866,7 +898,21 @@ impl Mollusk {
         accounts: &[(Pubkey, Account)],
         checks: &[Check],
     ) -> InstructionResult {
-        let result = self.process_instruction(instruction, accounts);
+        let loader_key = self.get_loader_key(&instruction.program_id);
+
+        let CompiledAccounts {
+            program_id_index,
+            instruction_accounts,
+            transaction_accounts,
+        } = crate::compile_accounts::compile_accounts(instruction, accounts.iter(), loader_key);
+
+        let result = self.process_instruction_inner(
+            instruction,
+            accounts,
+            program_id_index,
+            instruction_accounts,
+            transaction_accounts,
+        );
 
         #[cfg(any(feature = "fuzz", feature = "fuzz-fd"))]
         fuzz::generate_fixtures_from_mollusk_test(self, instruction, accounts, &result);
@@ -904,26 +950,42 @@ impl Mollusk {
         instructions: &[(&Instruction, &[Check])],
         accounts: &[(Pubkey, Account)],
     ) -> InstructionResult {
-        let mut result = InstructionResult {
+        let mut composite_result = InstructionResult {
             resulting_accounts: accounts.to_vec(),
             ..Default::default()
         };
 
         for (instruction, checks) in instructions.iter() {
-            let this_result = self.process_and_validate_instruction(
+            let loader_key = self.get_loader_key(&instruction.program_id);
+            let accounts = &composite_result.resulting_accounts;
+
+            let CompiledAccounts {
+                program_id_index,
+                instruction_accounts,
+                transaction_accounts,
+            } = crate::compile_accounts::compile_accounts(instruction, accounts.iter(), loader_key);
+
+            let this_result = self.process_instruction_inner(
                 instruction,
-                &result.resulting_accounts,
-                checks,
+                accounts,
+                program_id_index,
+                instruction_accounts,
+                transaction_accounts,
             );
 
-            result.absorb(this_result);
+            #[cfg(any(feature = "fuzz", feature = "fuzz-fd"))]
+            fuzz::generate_fixtures_from_mollusk_test(self, instruction, accounts, &this_result);
+
+            this_result.run_checks(checks, &self.config, self);
 
-            if result.program_result.is_err() {
+            composite_result.absorb(this_result);
+
+            if composite_result.program_result.is_err() {
                 break;
             }
         }
 
-        result
+        composite_result
     }
 
     #[cfg(feature = "fuzz")]
diff --git a/keys/src/accounts.rs b/keys/src/accounts.rs
@@ -2,11 +2,12 @@
 
 use {
     crate::keys::KeyMap,
-    mollusk_svm_error::error::MolluskError,
+    mollusk_svm_error::error::{MolluskError, MolluskPanic},
     solana_account::{Account, AccountSharedData},
     solana_instruction::Instruction,
     solana_pubkey::Pubkey,
     solana_transaction_context::{IndexOfAccount, InstructionAccount, TransactionAccount},
+    std::collections::HashMap,
 };
 
 // Helper struct to avoid cloning instruction data.
@@ -47,60 +48,64 @@ pub fn compile_instruction_accounts(
         .collect()
 }
 
-pub fn compile_transaction_accounts_for_instruction(
+pub fn compile_transaction_accounts_for_instruction<'a>(
     key_map: &KeyMap,
     instruction: &Instruction,
-    accounts: &[(Pubkey, Account)],
+    accounts: impl Iterator<Item = &'a (Pubkey, Account)>,
     stub_out_program_account: Option<Box<dyn Fn() -> Account>>,
 ) -> Vec<TransactionAccount> {
+    let len = key_map.len();
+    let mut by_key: HashMap<Pubkey, AccountSharedData> = HashMap::with_capacity(len);
+
+    for (key, account) in accounts {
+        if key_map.contains_key(key) {
+            by_key.insert(*key, AccountSharedData::from(account.clone()));
+        }
+    }
+
     key_map
         .keys()
         .map(|key| {
-            let account = accounts
-                .iter()
-                .find(|(k, _)| k == key)
-                .map(|(_, account)| AccountSharedData::from(account.clone()));
-
-            if let Some(account) = account {
-                (*key, account)
-            } else if let Some(stub_out_program_account) = &stub_out_program_account {
+            if let Some(stub_out_program_account) = &stub_out_program_account {
                 if instruction.program_id == *key {
-                    (*key, stub_out_program_account().into())
-                } else {
-                    panic!("{}", MolluskError::AccountMissing(key))
+                    return (*key, stub_out_program_account().into());
                 }
-            } else {
-                panic!("{}", MolluskError::AccountMissing(key))
             }
+            let account = by_key
+                .remove(key)
+                .or_panic_with(MolluskError::AccountMissing(key));
+            (*key, account)
         })
         .collect()
 }
 
-pub fn compile_transaction_accounts(
+pub fn compile_transaction_accounts<'a>(
     key_map: &KeyMap,
     instructions: &[Instruction],
-    accounts: &[(Pubkey, Account)],
+    accounts: impl Iterator<Item = &'a (Pubkey, Account)>,
     stub_out_program_account: Option<Box<dyn Fn() -> Account>>,
 ) -> Vec<TransactionAccount> {
+    let len = key_map.len();
+    let mut by_key: HashMap<Pubkey, AccountSharedData> = HashMap::with_capacity(len);
+
+    for (key, account) in accounts {
+        if key_map.contains_key(key) {
+            by_key.insert(*key, AccountSharedData::from(account.clone()));
+        }
+    }
+
     key_map
         .keys()
         .map(|key| {
-            let account = accounts
-                .iter()
-                .find(|(k, _)| k == key)
-                .map(|(_, account)| AccountSharedData::from(account.clone()));
-
-            if let Some(account) = account {
-                (*key, account)
-            } else if let Some(stub_out_program_account) = &stub_out_program_account {
+            if let Some(stub_out_program_account) = &stub_out_program_account {
                 if instructions.iter().any(|ix| ix.program_id == *key) {
-                    (*key, stub_out_program_account().into())
-                } else {
-                    panic!("{}", MolluskError::AccountMissing(key))
+                    return (*key, stub_out_program_account().into());
                 }
-            } else {
-                panic!("{}", MolluskError::AccountMissing(key))
             }
+            let account = by_key
+                .remove(key)
+                .or_panic_with(MolluskError::AccountMissing(key));
+            (*key, account)
         })
         .collect()
 }
diff --git a/keys/src/keys.rs b/keys/src/keys.rs
@@ -181,6 +181,21 @@ impl KeyMap {
     pub fn position(&self, key: &Pubkey) -> Option<usize> {
         self.map.keys().position(|k| k == key)
     }
+
+    /// Return whether or not the map contains a key.
+    pub fn contains_key(&self, key: &Pubkey) -> bool {
+        self.map.contains_key(key)
+    }
+
+    /// Get the length of the key map.
+    pub fn len(&self) -> usize {
+        self.map.len()
+    }
+
+    /// Return whether or not the key map is empty.
+    pub fn is_empty(&self) -> bool {
+        self.map.is_empty()
+    }
 }
 
 #[cfg(test)]
