|
1 | 1 | use { |
2 | 2 | super::{SlotLeader, stake_weighted_slot_leaders}, |
3 | | - itertools::Itertools, |
4 | 3 | solana_clock::Epoch, |
5 | | - solana_pubkey::Pubkey, |
| 4 | + solana_pubkey::{Pubkey, PubkeyHasherBuilder}, |
6 | 5 | solana_vote::vote_account::VoteAccountsHashMap, |
7 | 6 | std::{collections::HashMap, iter, num::NonZeroUsize, ops::Index}, |
8 | 7 | }; |
|
11 | 10 | pub struct LeaderSchedule { |
12 | 11 | slot_leaders: Vec<SlotLeader>, |
13 | 12 | // Inverted index from leader id to indices where they are the leader. |
14 | | - leader_slots_map: HashMap<Pubkey, Vec<usize>>, |
| 13 | + leader_slots_map: HashMap<Pubkey, Vec<usize>, PubkeyHasherBuilder>, |
15 | 14 | repeat: NonZeroUsize, |
16 | 15 | } |
17 | 16 |
|
@@ -47,24 +46,42 @@ impl LeaderSchedule { |
47 | 46 | }) |
48 | 47 | .collect(); |
49 | 48 | let slot_leaders = stake_weighted_slot_leaders(slot_leader_stakes, epoch, len, repeat); |
50 | | - Self::new_from_schedule(slot_leaders, repeat) |
| 49 | + Self { |
| 50 | + leader_slots_map: Self::invert_slot_leaders( |
| 51 | + &slot_leaders, |
| 52 | + Some(vote_accounts_map.len()), |
| 53 | + ), |
| 54 | + slot_leaders, |
| 55 | + repeat, |
| 56 | + } |
51 | 57 | } |
52 | 58 |
|
53 | 59 | pub fn new_from_schedule(slot_leaders: Vec<SlotLeader>, repeat: NonZeroUsize) -> Self { |
54 | | - let leader_slots_map = Self::invert_slot_leaders(&slot_leaders); |
| 60 | + let leader_slots_map = Self::invert_slot_leaders(&slot_leaders, None); |
55 | 61 | Self { |
56 | 62 | slot_leaders, |
57 | 63 | leader_slots_map, |
58 | 64 | repeat, |
59 | 65 | } |
60 | 66 | } |
61 | 67 |
|
62 | | - fn invert_slot_leaders(slot_leaders: &[SlotLeader]) -> HashMap<Pubkey, Vec<usize>> { |
63 | | - slot_leaders |
64 | | - .iter() |
65 | | - .enumerate() |
66 | | - .map(|(i, leader)| (leader.id, i)) |
67 | | - .into_group_map() |
| 68 | + fn invert_slot_leaders( |
| 69 | + slot_leaders: &[SlotLeader], |
| 70 | + nodes_len: Option<usize>, |
| 71 | + ) -> HashMap<Pubkey, Vec<usize>, PubkeyHasherBuilder> { |
| 72 | + let mut grouped_slot_leaders = match nodes_len { |
| 73 | + Some(nodes_len) => { |
| 74 | + HashMap::with_capacity_and_hasher(nodes_len, PubkeyHasherBuilder::default()) |
| 75 | + } |
| 76 | + None => HashMap::with_hasher(PubkeyHasherBuilder::default()), |
| 77 | + }; |
| 78 | + for (slot, leader) in slot_leaders.iter().enumerate() { |
| 79 | + grouped_slot_leaders |
| 80 | + .entry(leader.id) |
| 81 | + .and_modify(|slots: &mut Vec<usize>| slots.push(slot)) |
| 82 | + .or_insert(vec![slot]); |
| 83 | + } |
| 84 | + grouped_slot_leaders |
68 | 85 | } |
69 | 86 |
|
70 | 87 | pub fn get_slot_leaders(&self) -> impl Iterator<Item = &SlotLeader> { |
|
0 commit comments