diff --git a/harness/src/lib.rs b/harness/src/lib.rs index 0696f4e..b62434c 100644 --- a/harness/src/lib.rs +++ b/harness/src/lib.rs @@ -34,6 +34,7 @@ pub mod sysvar; use { crate::{ + program::ProgramCache, result::{Check, InstructionResult}, sysvar::Sysvars, }, @@ -65,7 +66,7 @@ pub struct Mollusk { pub compute_budget: ComputeBudget, pub feature_set: FeatureSet, pub program_account: AccountSharedData, - pub program_cache: LoadedProgramsForTxBatch, + pub program_cache: ProgramCache, pub program_id: Pubkey, pub sysvars: Sysvars, } @@ -83,7 +84,7 @@ impl Default for Mollusk { compute_budget: ComputeBudget::default(), feature_set: FeatureSet::all_enabled(), program_account, - program_cache: program::default_program_cache(), + program_cache: ProgramCache::default(), program_id, sysvars: Sysvars::default(), } @@ -97,22 +98,12 @@ impl Mollusk { /// Once loaded, adds the program to the program cache and updates the /// Mollusk instance with the program's ID and account. pub fn new(program_id: &Pubkey, program_name: &'static str) -> Self { - let elf = file::load_program_elf(program_name); - let mut mollusk = Self { program_id: *program_id, program_account: program::program_account(program_id), ..Default::default() }; - - program::add_program_to_cache( - &mut mollusk.program_cache, - program_id, - &elf, - &mollusk.compute_budget, - &mollusk.feature_set, - ); - + mollusk.add_program(program_id, program_name); mollusk } @@ -121,13 +112,8 @@ impl Mollusk { /// If you intend to CPI to a program, this is likely what you want to use. pub fn add_program(&mut self, program_id: &Pubkey, program_name: &'static str) { let elf = file::load_program_elf(program_name); - program::add_program_to_cache( - &mut self.program_cache, - program_id, - &elf, - &self.compute_budget, - &self.feature_set, - ); + self.program_cache + .add_program(program_id, &elf, &self.compute_budget, &self.feature_set); } /// Warp the test environment to a slot by updating sysvars. @@ -182,7 +168,7 @@ impl Mollusk { &sysvar_cache, None, self.compute_budget, - &self.program_cache, + self.program_cache.cache(), &mut programs_modified_by_tx, Arc::new(self.feature_set.clone()), Hash::default(), diff --git a/harness/src/program.rs b/harness/src/program.rs index 5f89ff9..7524f8d 100644 --- a/harness/src/program.rs +++ b/harness/src/program.rs @@ -18,6 +18,58 @@ use { std::sync::Arc, }; +pub struct ProgramCache { + cache: LoadedProgramsForTxBatch, +} + +impl Default for ProgramCache { + fn default() -> Self { + let mut cache = LoadedProgramsForTxBatch::default(); + BUILTINS.iter().for_each(|builtin| { + let program_id = builtin.program_id; + let entry = builtin.program_cache_entry(); + cache.replenish(program_id, entry); + }); + Self { cache } + } +} + +impl ProgramCache { + pub(crate) fn cache(&self) -> &LoadedProgramsForTxBatch { + &self.cache + } + + /// Add a program to a program cache. + pub fn add_program( + &mut self, + program_id: &Pubkey, + elf: &[u8], + compute_budget: &ComputeBudget, + feature_set: &FeatureSet, + ) { + let environment = Arc::new( + create_program_runtime_environment_v1(feature_set, compute_budget, false, false) + .unwrap(), + ); + self.cache.replenish( + *program_id, + Arc::new( + LoadedProgram::new( + &bpf_loader_upgradeable::id(), + environment, + 0, + 0, + None, + elf, + elf.len(), + &mut LoadProgramMetrics::default(), + ) + .unwrap(), + ), + ); + } +} + struct Builtin { program_id: Pubkey, name: &'static str, @@ -125,43 +177,3 @@ pub fn program_data_account(elf: &[u8]) -> AccountSharedData { pub fn program_accounts(program_id: &Pubkey, elf: &[u8]) -> (AccountSharedData, AccountSharedData) { (program_account(program_id), program_data_account(elf)) } - -/// Create a default program cache instance. -pub fn default_program_cache() -> LoadedProgramsForTxBatch { - let mut cache = LoadedProgramsForTxBatch::default(); - BUILTINS.iter().for_each(|builtin| { - let program_id = builtin.program_id; - let entry = builtin.program_cache_entry(); - cache.replenish(program_id, entry); - }); - cache -} - -/// Add a program to a program cache. -pub fn add_program_to_cache( - cache: &mut LoadedProgramsForTxBatch, - program_id: &Pubkey, - elf: &[u8], - compute_budget: &ComputeBudget, - feature_set: &FeatureSet, -) { - let environment = Arc::new( - create_program_runtime_environment_v1(feature_set, compute_budget, false, false).unwrap(), - ); - cache.replenish( - *program_id, - Arc::new( - LoadedProgram::new( - &bpf_loader_upgradeable::id(), - environment, - 0, - 0, - None, - elf, - elf.len(), - &mut LoadProgramMetrics::default(), - ) - .unwrap(), - ), - ); -}