From 7a753500dd8da091cd000f0bdbd2bd9f8213391a Mon Sep 17 00:00:00 2001
From: Joe C <joe.caulfield@anza.xyz>
Date: Mon, 24 Jun 2024 20:54:43 -0500
Subject: [PATCH] Program cache API (#15)

---
 harness/src/lib.rs     | 28 ++++---------
 harness/src/program.rs | 92 ++++++++++++++++++++++++------------------
 2 files changed, 59 insertions(+), 61 deletions(-)

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