diff --git a/runtime/src/account_overrides.rs b/runtime/src/account_overrides.rs index 62ee494ff2..ee8e7ec9e2 100644 --- a/runtime/src/account_overrides.rs +++ b/runtime/src/account_overrides.rs @@ -1,25 +1,31 @@ -use solana_sdk::{account::AccountSharedData, pubkey::Pubkey, sysvar}; +use { + solana_sdk::{account::AccountSharedData, pubkey::Pubkey, sysvar}, + std::collections::HashMap, +}; /// Encapsulates overridden accounts, typically used for transaction simulations #[derive(Default)] pub struct AccountOverrides { - pub slot_history: Option, + accounts: HashMap, } impl AccountOverrides { + pub fn set_account(&mut self, pubkey: &Pubkey, account: Option) { + match account { + Some(account) => self.accounts.insert(*pubkey, account), + None => self.accounts.remove(pubkey), + }; + } + /// Sets in the slot history /// /// Note: no checks are performed on the correctness of the contained data pub fn set_slot_history(&mut self, slot_history: Option) { - self.slot_history = slot_history; + self.set_account(&sysvar::slot_history::id(), slot_history); } /// Gets the account if it's found in the list of overrides pub fn get(&self, pubkey: &Pubkey) -> Option<&AccountSharedData> { - if pubkey == &sysvar::slot_history::id() { - self.slot_history.as_ref() - } else { - None - } + self.accounts.get(pubkey) } }