diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index 2964c09157..b46db0b8cc 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -113,7 +113,7 @@ pub struct Bank { /// staked nodes on epoch boundaries, saved off when a bank.slot() is at /// a leader schedule boundary - epoch_vote_states: HashMap>, + epoch_vote_accounts: HashMap>, } impl Bank { @@ -129,9 +129,9 @@ impl Bank { // genesis needs stakes for all epochs up to the epoch implied by // slot = 0 and genesis configuration - let vote_states = bank.vote_states(|_, _| true); + let vote_accounts = bank.vote_accounts(|k, v| Some((*k, v.clone()))); for i in 0..=bank.epoch_from_stakers_slot_offset() { - bank.epoch_vote_states.insert(i, vote_states.clone()); + bank.epoch_vote_accounts.insert(i, vote_accounts.clone()); } bank @@ -154,17 +154,17 @@ impl Bank { bank.accounts = Some(parent.accounts()); bank.accounts().new_from_parent(bank.slot, parent.slot); - bank.epoch_vote_states = { - let mut epoch_vote_states = parent.epoch_vote_states.clone(); + bank.epoch_vote_accounts = { + let mut epoch_vote_accounts = parent.epoch_vote_accounts.clone(); let epoch = bank.epoch_from_stakers_slot_offset(); // update epoch_vote_states cache // if my parent didn't populate for this epoch, we've // crossed a boundary - if epoch_vote_states.get(&epoch).is_none() { - epoch_vote_states.insert(epoch, bank.vote_states(|_, _| true)); + if epoch_vote_accounts.get(&epoch).is_none() { + epoch_vote_accounts.insert(epoch, bank.vote_accounts(|k, v| Some((*k, v.clone())))); } - epoch_vote_states + epoch_vote_accounts }; bank @@ -754,14 +754,27 @@ impl Bank { (self.slot + self.stakers_slot_offset) / self.slots_per_epoch } - pub fn epoch_vote_states(&self, epoch: u64, check: F) -> Option> + /// current vote accounts for this bank + pub fn vote_accounts(&self, filter: F) -> HashMap where - F: Fn(&Pubkey, &VoteState) -> Option<(Pubkey, T)>, + F: Fn(&Pubkey, &Account) -> Option<(Pubkey, T)>, { - self.epoch_vote_states.get(&epoch).map(|vote_states| { - vote_states + self.accounts() + .get_vote_accounts(self.slot) + .iter() + .filter_map(|(pubkey, account)| filter(pubkey, account)) + .collect() + } + + /// vote accounts for the specific epoch + pub fn epoch_vote_accounts(&self, epoch: u64, filter: F) -> Option> + where + F: Fn(&Pubkey, &Account) -> Option<(Pubkey, T)>, + { + self.epoch_vote_accounts.get(&epoch).map(|accounts| { + accounts .iter() - .filter_map(|(p, vote_state)| check(p, vote_state)) + .filter_map(|(pubkey, account)| filter(pubkey, account)) .collect() }) } @@ -1457,7 +1470,7 @@ mod tests { } #[test] - fn test_bank_epoch_vote_states() { + fn test_bank_epoch_vote_accounts() { let leader_id = Keypair::new().pubkey(); let leader_tokens = 2; let (mut genesis_block, _) = GenesisBlock::new_with_leader(5, leader_id, leader_tokens); @@ -1473,17 +1486,21 @@ mod tests { let parent = Arc::new(Bank::new(&genesis_block)); - let vote_states0 = parent.epoch_vote_states(0, |pubkey, vote_state| { - if vote_state.delegate_id == leader_id { - Some((*pubkey, true)) + let vote_accounts0 = parent.epoch_vote_accounts(0, |pubkey, account| { + if let Ok(vote_state) = VoteState::deserialize(&account.userdata) { + if vote_state.delegate_id == leader_id { + Some((*pubkey, true)) + } else { + None + } } else { None } }); - assert!(vote_states0.is_some()); - assert!(vote_states0.iter().len() != 0); + assert!(vote_accounts0.is_some()); + assert!(vote_accounts0.iter().len() != 0); - fn all(key: &Pubkey, _vote_state: &VoteState) -> Option<(Pubkey, Option<()>)> { + fn all(key: &Pubkey, _account: &Account) -> Option<(Pubkey, Option<()>)> { Some((*key, None)) } @@ -1492,7 +1509,7 @@ mod tests { if i > STAKERS_SLOT_OFFSET / SLOTS_PER_EPOCH { break; } - assert!(parent.epoch_vote_states(i, all).is_some()); + assert!(parent.epoch_vote_accounts(i, all).is_some()); i += 1; } @@ -1503,7 +1520,7 @@ mod tests { SLOTS_PER_EPOCH - (STAKERS_SLOT_OFFSET % SLOTS_PER_EPOCH), ); - assert!(child.epoch_vote_states(i, all).is_some()); + assert!(child.epoch_vote_accounts(i, all).is_some()); // child crosses epoch boundary but isn't the first slot in the epoch let child = Bank::new_from_parent( @@ -1511,7 +1528,7 @@ mod tests { leader_id, SLOTS_PER_EPOCH - (STAKERS_SLOT_OFFSET % SLOTS_PER_EPOCH) + 1, ); - assert!(child.epoch_vote_states(i, all).is_some()); + assert!(child.epoch_vote_accounts(i, all).is_some()); } }