replace some btreemaps with hashmaps

This allows us to use `with_capacity` which hopefully prevents allocations

Author: akhi3030

core/src/bls_sigverify/bls_sigverifier.rs

@@ -256,7 +256,10 @@ impl BLSSigVerifier {
             let mut guard = self.consensus_rewards.write();
             let root_slot = root_bank.slot();
             for v in rewards_votes {
-                guard.add_vote_message(root_slot, v);
+                if let Some(rank_map) = get_key_to_rank_map(&root_bank, v.vote.slot()) {
+                    let max_validators = rank_map.len();
+                    guard.add_vote(root_slot, max_validators, v);
+                }
             }
         }
 
@@ -306,7 +309,7 @@ impl BLSSigVerifier {
                 .filter_map(|vote| {
                     guard
                         .wants_vote(root_slot, &vote.vote_message)
-                        .then_some(vote.vote_message.clone())
+                        .then_some(vote.vote_message)
                 })
                 .collect()
         };

votor/src/consensus_rewards.rs

@@ -58,7 +58,7 @@ impl ConsensusRewards {
     }
 
     /// Adds received [`VoteMessage`]s from other nodes.
-    pub fn add_vote_message(&mut self, root_slot: Slot, vote: VoteMessage) {
+    pub fn add_vote(&mut self, root_slot: Slot, max_validators: usize, vote: VoteMessage) {
         // drop old state no longer needed
         self.votes = self.votes.split_off(
             &(root_slot
@@ -72,7 +72,7 @@ impl ConsensusRewards {
         }
         self.votes
             .entry(vote.vote.slot())
-            .or_default()
+            .or_insert(Entry::new(max_validators))
             .add_vote(vote);
     }
 

votor/src/consensus_rewards/entry.rs

@@ -9,12 +9,14 @@ use {
         rewards_certificate::{NotarRewardCertificate, SkipRewardCertificate},
         vote::Vote,
     },
-    std::collections::BTreeMap,
+    std::collections::HashMap,
 };
 
 /// Builds a signature and bitmap suitable for creating a rewards certificate.
-fn build_sig_bitmap(votes: &BTreeMap<u16, VoteMessage>) -> Option<(BLSSignature, Vec<u8>)> {
-    let max_rank = votes.last_key_value().map(|(rank, _)| rank).cloned()?;
+fn build_sig_bitmap(
+    votes: &HashMap<u16, VoteMessage>,
+    max_rank: u16,
+) -> Option<(BLSSignature, Vec<u8>)> {
     let mut bitmap = BitVec::repeat(false, max_rank as usize);
     for vote in votes.keys() {
         bitmap.set(*vote as usize, true);
@@ -29,20 +31,35 @@ fn build_sig_bitmap(votes: &BTreeMap<u16, VoteMessage>) -> Option<(BLSSignature,
 }
 
 /// Per slot container for storing notar and skip votes for creating rewards certificates.
-#[derive(Default)]
 pub(super) struct Entry {
-    /// map from validator rank to the skip vote.
-    skip: BTreeMap<u16, VoteMessage>,
-    /// notar votes are indexed by block id as different validators may vote for different blocks.
-    notar: BTreeMap<Hash, BTreeMap<u16, VoteMessage>>,
+    /// Map from validator rank to the skip vote.
+    skip: HashMap<u16, VoteMessage>,
+    /// Largest ranked validator that voted skip.
+    skip_max_rank: u16,
+    /// Notar votes are indexed by block id as different validators may vote for different blocks.
+    /// Per block id, store a map from validator rank to notar vote and also the largest ranked validator that voted notar.
+    notar: HashMap<Hash, (HashMap<u16, VoteMessage>, u16)>,
+    /// Maximum number of validators for the slot this entry is working on.
+    max_validators: usize,
 }
 
 impl Entry {
+    /// Creates a new instance of [`Entry`].
+    pub(super) fn new(max_validators: usize) -> Self {
+        Self {
+            skip: HashMap::with_capacity(max_validators),
+            skip_max_rank: 0,
+            // under normal operations, all validators should vote for a single block id, still allocate space for a few more to hopefully avoid allocations.
+            notar: HashMap::with_capacity(5),
+            max_validators,
+        }
+    }
+
     pub(super) fn wants_vote(&self, vote: &VoteMessage) -> bool {
         match vote.vote {
             Vote::Skip(_) => !self.skip.contains_key(&vote.rank),
             Vote::Notarize(notar) => {
-                let Some(notar) = self.notar.get(&notar.block_id) else {
+                let Some((notar, _)) = self.notar.get(&notar.block_id) else {
                     return true;
                 };
                 !notar.contains_key(&vote.rank)
@@ -57,13 +74,16 @@ impl Entry {
     pub(super) fn add_vote(&mut self, vote: VoteMessage) {
         match vote.vote {
             Vote::Notarize(notar) => {
-                self.notar
+                let (map, max_rank) = self
+                    .notar
                     .entry(notar.block_id)
-                    .or_default()
-                    .insert(vote.rank, vote);
+                    .or_insert((HashMap::with_capacity(self.max_validators), 0));
+                map.insert(vote.rank, vote);
+                *max_rank = (*max_rank).max(vote.rank);
             }
             Vote::Skip(_) => {
                 self.skip.insert(vote.rank, vote);
+                self.skip_max_rank = self.skip_max_rank.max(vote.rank);
             }
             _ => (),
         }
@@ -76,29 +96,31 @@ impl Entry {
         Option<SkipRewardCertificate>,
         Option<NotarRewardCertificate>,
     ) {
-        let skip = build_sig_bitmap(&self.skip).map(|(signature, bitmap)| SkipRewardCertificate {
-            slot,
-            signature,
-            bitmap,
+        let skip = build_sig_bitmap(&self.skip, self.skip_max_rank).map(|(signature, bitmap)| {
+            SkipRewardCertificate {
+                slot,
+                signature,
+                bitmap,
+            }
         });
 
         // we can only submit one notar rewards certificate but different validators may vote for different blocks and we cannot combine notar votes for different blocks together in one cert.
         // pick the block_id with most votes.
         // in practice, all validators should have voted for the same block otherwise, we have evidence of leader equivocation.
         // TODO: collect metrics for when equivocation is detected.
         let mut notar = None;
-        for (block_id, map) in &self.notar {
+        for (block_id, (map, max_rank)) in &self.notar {
             match notar {
-                None => notar = Some((block_id, map)),
-                Some((_, notar_map)) => {
+                None => notar = Some((block_id, map, max_rank)),
+                Some((_, notar_map, _)) => {
                     if map.len() > notar_map.len() {
-                        notar = Some((block_id, map));
+                        notar = Some((block_id, map, max_rank));
                     }
                 }
             }
         }
-        let notar = notar.and_then(|(block_id, votes)| {
-            build_sig_bitmap(votes).map(|(signature, bitmap)| NotarRewardCertificate {
+        let notar = notar.and_then(|(block_id, votes, max_rank)| {
+            build_sig_bitmap(votes, *max_rank).map(|(signature, bitmap)| NotarRewardCertificate {
                 slot,
                 block_id: *block_id,
                 signature,