use callback to get epoch vote stake

program-runtime/src/invoke_context.rs

@@ -43,7 +43,6 @@ use {
     std::{
         alloc::Layout,
         cell::RefCell,
-        collections::HashMap,
         fmt::{self, Debug},
         rc::Rc,
     },
@@ -152,7 +151,7 @@ pub struct EnvironmentConfig<'a> {
     pub blockhash: Hash,
     pub blockhash_lamports_per_signature: u64,
     epoch_total_stake: u64,
-    epoch_vote_stake: &'a HashMap<Pubkey, u64>,
+    get_epoch_stake_callback: &'a dyn Fn(&'a Pubkey) -> u64,
     pub feature_set: Arc<FeatureSet>,
     sysvar_cache: &'a SysvarCache,
 }
@@ -161,15 +160,15 @@ impl<'a> EnvironmentConfig<'a> {
         blockhash: Hash,
         blockhash_lamports_per_signature: u64,
         epoch_total_stake: u64,
-        epoch_vote_stake: &'a HashMap<Pubkey, u64>,
+        get_epoch_stake_callback: &'a dyn Fn(&'a Pubkey) -> u64,
         feature_set: Arc<FeatureSet>,
         sysvar_cache: &'a SysvarCache,
     ) -> Self {
         Self {
             blockhash,
             blockhash_lamports_per_signature,
             epoch_total_stake,
-            epoch_vote_stake,
+            get_epoch_stake_callback,
             feature_set,
             sysvar_cache,
         }
@@ -662,9 +661,9 @@ impl<'a> InvokeContext<'a> {
         self.environment_config.epoch_total_stake
     }
 
-    /// Get cached stake for epoch vote accounts.
-    pub fn get_epoch_vote_stake(&self) -> &HashMap<Pubkey, u64> {
-        self.environment_config.epoch_vote_stake
+    /// Get cached stake for the epoch vote account.
+    pub fn get_epoch_vote_stake(&self, pubkey: &'a Pubkey) -> u64 {
+        (self.environment_config.get_epoch_stake_callback)(pubkey)
     }
 
     // Should alignment be enforced during user pointer translation
@@ -762,12 +761,11 @@ macro_rules! with_mock_invoke_context {
                 }
             }
         });
-        let epoch_vote_stake = HashMap::default();
         let environment_config = EnvironmentConfig::new(
             Hash::default(),
             0,
             0,
-            &epoch_vote_stake,
+            &|_| 0,
             Arc::new(FeatureSet::all_enabled()),
             &sysvar_cache,
         );

programs/bpf_loader/src/lib.rs

@@ -1642,7 +1642,7 @@ mod tests {
             rent::Rent,
             system_program, sysvar,
         },
-        std::{collections::HashMap, fs::File, io::Read, ops::Range, sync::atomic::AtomicU64},
+        std::{fs::File, io::Read, ops::Range, sync::atomic::AtomicU64},
     };
 
     fn process_instruction(

programs/bpf_loader/src/serialization.rs

@@ -632,7 +632,6 @@ mod tests {
         },
         std::{
             cell::RefCell,
-            collections::HashMap,
             mem::transmute,
             rc::Rc,
             slice::{self, from_raw_parts, from_raw_parts_mut},

programs/bpf_loader/src/syscalls/cpi.rs

@@ -1614,7 +1614,6 @@ mod tests {
         },
         std::{
             cell::{Cell, RefCell},
-            collections::HashMap,
             mem, ptr,
             rc::Rc,
             slice,

programs/bpf_loader/src/syscalls/mod.rs

@@ -2103,11 +2103,7 @@ declare_builtin_function!(
             let check_aligned = invoke_context.get_check_aligned();
             let vote_address = translate_type::<Pubkey>(memory_mapping, var_addr, check_aligned)?;
 
-            Ok(invoke_context
-                .get_epoch_vote_stake()
-                .get(vote_address)
-                .copied()
-                .unwrap_or(0))
+            Ok(invoke_context.get_epoch_vote_stake(vote_address))
         }
     }
 );
@@ -2142,7 +2138,7 @@ mod tests {
                 last_restart_slot::LastRestartSlot,
             },
         },
-        std::{collections::HashMap, mem, str::FromStr},
+        std::{mem, str::FromStr},
         test_case::test_case,
     };
 
