diff --git a/leader-schedule/Cargo.toml b/leader-schedule/Cargo.toml index 0dc14348e273a9..4e2be25c3b4f9b 100644 --- a/leader-schedule/Cargo.toml +++ b/leader-schedule/Cargo.toml @@ -24,14 +24,13 @@ agave-random = { workspace = true } itertools = { workspace = true } rand_chacha = { workspace = true } solana-clock = { workspace = true } -solana-pubkey = { workspace = true } +solana-pubkey = { workspace = true, features = ["rand"] } solana-vote = { workspace = true } [dev-dependencies] bs58 = { workspace = true, features = ["alloc"] } rand = { workspace = true } sha2 = { workspace = true } -solana-pubkey = { workspace = true, features = ["rand"] } solana-vote = { path = "../vote", features = ["agave-unstable-api", "dev-context-only-utils"] } test-case = { workspace = true } diff --git a/leader-schedule/src/vote_keyed.rs b/leader-schedule/src/vote_keyed.rs index eae938156687b3..c82467566fe163 100644 --- a/leader-schedule/src/vote_keyed.rs +++ b/leader-schedule/src/vote_keyed.rs @@ -1,8 +1,7 @@ use { super::{SlotLeader, stake_weighted_slot_leaders}, - itertools::Itertools, solana_clock::Epoch, - solana_pubkey::Pubkey, + solana_pubkey::{Pubkey, PubkeyHasherBuilder}, solana_vote::vote_account::VoteAccountsHashMap, std::{collections::HashMap, iter, num::NonZeroUsize, ops::Index}, }; @@ -11,7 +10,7 @@ use { pub struct LeaderSchedule { slot_leaders: Vec, // Inverted index from leader id to indices where they are the leader. - leader_slots_map: HashMap>, + leader_slots_map: HashMap, PubkeyHasherBuilder>, repeat: NonZeroUsize, } @@ -47,11 +46,18 @@ impl LeaderSchedule { }) .collect(); let slot_leaders = stake_weighted_slot_leaders(slot_leader_stakes, epoch, len, repeat); - Self::new_from_schedule(slot_leaders, repeat) + Self { + leader_slots_map: Self::invert_slot_leaders( + &slot_leaders, + Some(vote_accounts_map.len()), + ), + slot_leaders, + repeat, + } } pub fn new_from_schedule(slot_leaders: Vec, repeat: NonZeroUsize) -> Self { - let leader_slots_map = Self::invert_slot_leaders(&slot_leaders); + let leader_slots_map = Self::invert_slot_leaders(&slot_leaders, None); Self { slot_leaders, leader_slots_map, @@ -59,12 +65,23 @@ impl LeaderSchedule { } } - fn invert_slot_leaders(slot_leaders: &[SlotLeader]) -> HashMap> { - slot_leaders - .iter() - .enumerate() - .map(|(i, leader)| (leader.id, i)) - .into_group_map() + fn invert_slot_leaders( + slot_leaders: &[SlotLeader], + nodes_len: Option, + ) -> HashMap, PubkeyHasherBuilder> { + let mut grouped_slot_leaders = match nodes_len { + Some(nodes_len) => { + HashMap::with_capacity_and_hasher(nodes_len, PubkeyHasherBuilder::default()) + } + None => HashMap::default(), + }; + for (index, leader) in slot_leaders.iter().enumerate() { + grouped_slot_leaders + .entry(leader.id) + .and_modify(|indices: &mut Vec<_>| indices.push(index)) + .or_insert(vec![index]); + } + grouped_slot_leaders } pub fn get_slot_leaders(&self) -> impl Iterator {