refactor for total stake

program-runtime/src/invoke_context.rs

@@ -147,24 +147,27 @@ impl BpfAllocator {
 
 pub struct EnvironmentConfig<'a> {
     pub blockhash: Hash,
+    epoch_total_stake: Option<u64>,
+    epoch_vote_accounts: Option<&'a VoteAccountsHashMap>,
     pub feature_set: Arc<FeatureSet>,
     pub lamports_per_signature: u64,
-    vote_accounts: Option<&'a VoteAccountsHashMap>,
     sysvar_cache: &'a SysvarCache,
 }
 impl<'a> EnvironmentConfig<'a> {
     pub fn new(
         blockhash: Hash,
+        epoch_total_stake: Option<u64>,
+        epoch_vote_accounts: Option<&'a VoteAccountsHashMap>,
         feature_set: Arc<FeatureSet>,
         lamports_per_signature: u64,
-        vote_accounts: Option<&'a VoteAccountsHashMap>,
         sysvar_cache: &'a SysvarCache,
     ) -> Self {
         Self {
             blockhash,
+            epoch_total_stake,
+            epoch_vote_accounts,
             feature_set,
             lamports_per_signature,
-            vote_accounts,
             sysvar_cache,
         }
     }
@@ -618,9 +621,14 @@ impl<'a> InvokeContext<'a> {
         self.environment_config.sysvar_cache
     }
 
-    /// Get cached vote accounts.
-    pub fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap> {
-        self.environment_config.vote_accounts
+    /// Get cached epoch total stake.
+    pub fn get_epoch_total_stake(&self) -> Option<u64> {
+        self.environment_config.epoch_total_stake
+    }
+
+    /// Get cached epoch vote accounts.
+    pub fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap> {
+        self.environment_config.epoch_vote_accounts
     }
 
     // Should alignment be enforced during user pointer translation
@@ -719,9 +727,10 @@ macro_rules! with_mock_invoke_context {
         });
         let environment_config = EnvironmentConfig::new(
             Hash::default(),
+            None,
+            None,
             Arc::new(FeatureSet::all_enabled()),
             0,
-            /* vote_accounts */ None,
             &sysvar_cache,
         );
         let program_cache_for_tx_batch = ProgramCacheForTxBatch::default();

programs/bpf_loader/src/syscalls/mod.rs

@@ -2024,51 +2024,75 @@ declare_builtin_function!(
     SyscallGetEpochStake,
     fn rust(
         invoke_context: &mut InvokeContext,
-        vote_address: u64,
+        var_addr: u64,
         _arg2: u64,
         _arg3: u64,
         _arg4: u64,
         _arg5: u64,
         memory_mapping: &mut MemoryMapping,
     ) -> Result<u64, Error> {
-        // Compute units, as specified by SIMD-0133.
-        // cu = syscall_base_cost
-        //     + floor(32/cpi_bytes_per_unit)
-        //     + mem_op_base_cost
         let compute_budget = invoke_context.get_compute_budget();
-        let compute_units = compute_budget
-            .syscall_base_cost
-            .saturating_add(
-                (PUBKEY_BYTES as u64)
-                    .checked_div(compute_budget.cpi_bytes_per_unit)
-                    .unwrap_or(u64::MAX),
+
+        if var_addr == 0 {
+            // As specified by SIMD-0133: If `var_addr` is a null pointer:
+            //
+            // Compute units:
+            //
+            // ```
+            // syscall_base
+            // ```
+            let compute_units = compute_budget.syscall_base_cost;
+            consume_compute_meter(invoke_context, compute_units)?;
+            //
+            // Control flow:
+            //
+            // - The syscall aborts the virtual machine if:
+            //     - Compute budget is exceeded.
+            // - Otherwise, the syscall returns a `u64` integer representing the total active
+            //   stake on the cluster for the current epoch.
+            Ok(invoke_context.get_epoch_total_stake().unwrap_or(0))
+        } else {
+            // As specified by SIMD-0133: If `var_addr` is _not_ a null pointer:
+            //
+            // Compute units:
+            //
+            // ```
+            // syscall_base + floor(PUBKEY_BYTES/cpi_bytes_per_unit) + mem_op_base
+            // ```
+            let compute_units = compute_budget
+                .syscall_base_cost
+                .saturating_add(
+                    (PUBKEY_BYTES as u64)
+                        .checked_div(compute_budget.cpi_bytes_per_unit)
+                        .unwrap_or(u64::MAX),
+                )
+                .saturating_add(compute_budget.mem_op_base_cost);
+            consume_compute_meter(invoke_context, compute_units)?;
+            //
+            // Control flow:
+            //
+            // - The syscall aborts the virtual machine if:
+            //     - Not all bytes in VM memory range `[vote_addr, vote_addr + 32)` are
+            //       readable.
+            //     - Compute budget is exceeded.
+            // - Otherwise, the syscall returns a `u64` integer representing the total active
+            //   stake delegated to the vote account at the provided address.
+            //   If the provided vote address corresponds to an account that is not a vote
+            //   account or does not exist, the syscall will return `0` for active stake.
+            let check_aligned = invoke_context.get_check_aligned();
+            let vote_address = translate_type::<Pubkey>(memory_mapping, var_addr, check_aligned)?;
+
+            Ok(
+                if let Some(vote_accounts) = invoke_context.get_epoch_vote_accounts() {
+                    vote_accounts
+                        .get(vote_address)
+                        .map(|(stake, _)| *stake)
+                        .unwrap_or(0)
+                } else {
+                    0
+                },
             )
-            .saturating_add(compute_budget.mem_op_base_cost);
-
-        consume_compute_meter(invoke_context, compute_units)?;
-
-        // Control flow, as specified by SIMD-0133.
-        // * The syscall aborts the virtual machine if not all bytes in VM
-        //   memory range `[vote_addr, vote_addr + 32)` are readable.
-        // * Otherwise, the syscall returns a `u64` integer representing the
-        //   total active stake delegated to the vote account at the provided
-        //   address.
-        //   * If the provided vote address corresponds to an account that is
-        //     not a vote account or does not exist, the syscall will return
-        //     `0` for active stake.
-        let check_aligned = invoke_context.get_check_aligned();
-        let vote_address = translate_type::<Pubkey>(memory_mapping, vote_address, check_aligned)?;
-
-        Ok(
-            if let Some(vote_accounts) = invoke_context.get_vote_accounts() {
-                vote_accounts
-                    .get(vote_address)
-                    .map(|(stake, _)| *stake)
-                    .unwrap_or(0)
-            } else {
-                0
-            },
-        )
+        }
     }
 );
 
@@ -4744,7 +4768,50 @@ mod tests {
     }
 
     #[test]
-    fn test_syscall_get_epoch_stake() {
+    fn test_syscall_get_epoch_stake_total_stake() {
+        let config = Config::default();
+        let mut compute_budget = ComputeBudget::default();
+        let sysvar_cache = Arc::<SysvarCache>::default();
+
+        let expected_total_stake = 200_000_000_000_000u64;
+        // Compute units, as specified by SIMD-0133.
+        // cu = syscall_base_cost
+        let expected_cus = compute_budget.syscall_base_cost;
+
+        // Set the compute budget to the expected CUs to ensure the syscall
+        // doesn't exceed the expected usage.
+        compute_budget.compute_unit_limit = expected_cus;
+
+        with_mock_invoke_context!(invoke_context, transaction_context, vec![]);
+        invoke_context.environment_config = EnvironmentConfig::new(
+            Hash::default(),
+            Some(expected_total_stake),
+            None, // Vote accounts are not needed for this test.
+            Arc::<FeatureSet>::default(),
+            0,
+            &sysvar_cache,
+        );
+
+        let null_pointer_var = std::ptr::null::<Pubkey>() as u64;
+
+        let mut memory_mapping = MemoryMapping::new(vec![], &config, &SBPFVersion::V2).unwrap();
+
+        let result = SyscallGetEpochStake::rust(
+            &mut invoke_context,
+            null_pointer_var,
+            0,
+            0,
+            0,
+            0,
+            &mut memory_mapping,
+        )
+        .unwrap();
+
+        assert_eq!(result, expected_total_stake);
+    }
+
+    #[test]
+    fn test_syscall_get_epoch_stake_vote_account_stake() {
         let config = Config::default();
         let mut compute_budget = ComputeBudget::default();
         let sysvar_cache = Arc::<SysvarCache>::default();
@@ -4780,9 +4847,10 @@ mod tests {
         with_mock_invoke_context!(invoke_context, transaction_context, vec![]);
         invoke_context.environment_config = EnvironmentConfig::new(
             Hash::default(),
+            None, // Total stake is not needed for this test.
+            Some(&vote_accounts_map),
             Arc::<FeatureSet>::default(),
             0,
-            Some(&vote_accounts_map),
             &sysvar_cache,
         );
 

runtime/src/bank.rs

@@ -6155,6 +6155,13 @@ impl Bank {
         Some(self.epoch_stakes.get(&epoch)?.stakes().staked_nodes())
     }
 
+    /// Get the total epoch stake for the given epoch.
+    pub fn epoch_total_stake(&self, epoch: Epoch) -> Option<u64> {
+        self.epoch_stakes
+            .get(&epoch)
+            .map(|epoch_stakes| epoch_stakes.total_stake())
+    }
+
     /// vote accounts for the specific epoch along with the stake
     ///   attributed to each account
     pub fn epoch_vote_accounts(&self, epoch: Epoch) -> Option<&VoteAccountsHashMap> {
@@ -6838,7 +6845,11 @@ impl TransactionProcessingCallback for Bank {
         self.feature_set.clone()
     }
 
-    fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap> {
+    fn get_epoch_total_stake(&self) -> Option<u64> {
+        self.epoch_total_stake(self.epoch())
+    }
+
+    fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap> {
         self.epoch_vote_accounts(self.epoch())
     }
 

runtime/src/bank/builtins/core_bpf_migration/mod.rs

@@ -205,9 +205,10 @@ impl Bank {
                 &program_cache_for_tx_batch,
                 EnvironmentConfig::new(
                     Hash::default(),
+                    None,
+                    None,
                     self.feature_set.clone(),
                     0,
-                    None,
                     &sysvar_cache,
                 ),
                 None,

svm/src/account_loader.rs

@@ -545,7 +545,11 @@ mod tests {
             self.feature_set.clone()
         }
 
-        fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap> {
+        fn get_epoch_total_stake(&self) -> Option<u64> {
+            None
+        }
+
+        fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap> {
             None
         }
     }

svm/src/message_processor.rs

@@ -275,9 +275,10 @@ mod tests {
         let mut programs_modified_by_tx = ProgramCacheForTxBatch::default();
         let environment_config = EnvironmentConfig::new(
             Hash::default(),
+            None,
+            None,
             Arc::new(FeatureSet::all_enabled()),
             0,
-            None,
             &sysvar_cache,
         );
         let mut invoke_context = InvokeContext::new(
@@ -330,9 +331,10 @@ mod tests {
         let mut programs_modified_by_tx = ProgramCacheForTxBatch::default();
         let environment_config = EnvironmentConfig::new(
             Hash::default(),
+            None,
+            None,
             Arc::new(FeatureSet::all_enabled()),
             0,
-            None,
             &sysvar_cache,
         );
         let mut invoke_context = InvokeContext::new(
@@ -375,9 +377,10 @@ mod tests {
         let mut programs_modified_by_tx = ProgramCacheForTxBatch::default();
         let environment_config = EnvironmentConfig::new(
             Hash::default(),
+            None,
+            None,
             Arc::new(FeatureSet::all_enabled()),
             0,
-            None,
             &sysvar_cache,
         );
         let mut invoke_context = InvokeContext::new(
@@ -511,9 +514,10 @@ mod tests {
         let mut programs_modified_by_tx = ProgramCacheForTxBatch::default();
         let environment_config = EnvironmentConfig::new(
             Hash::default(),
+            None,
+            None,
             Arc::new(FeatureSet::all_enabled()),
             0,
-            None,
             &sysvar_cache,
         );
         let mut invoke_context = InvokeContext::new(
@@ -551,9 +555,10 @@ mod tests {
         let mut programs_modified_by_tx = ProgramCacheForTxBatch::default();
         let environment_config = EnvironmentConfig::new(
             Hash::default(),
+            None,
+            None,
             Arc::new(FeatureSet::all_enabled()),
             0,
-            None,
             &sysvar_cache,
         );
         let mut invoke_context = InvokeContext::new(
@@ -588,9 +593,10 @@ mod tests {
         let mut programs_modified_by_tx = ProgramCacheForTxBatch::default();
         let environment_config = EnvironmentConfig::new(
             Hash::default(),
+            None,
+            None,
             Arc::new(FeatureSet::all_enabled()),
             0,
-            None,
             &sysvar_cache,
         );
         let mut invoke_context = InvokeContext::new(
@@ -686,9 +692,10 @@ mod tests {
         let mut programs_modified_by_tx = ProgramCacheForTxBatch::default();
         let environment_config = EnvironmentConfig::new(
             Hash::default(),
+            None,
+            None,
             Arc::new(FeatureSet::all_enabled()),
             0,
-            None,
             &sysvar_cache,
         );
         let mut invoke_context = InvokeContext::new(

svm/src/program_loader.rs

@@ -286,7 +286,11 @@ mod tests {
             self.feature_set.clone()
         }
 
-        fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap> {
+        fn get_epoch_total_stake(&self) -> Option<u64> {
+            None
+        }
+
+        fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap> {
             None
         }
 

svm/src/transaction_processing_callback.rs

@@ -20,7 +20,9 @@ pub trait TransactionProcessingCallback {
 
     fn get_feature_set(&self) -> Arc<FeatureSet>;
 
-    fn get_vote_accounts(&self) -> Option<&VoteAccountsHashMap>;
+    fn get_epoch_total_stake(&self) -> Option<u64>;
+
+    fn get_epoch_vote_accounts(&self) -> Option<&VoteAccountsHashMap>;
 
     fn get_program_match_criteria(&self, _program: &Pubkey) -> ProgramCacheMatchCriteria {
         ProgramCacheMatchCriteria::NoCriteria