@@ -4792,12 +4788,11 @@ mod tests {
         compute_budget.compute_unit_limit = expected_cus;
 
         with_mock_invoke_context!(invoke_context, transaction_context, vec![]);
-        let epoch_vote_stake = HashMap::default();
         invoke_context.environment_config = EnvironmentConfig::new(
             Hash::default(),
             0,
             expected_total_stake,
-            &epoch_vote_stake, // Vote accounts are not needed for this test.
+            &|_| 0, // Vote accounts are not needed for this test.
             Arc::<FeatureSet>::default(),
             &sysvar_cache,
         );
@@ -4840,15 +4835,19 @@ mod tests {
         compute_budget.compute_unit_limit = expected_cus;
 
         let vote_address = Pubkey::new_unique();
-        let mut vote_accounts_map = HashMap::new();
-        vote_accounts_map.insert(vote_address, expected_epoch_stake);
-
         with_mock_invoke_context!(invoke_context, transaction_context, vec![]);
+        let callback = |pubkey: &Pubkey| {
+            if *pubkey == vote_address {
+                expected_epoch_stake
+            } else {
+                0
+            }
+        };
         invoke_context.environment_config = EnvironmentConfig::new(
             Hash::default(),
             0,
             0, // Total stake is not needed for this test.
-            &vote_accounts_map,
+            &callback,
             Arc::<FeatureSet>::default(),
             &sysvar_cache,
         );

programs/sbf/benches/bpf_loader.rs

@@ -46,7 +46,7 @@ use {
         signature::Signer,
         transaction_context::InstructionAccount,
     },
-    std::{collections::HashMap, mem, sync::Arc},
+    std::{mem, sync::Arc},
     test::Bencher,
 };
 

programs/stake/src/stake_state.rs

@@ -1455,7 +1455,6 @@ mod tests {
             stake::state::warmup_cooldown_rate,
             sysvar::{epoch_schedule, SysvarId},
         },
-        std::collections::HashMap,
         test_case::test_case,
     };
 

programs/system/src/system_instruction.rs

@@ -257,7 +257,6 @@ mod test {
             system_program,
             transaction_context::InstructionAccount,
         },
-        std::collections::HashMap,
     };
 
     pub const NONCE_ACCOUNT_INDEX: IndexOfAccount = 0;

programs/system/src/system_processor.rs

