Update program API (#12)

harness/src/lib.rs

@@ -42,13 +42,16 @@ use {
     solana_sdk::{
         account::AccountSharedData,
         clock::Clock,
+        epoch_rewards::EpochRewards,
+        epoch_schedule::EpochSchedule,
         feature_set::FeatureSet,
         hash::Hash,
         instruction::Instruction,
         pubkey::Pubkey,
         rent::Rent,
         slot_hashes::SlotHashes,
-        system_program,
+        stake_history::StakeHistory,
+        sysvar::last_restart_slot::LastRestartSlot,
         transaction_context::{InstructionAccount, TransactionContext},
     },
     std::sync::Arc,
@@ -78,12 +81,13 @@ impl Default for Mollusk {
              solana_runtime::message_processor=debug,\
              solana_runtime::system_instruction_processor=trace",
         );
+        let (program_id, program_account) = program::system_program();
         Self {
             compute_budget: ComputeBudget::default(),
             feature_set: FeatureSet::all_enabled(),
-            program_account: program::system_program_account(),
+            program_account,
             program_cache: program::default_program_cache(),
-            program_id: system_program::id(),
+            program_id,
             sysvar_cache: sysvar::default_sysvar_cache(),
         }
     }
@@ -100,7 +104,7 @@ impl Mollusk {
 
         let mut mollusk = Self {
             program_id: *program_id,
-            program_account: program::create_program_account(program_id),
+            program_account: program::program_account(program_id),
             ..Default::default()
         };
 
@@ -129,40 +133,46 @@ impl Mollusk {
         );
     }
 
+    /// Get the current clock.
+    pub fn get_clock(&self) -> Arc<Clock> {
+        self.sysvar_cache.get_clock().unwrap_or_default()
+    }
+
+    /// Get the current epoch rewards.
+    pub fn get_epoch_rewards(&self) -> Arc<EpochRewards> {
+        self.sysvar_cache.get_epoch_rewards().unwrap_or_default()
+    }
+
+    /// Get the current epoch schedule.
+    pub fn get_epoch_schedule(&self) -> Arc<EpochSchedule> {
+        self.sysvar_cache.get_epoch_schedule().unwrap_or_default()
+    }
+
+    /// Get the current last restart slot.
+    pub fn get_last_restart_slot(&self) -> Arc<LastRestartSlot> {
+        self.sysvar_cache
+            .get_last_restart_slot()
+            .unwrap_or_default()
+    }
+
     /// Get the current rent.
     pub fn get_rent(&self) -> Arc<Rent> {
         self.sysvar_cache.get_rent().unwrap_or_default()
     }
 
+    /// Get the current slot hashes.
+    pub fn get_slot_hashes(&self) -> Arc<SlotHashes> {
+        self.sysvar_cache.get_slot_hashes().unwrap_or_default()
+    }
+
+    /// Get the current stake history.
+    pub fn get_stake_history(&self) -> Arc<StakeHistory> {
+        self.sysvar_cache.get_stake_history().unwrap_or_default()
+    }
+
     /// Warp to a slot by updating the `Clock` and `SlotHashes` sysvars.
     pub fn warp_to_slot(&mut self, slot: u64) {
-        // First update `Clock`.
-        let epoch_schedule = self.sysvar_cache.get_epoch_schedule().unwrap_or_default();
-        let epoch = epoch_schedule.get_epoch(slot);
-        let leader_schedule_epoch = epoch_schedule.get_leader_schedule_epoch(slot);
-        self.sysvar_cache.set_clock(Clock {
-            slot,
-            epoch,
-            leader_schedule_epoch,
-            ..Default::default()
-        });
-
-        // Then update `SlotHashes`.
-        let mut i = 0;
-        if let Some(most_recent_slot_hash) = self
-            .sysvar_cache
-            .get_slot_hashes()
-            .unwrap_or_default()
-            .first()
-        {
-            i = most_recent_slot_hash.0;
-        }
-        let mut new_slot_hashes = vec![];
-        for slot in i..slot + 1 {
-            new_slot_hashes.push((slot, Hash::default()));
-        }
-        self.sysvar_cache
-            .set_slot_hashes(SlotHashes::new(&new_slot_hashes));
+        sysvar::warp_sysvar_cache_to_slot(&mut self.sysvar_cache, slot);
     }
 
     /// The main Mollusk API method.

harness/src/program.rs

@@ -25,18 +25,6 @@ struct Builtin {
 }
 
 impl Builtin {
-    fn program_account(&self) -> AccountSharedData {
-        let data = self.name.as_bytes().to_vec();
-        let lamports = Rent::default().minimum_balance(data.len());
-        AccountSharedData::from(Account {
-            lamports,
-            data,
-            owner: native_loader::id(),
-            executable: true,
-            rent_epoch: 0,
-        })
-    }
-
     fn program_cache_entry(&self) -> Arc<LoadedProgram> {
         Arc::new(LoadedProgram::new_builtin(
             0,
@@ -60,20 +48,33 @@ static BUILTINS: &[Builtin] = &[
     /* ... */
 ];
 
-/// Get the account for the system program.
-pub fn system_program_account() -> AccountSharedData {
-    BUILTINS[0].program_account()
+fn builtin_program_account(program_id: &Pubkey, name: &str) -> (Pubkey, AccountSharedData) {
+    let data = name.as_bytes().to_vec();
+    let lamports = Rent::default().minimum_balance(data.len());
+    let account = AccountSharedData::from(Account {
+        lamports,
+        data,
+        owner: native_loader::id(),
+        executable: true,
+        rent_epoch: 0,
+    });
+    (*program_id, account)
+}
+
+/// Get the key and account for the system program.
+pub fn system_program() -> (Pubkey, AccountSharedData) {
+    builtin_program_account(&BUILTINS[0].program_id, BUILTINS[0].name)
 }
 
-/// Get the account for the BPF Loader Upgradeable program.
-pub fn bpf_loader_upgradeable_program_account() -> AccountSharedData {
-    BUILTINS[1].program_account()
+/// Get the key and account for the BPF Loader Upgradeable program.
+pub fn bpf_loader_upgradeable_program() -> (Pubkey, AccountSharedData) {
+    builtin_program_account(&BUILTINS[1].program_id, BUILTINS[1].name)
 }
 
 /* ... */
 
 /// Create a BPF Loader Upgradeable program account.
-pub fn create_program_account(program_id: &Pubkey) -> AccountSharedData {
+pub fn program_account(program_id: &Pubkey) -> AccountSharedData {
     let programdata_address =
         Pubkey::find_program_address(&[program_id.as_ref()], &bpf_loader_upgradeable::id()).0;
     let data = bincode::serialize(&UpgradeableLoaderState::Program {
@@ -91,7 +92,7 @@ pub fn create_program_account(program_id: &Pubkey) -> AccountSharedData {
 }
 
 /// Create a BPF Loader Upgradeable program data account.
-pub fn create_program_data_account(elf: &[u8]) -> AccountSharedData {
+pub fn program_data_account(elf: &[u8]) -> AccountSharedData {
     let data = {
         let elf_offset = UpgradeableLoaderState::size_of_programdata_metadata();
         let data_len = elf_offset + elf.len();
@@ -121,14 +122,8 @@ pub fn create_program_data_account(elf: &[u8]) -> AccountSharedData {
 ///
 /// Returns a tuple, where the first element is the program account and the
 /// second element is the program data account.
-pub fn create_program_accounts(
-    program_id: &Pubkey,
-    elf: &[u8],
-) -> (AccountSharedData, AccountSharedData) {
-    (
-        create_program_account(program_id),
-        create_program_data_account(elf),
-    )
+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.

harness/src/sysvar.rs

@@ -3,16 +3,76 @@
 use {
     solana_program_runtime::sysvar_cache::SysvarCache,
     solana_sdk::{
-        clock::Clock, epoch_schedule::EpochSchedule, rent::Rent, slot_hashes::SlotHashes,
+        clock::{Clock, Slot},
+        epoch_rewards::EpochRewards,
+        epoch_schedule::EpochSchedule,
+        hash::Hash,
+        rent::Rent,
+        slot_hashes::SlotHashes,
+        stake_history::StakeHistory,
+        sysvar::{last_restart_slot::LastRestartSlot, SysvarId},
     },
 };
 
+pub(crate) fn warp_sysvar_cache_to_slot(sysvar_cache: &mut SysvarCache, slot: Slot) {
+    // First update `Clock`.
+    let epoch_schedule = sysvar_cache.get_epoch_schedule().unwrap_or_default();
+    let epoch = epoch_schedule.get_epoch(slot);
+    let leader_schedule_epoch = epoch_schedule.get_leader_schedule_epoch(slot);
+    let new_clock = Clock {
+        slot,
+        epoch,
+        leader_schedule_epoch,
+        ..Default::default()
+    };
+
+    // Then update `SlotHashes`.
+    let mut i = 0;
+    if let Some(most_recent_slot_hash) = sysvar_cache.get_slot_hashes().unwrap_or_default().first()
+    {
+        i = most_recent_slot_hash.0;
+    }
+    let mut new_slot_hashes = vec![];
+    for slot in i..slot + 1 {
+        new_slot_hashes.push((slot, Hash::default()));
+    }
+    let new_slot_hashes = SlotHashes::new(&new_slot_hashes);
+
+    sysvar_cache.fill_missing_entries(|pubkey, set_sysvar| {
+        if pubkey.eq(&Clock::id()) {
+            set_sysvar(&bincode::serialize(&new_clock).unwrap());
+        }
+        if pubkey.eq(&SlotHashes::id()) {
+            set_sysvar(&bincode::serialize(&new_slot_hashes).unwrap());
+        }
+    });
+}
+
 /// Create a default sysvar cache instance.
 pub fn default_sysvar_cache() -> SysvarCache {
     let mut cache = SysvarCache::default();
-    cache.set_clock(Clock::default());
-    cache.set_epoch_schedule(EpochSchedule::default());
-    cache.set_rent(Rent::default());
-    cache.set_slot_hashes(SlotHashes::default());
+    cache.fill_missing_entries(|pubkey, set_sysvar| {
+        if pubkey.eq(&Clock::id()) {
+            set_sysvar(&bincode::serialize(&Clock::default()).unwrap());
+        }
+        if pubkey.eq(&EpochRewards::id()) {
+            set_sysvar(&bincode::serialize(&EpochRewards::default()).unwrap());
+        }
+        if pubkey.eq(&EpochSchedule::id()) {
+            set_sysvar(&bincode::serialize(&EpochSchedule::default()).unwrap());
+        }
+        if pubkey.eq(&LastRestartSlot::id()) {
+            set_sysvar(&bincode::serialize(&LastRestartSlot::default()).unwrap());
+        }
+        if pubkey.eq(&Rent::id()) {
+            set_sysvar(&bincode::serialize(&Rent::default()).unwrap());
+        }
+        if pubkey.eq(&SlotHashes::id()) {
+            set_sysvar(&bincode::serialize(&SlotHashes::default()).unwrap());
+        }
+        if pubkey.eq(&StakeHistory::id()) {
+            set_sysvar(&bincode::serialize(&StakeHistory::default()).unwrap());
+        }
+    });
     cache
 }

harness/tests/bpf_program.rs

@@ -1,6 +1,6 @@
 use {
     mollusk_svm::{
-        program::{create_program_account, system_program_account},
+        program::{program_account, system_program},
         result::Check,
         Mollusk,
     },
@@ -128,7 +128,7 @@ fn test_transfer() {
             &[
                 (payer, payer_account.clone()),
                 (recipient, recipient_account.clone()),
-                (system_program::id(), system_program_account()),
+                system_program(),
             ],
             &[
                 Check::err(ProgramError::MissingRequiredSignature),
@@ -144,7 +144,7 @@ fn test_transfer() {
             &[
                 (payer, AccountSharedData::default()),
                 (recipient, recipient_account.clone()),
-                (system_program::id(), system_program_account()),
+                system_program(),
             ],
             &[
                 Check::err(ProgramError::Custom(
@@ -161,7 +161,7 @@ fn test_transfer() {
         &[
             (payer, payer_account.clone()),
             (recipient, recipient_account.clone()),
-            (system_program::id(), system_program_account()),
+            system_program(),
         ],
         &[
             Check::success(),
@@ -207,7 +207,7 @@ fn test_close_account() {
             &[
                 (key, account.clone()),
                 (incinerator::id(), AccountSharedData::default()),
-                (system_program::id(), system_program_account()),
+                system_program(),
             ],
             &[
                 Check::err(ProgramError::MissingRequiredSignature),
@@ -222,7 +222,7 @@ fn test_close_account() {
         &[
             (key, account.clone()),
             (incinerator::id(), AccountSharedData::default()),
-            (system_program::id(), system_program_account()),
+            system_program(),
         ],
         &[
             Check::success(),
@@ -287,7 +287,7 @@ fn test_cpi() {
                 (key, account.clone()),
                 (
                     cpi_target_program_id,
-                    create_program_account(&cpi_target_program_id),
+                    program_account(&cpi_target_program_id),
                 ),
             ],
             &[
@@ -312,7 +312,7 @@ fn test_cpi() {
                 (key, account.clone()),
                 (
                     cpi_target_program_id,
-                    create_program_account(&cpi_target_program_id),
+                    program_account(&cpi_target_program_id),
                 ),
             ],
             &[
@@ -336,7 +336,7 @@ fn test_cpi() {
                 (key, account.clone()),
                 (
                     cpi_target_program_id,
-                    create_program_account(&cpi_target_program_id),
+                    program_account(&cpi_target_program_id),
                 ),
             ],
             &[
@@ -353,7 +353,7 @@ fn test_cpi() {
             (key, account.clone()),
             (
                 cpi_target_program_id,
-                create_program_account(&cpi_target_program_id),
+                program_account(&cpi_target_program_id),
             ),
         ],
         &[