make account compilation take an iter

Author: buffalojoec

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 || {

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()

harness/src/lib.rs

@@ -812,7 +812,7 @@ impl Mollusk {
             program_id_index,
             instruction_accounts,
             transaction_accounts,
-        } = crate::compile_accounts::compile_accounts(instruction, accounts, loader_key);
+        } = crate::compile_accounts::compile_accounts(instruction, accounts.iter(), loader_key);
 
         self.process_instruction_inner(
             instruction,
@@ -850,7 +850,7 @@ impl Mollusk {
                 program_id_index,
                 instruction_accounts,
                 transaction_accounts,
-            } = crate::compile_accounts::compile_accounts(instruction, accounts, loader_key);
+            } = crate::compile_accounts::compile_accounts(instruction, accounts.iter(), loader_key);
 
             let this_result = self.process_instruction_inner(
                 instruction,
@@ -904,7 +904,7 @@ impl Mollusk {
             program_id_index,
             instruction_accounts,
             transaction_accounts,
-        } = crate::compile_accounts::compile_accounts(instruction, accounts, loader_key);
+        } = crate::compile_accounts::compile_accounts(instruction, accounts.iter(), loader_key);
 
         let result = self.process_instruction_inner(
             instruction,
@@ -963,7 +963,7 @@ impl Mollusk {
                 program_id_index,
                 instruction_accounts,
                 transaction_accounts,
-            } = crate::compile_accounts::compile_accounts(instruction, accounts, loader_key);
+            } = crate::compile_accounts::compile_accounts(instruction, accounts.iter(), loader_key);
 
             let this_result = self.process_instruction_inner(
                 instruction,

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()
 }

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)]