[votor] Replace VoteCertificate with either VoteCertificateBuilder or CertificateMessage (anza-xyz#304)

Author: samkim-crypto

votor/src/certificate_pool.rs

@@ -4,7 +4,7 @@ use {
         certificate_pool::{
             parent_ready_tracker::ParentReadyTracker,
             stats::CertificatePoolStats,
-            vote_certificate::{CertificateError, VoteCertificate},
+            vote_certificate_builder::{CertificateError, VoteCertificateBuilder},
             vote_pool::{DuplicateBlockVotePool, SimpleVotePool, VotePool, VotePoolType},
         },
         conflicting_types,
@@ -37,7 +37,7 @@ use {
 
 pub mod parent_ready_tracker;
 mod stats;
-mod vote_certificate;
+mod vote_certificate_builder;
 mod vote_pool;
 
 impl VoteType {
@@ -233,16 +233,16 @@ impl CertificatePool {
             if accumulated_stake as f64 / (total_stake as f64) < limit {
                 continue;
             }
-            let mut vote_certificate = VoteCertificate::new(cert_id);
+            let mut vote_certificate_builder = VoteCertificateBuilder::new(cert_id);
             vote_types.iter().for_each(|vote_type| {
                 if let Some(vote_pool) = self.vote_pools.get(&(slot, *vote_type)) {
                 match vote_pool {
-                    VotePoolType::SimpleVotePool(pool) => pool.add_to_certificate(&mut vote_certificate),
-                    VotePoolType::DuplicateBlockVotePool(pool) => pool.add_to_certificate(block_id.as_ref().expect("Duplicate block pool for {vote_type:?} expects a block id for certificate {cert_id:?}"), &mut vote_certificate),
+                    VotePoolType::SimpleVotePool(pool) => pool.add_to_certificate(&mut vote_certificate_builder),
+                    VotePoolType::DuplicateBlockVotePool(pool) => pool.add_to_certificate(block_id.as_ref().expect("Duplicate block pool for {vote_type:?} expects a block id for certificate {cert_id:?}"), &mut vote_certificate_builder),
                 };
             }
             });
-            let new_cert = Arc::new(vote_certificate.certificate());
+            let new_cert = Arc::new(vote_certificate_builder.build());
             self.send_and_insert_certificate(cert_id, new_cert.clone(), events)?;
             self.stats
                 .incr_cert_type(new_cert.certificate.certificate_type, true);

votor/src/certificate_pool/vote_certificate.rs renamed to votor/src/certificate_pool/vote_certificate_builder.rs

@@ -1,9 +1,7 @@
 use {
     crate::CertificateId,
     bitvec::prelude::*,
-    solana_bls_signatures::{
-        BlsError, Pubkey as BlsPubkey, PubkeyProjective, Signature, SignatureProjective,
-    },
+    solana_bls_signatures::{BlsError, Pubkey as BlsPubkey, PubkeyProjective, SignatureProjective},
     solana_runtime::epoch_stakes::BLSPubkeyToRankMap,
     solana_vote::alpenglow::{
         bls_message::{CertificateMessage, VoteMessage},
@@ -29,70 +27,56 @@ pub enum CertificateError {
     ValidatorDoesNotExist(usize),
 }
 
-//TODO(wen): Maybe we can merge all the below functions into CertificateMessage.
+/// A builder for creating a `CertificateMessage` by efficiently aggregating BLS signatures.
 #[derive(Clone)]
-pub struct VoteCertificate(CertificateMessage);
+pub struct VoteCertificateBuilder {
+    certificate: Certificate,
+    signature: SignatureProjective,
+    bitmap: BitVec<u8, Lsb0>,
+}
 
-impl From<CertificateMessage> for VoteCertificate {
-    fn from(certificate_message: CertificateMessage) -> Self {
-        Self(certificate_message)
+impl TryFrom<CertificateMessage> for VoteCertificateBuilder {
+    type Error = CertificateError;
+
+    fn try_from(message: CertificateMessage) -> Result<Self, Self::Error> {
+        let projective_signature = SignatureProjective::try_from(message.signature)?;
+        Ok(VoteCertificateBuilder {
+            certificate: message.certificate,
+            signature: projective_signature,
+            bitmap: message.bitmap,
+        })
     }
 }
 
-impl VoteCertificate {
+impl VoteCertificateBuilder {
     pub fn new(certificate_id: CertificateId) -> Self {
-        VoteCertificate(CertificateMessage {
+        Self {
             certificate: certificate_id.into(),
-            signature: Signature::default(),
+            signature: SignatureProjective::identity(),
             bitmap: BitVec::<u8, Lsb0>::repeat(false, MAXIMUM_VALIDATORS),
-        })
+        }
     }
 
-    pub fn aggregate<'a, 'b, T>(&mut self, messages: T)
-    where
-        T: Iterator<Item = &'a VoteMessage>,
-        Self: 'b,
-        'b: 'a,
-    {
-        let signature = &mut self.0.signature;
-        // TODO: signature aggregation can be done out-of-order;
-        // consider aggregating signatures separately in parallel
-        let mut current_signature = if signature == &Signature::default() {
-            SignatureProjective::identity()
-        } else {
-            SignatureProjective::try_from(*signature).expect("Invalid signature")
-        };
-
-        // aggregate the votes
-        let bitmap = &mut self.0.bitmap;
+    /// Aggregates a slice of `VoteMessage`s into the builder.
+    pub fn aggregate(&mut self, messages: &[VoteMessage]) -> Result<(), CertificateError> {
         for vote_message in messages {
-            // set bit-vector for the validator
-            //
-            // TODO: This only accounts for one type of vote. Update this after
-            // we have a base3 encoding implementation.
-            assert!(
-                bitmap.len() > vote_message.rank as usize,
-                "Vote rank {} exceeds bitmap length {}",
-                vote_message.rank,
-                bitmap.len()
-            );
-            assert!(
-                bitmap.get(vote_message.rank as usize).as_deref() != Some(&true),
-                "Conflicting vote check should make this unreachable {vote_message:?}"
-            );
-            bitmap.set(vote_message.rank as usize, true);
-            // aggregate the signature
-            current_signature
-                .aggregate_with([&vote_message.signature])
-                .expect(
-                    "Failed to aggregate signature: {vote_message.signature:?} into {current_signature:?}"
-                );
+            if self.bitmap.len() <= vote_message.rank as usize {
+                return Err(CertificateError::ValidatorDoesNotExist(
+                    vote_message.rank as usize,
+                ));
+            }
+            self.bitmap.set(vote_message.rank as usize, true);
         }
-        *signature = Signature::from(current_signature);
+        let signature_iter = messages.iter().map(|vote_message| &vote_message.signature);
+        Ok(self.signature.aggregate_with(signature_iter)?)
     }
 
-    pub fn certificate(self) -> CertificateMessage {
-        self.0
+    pub fn build(self) -> CertificateMessage {
+        CertificateMessage {
+            certificate: self.certificate,
+            signature: self.signature.into(),
+            bitmap: self.bitmap,
+        }
     }
 }
 

votor/src/certificate_pool/vote_pool.rs

@@ -1,5 +1,5 @@
 use {
-    crate::{certificate_pool::vote_certificate::VoteCertificate, Stake},
+    crate::{certificate_pool::vote_certificate_builder::VoteCertificateBuilder, Stake},
     solana_pubkey::Pubkey,
     solana_sdk::hash::Hash,
     solana_vote::alpenglow::bls_message::VoteMessage,
@@ -96,8 +96,10 @@ impl SimpleVotePool {
         true
     }
 
-    pub fn add_to_certificate(&self, output: &mut VoteCertificate) {
-        output.aggregate(self.vote_entry.transactions.iter())
+    pub fn add_to_certificate(&self, output: &mut VoteCertificateBuilder) {
+        output
+            .aggregate(&self.vote_entry.transactions)
+            .expect("Incoming vote message signatures are assumed to be valid")
     }
 }
 
@@ -175,9 +177,11 @@ impl DuplicateBlockVotePool {
             .map_or(0, |vote_entries| vote_entries.total_stake_by_key)
     }
 
-    pub fn add_to_certificate(&self, block_id: &Hash, output: &mut VoteCertificate) {
+    pub fn add_to_certificate(&self, block_id: &Hash, output: &mut VoteCertificateBuilder) {
         if let Some(vote_entries) = self.votes.get(block_id) {
-            output.aggregate(vote_entries.transactions.iter())
+            output
+                .aggregate(&vote_entries.transactions)
+                .expect("Incoming vote message signatures are assumed to be valid")
         }
     }