invoke context: add epoch stake

anza-xyz · Apr 18, 2024 · a233266 · a233266
1 parent d5b979c
commit a233266
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 0 deletions.
diff --git a/program-runtime/src/invoke_context.rs b/program-runtime/src/invoke_context.rs
@@ -169,6 +169,7 @@ pub struct InvokeContext<'a> {
     pub programs_loaded_for_tx_batch: &'a ProgramCacheForTxBatch,
     pub programs_modified_by_tx: &'a mut ProgramCacheForTxBatch,
     pub feature_set: Arc<FeatureSet>,
+    pub get_epoch_stake_callback: &'a dyn Fn(&'a Pubkey) -> u64,
     pub timings: ExecuteDetailsTimings,
     pub blockhash: Hash,
     pub lamports_per_signature: u64,
@@ -186,6 +187,7 @@ impl<'a> InvokeContext<'a> {
         programs_loaded_for_tx_batch: &'a ProgramCacheForTxBatch,
         programs_modified_by_tx: &'a mut ProgramCacheForTxBatch,
         feature_set: Arc<FeatureSet>,
+        get_epoch_stake_callback: &'a dyn Fn(&'a Pubkey) -> u64,
         blockhash: Hash,
         lamports_per_signature: u64,
     ) -> Self {
@@ -199,6 +201,7 @@ impl<'a> InvokeContext<'a> {
             programs_loaded_for_tx_batch,
             programs_modified_by_tx,
             feature_set,
+            get_epoch_stake_callback,
             timings: ExecuteDetailsTimings::default(),
             blockhash,
             lamports_per_signature,
@@ -685,6 +688,7 @@ macro_rules! with_mock_invoke_context {
             &programs_loaded_for_tx_batch,
             &mut programs_modified_by_tx,
             Arc::new(FeatureSet::all_enabled()),
+            &|_| 0,
             Hash::default(),
             0,
         );

diff --git a/runtime/src/bank/builtins/core_bpf_migration/mod.rs b/runtime/src/bank/builtins/core_bpf_migration/mod.rs
@@ -195,6 +195,7 @@ impl Bank {
                 &programs_loaded,
                 &mut programs_modified,
                 self.feature_set.clone(),
+                &|_| 0,
                 Hash::default(),
                 0,
             );

diff --git a/svm/src/message_processor.rs b/svm/src/message_processor.rs
@@ -280,6 +280,7 @@ mod tests {
             &programs_loaded_for_tx_batch,
             &mut programs_modified_by_tx,
             Arc::new(FeatureSet::all_enabled()),
+            &|_| 0,
             Hash::default(),
             0,
         );
@@ -331,6 +332,7 @@ mod tests {
             &programs_loaded_for_tx_batch,
             &mut programs_modified_by_tx,
             Arc::new(FeatureSet::all_enabled()),
+            &|_| 0,
             Hash::default(),
             0,
         );
@@ -372,6 +374,7 @@ mod tests {
             &programs_loaded_for_tx_batch,
             &mut programs_modified_by_tx,
             Arc::new(FeatureSet::all_enabled()),
+            &|_| 0,
             Hash::default(),
             0,
         );
@@ -504,6 +507,7 @@ mod tests {
             &programs_loaded_for_tx_batch,
             &mut programs_modified_by_tx,
             Arc::new(FeatureSet::all_enabled()),
+            &|_| 0,
             Hash::default(),
             0,
         );
@@ -540,6 +544,7 @@ mod tests {
             &programs_loaded_for_tx_batch,
             &mut programs_modified_by_tx,
             Arc::new(FeatureSet::all_enabled()),
+            &|_| 0,
             Hash::default(),
             0,
         );
@@ -573,6 +578,7 @@ mod tests {
             &programs_loaded_for_tx_batch,
             &mut programs_modified_by_tx,
             Arc::new(FeatureSet::all_enabled()),
+            &|_| 0,
             Hash::default(),
             0,
         );
@@ -667,6 +673,7 @@ mod tests {
             &programs_loaded_for_tx_batch,
             &mut programs_modified_by_tx,
             Arc::new(FeatureSet::all_enabled()),
+            &|_| 0,
             Hash::default(),
             0,
         );

diff --git a/svm/src/transaction_processor.rs b/svm/src/transaction_processor.rs
@@ -673,6 +673,7 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
             programs_loaded_for_tx_batch.latest_root_epoch,
         );
         let sysvar_cache = &self.sysvar_cache.read().unwrap();
+        let epoch_stake_callback = |pubkey| callback.get_epoch_stake(pubkey);
 
         let mut invoke_context = InvokeContext::new(
             &mut transaction_context,
@@ -682,6 +683,7 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
             programs_loaded_for_tx_batch,
             &mut programs_modified_by_tx,
             callback.get_feature_set(),
+            &epoch_stake_callback,
             blockhash,
             lamports_per_signature,
         );