@@ -570,7 +570,7 @@ mod tests {
         solana_program_runtime::{
             invoke_context::mock_process_instruction, with_mock_invoke_context,
         },
-        std::collections::{BinaryHeap, HashMap},
+        std::collections::BinaryHeap,
     };
 
     impl From<Pubkey> for Address {

runtime/src/bank.rs

@@ -1150,10 +1150,6 @@ impl Bank {
         bank.update_last_restart_slot();
         bank.transaction_processor
             .fill_missing_sysvar_cache_entries(&bank);
-        bank.transaction_processor.set_epoch_stake_information(
-            bank.get_current_epoch_total_stake(),
-            Arc::new(bank.get_epoch_vote_stake()),
-        );
         bank
     }
 
@@ -1468,13 +1464,6 @@ impl Bank {
             .stats
             .reset();
 
-        if parent.epoch() < new.epoch() {
-            new.transaction_processor.set_epoch_stake_information(
-                new.get_current_epoch_total_stake(),
-                Arc::new(new.get_epoch_vote_stake()),
-            )
-        }
-
         new
     }
 
@@ -1860,11 +1849,6 @@ impl Bank {
         assert_eq!(bank.epoch_schedule, genesis_config.epoch_schedule);
         assert_eq!(bank.epoch, bank.epoch_schedule.get_epoch(bank.slot));
 
-        bank.transaction_processor.set_epoch_stake_information(
-            bank.get_current_epoch_total_stake(),
-            Arc::new(bank.get_epoch_vote_stake()),
-        );
-
         datapoint_info!(
             "bank-new-from-fields",
             (
@@ -3784,6 +3768,7 @@ impl Bank {
         let processing_environment = TransactionProcessingEnvironment {
             blockhash,
             blockhash_lamports_per_signature,
+            epoch_total_stake: self.get_current_epoch_total_stake(),
             feature_set: Arc::clone(&self.feature_set),
             fee_lamports_per_signature: self.fee_structure.lamports_per_signature,
             rent_collector: Some(&rent_collector_with_metrics),
@@ -5227,13 +5212,6 @@ impl Bank {
             );
     }
 
-    fn get_epoch_vote_stake(&self) -> HashMap<Pubkey, u64> {
-        self.get_current_epoch_vote_accounts()
-            .iter()
-            .map(|(address, (stake, _))| (*address, *stake))
-            .collect()
-    }
-
     pub fn set_inflation(&self, inflation: Inflation) {
         *self.inflation.write().unwrap() = inflation;
     }
@@ -7151,6 +7129,13 @@ impl TransactionProcessingCallback for Bank {
             self.inspect_account_for_accounts_lt_hash(address, &account_state, is_writable);
         }
     }
+
+    fn get_epoch_stake(&self, vote_address: &Pubkey) -> u64 {
+        self.get_current_epoch_vote_accounts()
+            .get(vote_address)
+            .map(|(stake, _)| (*stake))
+            .unwrap_or(0)
+    }
 }
 
 #[cfg(feature = "dev-context-only-utils")]

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

@@ -21,7 +21,7 @@ use {
         transaction_context::TransactionContext,
     },
     source_buffer::SourceBuffer,
-    std::{cmp::Ordering, collections::HashMap, sync::atomic::Ordering::Relaxed},
+    std::{cmp::Ordering, sync::atomic::Ordering::Relaxed},
     target_builtin::TargetBuiltin,
     target_core_bpf::TargetCoreBpf,
 };
@@ -196,16 +196,14 @@ impl Bank {
                 compute_budget.max_instruction_trace_length,
             );
 
-            let epoch_vote_stake: HashMap<Pubkey, u64> = HashMap::default();
-
             let mut dummy_invoke_context = InvokeContext::new(
                 &mut dummy_transaction_context,
                 &mut program_cache_for_tx_batch,
                 EnvironmentConfig::new(
                     Hash::default(),
                     0,
                     0,
-                    &epoch_vote_stake,
+                    &|_| 0,
                     self.feature_set.clone(),
                     &sysvar_cache,
                 ),

svm/doc/spec.md

@@ -94,8 +94,6 @@ a `TransactionBatchProcessor` object the client need to specify the
 - `program_cache: Arc<RwLock<ProgramCache<FG>>>` is a reference to
     a ProgramCache instance. All on chain programs used in transaction
     batch execution are loaded from the program cache.
-- `epoch_total_stake`: The total stake for the current epoch.
-- `epoch_vote_stake`: The stake for vote accounts for the current epoch.
 
 In addition, `TransactionBatchProcessor` needs an instance of
 `SysvarCache` and a set of pubkeys of builtin program IDs.
@@ -135,6 +133,8 @@ pub trait TransactionProcessingCallback {
     fn get_account_shared_data(&self, pubkey: &Pubkey) -> Option<AccountSharedData>;
 
     fn add_builtin_account(&self, _name: &str, _program_id: &Pubkey) {}
+
+    fn get_epoch_stake(&self, _vote_address: &Pubkey) -> u64; 
 }
 ```
 
@@ -174,6 +174,7 @@ the runtime environment to use for processing transactions.
 
 - `blockhash`: The blockhash to use for the transaction batch.
 - `feature_set`: Runtime feature set to use for the transaction batch.
+- `epoch_total_stake`: The total stake for the current epoch.
 - `fee_structure`: Fee structure to use for assessing transaction fees.
 - `lamports_per_signature`: Lamports per signature to charge per transaction.
 - `rent_collector`: Rent collector to use for the transaction batch.

svm/examples/json-rpc/server/src/rpc_process.rs

@@ -538,6 +538,7 @@ impl JsonRpcRequestProcessor {
         let processing_environment = TransactionProcessingEnvironment {
             blockhash,
             blockhash_lamports_per_signature: lamports_per_signature,
+            epoch_total_stake: 0,
             feature_set: Arc::clone(&bank.feature_set),
             fee_lamports_per_signature: lamports_per_signature,
             rent_collector: None,

svm/examples/paytube/src/lib.rs

@@ -148,6 +148,7 @@ impl PayTubeChannel {
         let processing_environment = TransactionProcessingEnvironment {
             blockhash: Hash::default(),
             blockhash_lamports_per_signature: fee_structure.lamports_per_signature,
+            epoch_total_stake: 0,
             feature_set: Arc::new(feature_set),
             fee_lamports_per_signature: fee_structure.lamports_per_signature,
             rent_collector: Some(&rent_collector),