apollo_consensus: use Vote in Prevote/Precommit SMEvent

Author: asmaastarkware

crates/apollo_consensus/src/single_height_consensus.rs

@@ -240,38 +240,38 @@ impl SingleHeightConsensus {
             | StateMachineEvent::TimeoutPrecommit(_round) => {
                 self.handle_timeout(context, event).await
             }
-            StateMachineEvent::Prevote(_proposal_id, round) => {
+            StateMachineEvent::Prevote(vote) => {
                 let Some(last_vote) = &self.last_prevote else {
                     return Err(ConsensusError::InternalInconsistency(
                         "No prevote to send".to_string(),
                     ));
                 };
-                if last_vote.round > round {
+                if last_vote.round > vote.round {
                     // Only replay the newest prevote.
                     return Ok(ShcReturn::Tasks(Vec::new()));
                 }
                 trace_every_n_sec!(REBROADCAST_LOG_PERIOD_SECS, "Rebroadcasting {last_vote:?}");
                 context.broadcast(last_vote.clone()).await?;
                 Ok(ShcReturn::Tasks(vec![ShcTask::Prevote(
                     self.timeouts.get_prevote_timeout(0),
-                    event,
+                    StateMachineEvent::Prevote(last_vote.clone()),
                 )]))
             }
-            StateMachineEvent::Precommit(_proposal_id, round) => {
+            StateMachineEvent::Precommit(vote) => {
                 let Some(last_vote) = &self.last_precommit else {
                     return Err(ConsensusError::InternalInconsistency(
                         "No precommit to send".to_string(),
                     ));
                 };
-                if last_vote.round > round {
+                if last_vote.round > vote.round {
                     // Only replay the newest precommit.
                     return Ok(ShcReturn::Tasks(Vec::new()));
                 }
                 trace_every_n_sec!(REBROADCAST_LOG_PERIOD_SECS, "Rebroadcasting {last_vote:?}");
                 context.broadcast(last_vote.clone()).await?;
                 Ok(ShcReturn::Tasks(vec![ShcTask::Precommit(
                     self.timeouts.get_precommit_timeout(0),
-                    event,
+                    StateMachineEvent::Precommit(last_vote.clone()),
                 )]))
             }
             StateMachineEvent::Proposal(proposal_id, round, valid_round) => {
@@ -350,14 +350,10 @@ impl SingleHeightConsensus {
         }
 
         let (votes, sm_vote) = match vote.vote_type {
-            VoteType::Prevote => (
-                &mut self.prevotes,
-                StateMachineEvent::Prevote(vote.proposal_commitment, vote.round),
-            ),
-            VoteType::Precommit => (
-                &mut self.precommits,
-                StateMachineEvent::Precommit(vote.proposal_commitment, vote.round),
-            ),
+            VoteType::Prevote => (&mut self.prevotes, StateMachineEvent::Prevote(vote.clone())),
+            VoteType::Precommit => {
+                (&mut self.precommits, StateMachineEvent::Precommit(vote.clone()))
+            }
         };
 
         match votes.entry((vote.round, vote.voter)) {
@@ -407,27 +403,11 @@ impl SingleHeightConsensus {
                 StateMachineEvent::Decision(proposal_id, round) => {
                     return self.handle_state_machine_decision(proposal_id, round).await;
                 }
-                StateMachineEvent::Prevote(proposal_id, round) => {
-                    ret_val.extend(
-                        self.handle_state_machine_vote(
-                            context,
-                            proposal_id,
-                            round,
-                            VoteType::Prevote,
-                        )
-                        .await?,
-                    );
+                StateMachineEvent::Prevote(vote) => {
+                    ret_val.extend(self.handle_state_machine_vote(context, vote).await?);
                 }
-                StateMachineEvent::Precommit(proposal_id, round) => {
-                    ret_val.extend(
-                        self.handle_state_machine_vote(
-                            context,
-                            proposal_id,
-                            round,
-                            VoteType::Precommit,
-                        )
-                        .await?,
-                    );
+                StateMachineEvent::Precommit(vote) => {
+                    ret_val.extend(self.handle_state_machine_vote(context, vote).await?);
                 }
                 StateMachineEvent::TimeoutPropose(round) => {
                     ret_val.push(ShcTask::TimeoutPropose(
@@ -520,42 +500,34 @@ impl SingleHeightConsensus {
     async fn handle_state_machine_vote<ContextT: ConsensusContext>(
         &mut self,
         context: &mut ContextT,
-        proposal_id: Option<ProposalCommitment>,
-        round: Round,
-        vote_type: VoteType,
+        vote: Vote,
     ) -> Result<Vec<ShcTask>, ConsensusError> {
-        let prevote_timeout = self.timeouts.get_prevote_timeout(round);
-        let precommit_timeout = self.timeouts.get_precommit_timeout(round);
-        let (votes, last_vote, task) = match vote_type {
+        let prevote_timeout = self.timeouts.get_prevote_timeout(vote.round);
+        let precommit_timeout = self.timeouts.get_precommit_timeout(vote.round);
+        let (votes, last_vote, task) = match vote.vote_type {
             VoteType::Prevote => (
                 &mut self.prevotes,
                 &mut self.last_prevote,
-                ShcTask::Prevote(prevote_timeout, StateMachineEvent::Prevote(proposal_id, round)),
+                ShcTask::Prevote(prevote_timeout, StateMachineEvent::Prevote(vote.clone())),
             ),
             VoteType::Precommit => (
                 &mut self.precommits,
                 &mut self.last_precommit,
-                ShcTask::Precommit(
-                    precommit_timeout,
-                    StateMachineEvent::Precommit(proposal_id, round),
-                ),
+                ShcTask::Precommit(precommit_timeout, StateMachineEvent::Precommit(vote.clone())),
             ),
         };
-        let vote = Vote {
-            vote_type,
-            height: self.state_machine.height().0,
-            round,
-            proposal_commitment: proposal_id,
-            voter: self.state_machine.validator_id(),
-        };
-        if let Some(old) = votes.insert((round, self.state_machine.validator_id()), vote.clone()) {
+        // Ensure the voter matches this node.
+        assert_eq!(vote.voter, self.state_machine.validator_id());
+        if let Some(old) =
+            votes.insert((vote.round, self.state_machine.validator_id()), vote.clone())
+        {
             return Err(ConsensusError::InternalInconsistency(format!(
                 "State machine should not send repeat votes: old={old:?}, new={vote:?}"
             )));
         }
         *last_vote = match last_vote {
             None => Some(vote.clone()),
-            Some(last_vote) if round > last_vote.round => Some(vote.clone()),
+            Some(last_vote) if vote.round > last_vote.round => Some(vote.clone()),
             Some(_) => {
                 // According to the Tendermint paper, the state machine should only vote for its
                 // current round. It should monotonically increase its round. It should only vote

crates/apollo_consensus/src/single_height_consensus_test.rs

@@ -42,19 +42,16 @@ fn get_proposal_init_for_height(height: BlockNumber) -> ProposalInit {
     ProposalInit { height, ..*PROPOSAL_INIT }
 }
 
-fn prevote_task(block_felt: Option<Felt>, round: u32) -> ShcTask {
+fn prevote_task(block_felt: Option<Felt>, round: u32, voter: ValidatorId) -> ShcTask {
     let duration = TIMEOUTS.get_prevote_timeout(round);
-    ShcTask::Prevote(
-        duration,
-        StateMachineEvent::Prevote(block_felt.map(ProposalCommitment), round),
-    )
+    ShcTask::Prevote(duration, StateMachineEvent::Prevote(prevote(block_felt, 0, round, voter)))
 }
 
-fn precommit_task(block_felt: Option<Felt>, round: u32) -> ShcTask {
+fn precommit_task(block_felt: Option<Felt>, round: u32, voter: ValidatorId) -> ShcTask {
     let duration = TIMEOUTS.get_precommit_timeout(round);
     ShcTask::Precommit(
         duration,
-        StateMachineEvent::Precommit(block_felt.map(ProposalCommitment), round),
+        StateMachineEvent::Precommit(precommit(block_felt, 0, round, voter)),
     )
 }
 
@@ -115,7 +112,7 @@ async fn proposer() {
     assert_eq!(*shc_ret.as_tasks().unwrap()[0].as_build_proposal().unwrap().0, 0);
     assert_eq!(
         shc.handle_event(&mut context, StateMachineEvent::GetProposal(Some(BLOCK.id), 0)).await,
-        Ok(ShcReturn::Tasks(vec![prevote_task(Some(BLOCK.id.0), 0)]))
+        Ok(ShcReturn::Tasks(vec![prevote_task(Some(BLOCK.id.0), 0, *PROPOSER_ID)]))
     );
     assert_eq!(
         shc.handle_vote(&mut context, prevote(Some(BLOCK.id.0), 0, 0, *VALIDATOR_ID_1)).await,
@@ -130,7 +127,10 @@ async fn proposer() {
     // The Node got a Prevote quorum.
     assert_eq!(
         shc.handle_vote(&mut context, prevote(Some(BLOCK.id.0), 0, 0, *VALIDATOR_ID_2)).await,
-        Ok(ShcReturn::Tasks(vec![timeout_prevote_task(0), precommit_task(Some(BLOCK.id.0), 0),]))
+        Ok(ShcReturn::Tasks(vec![
+            timeout_prevote_task(0),
+            precommit_task(Some(BLOCK.id.0), 0, *PROPOSER_ID)
+        ]))
     );
 
     let precommits = vec![
@@ -194,7 +194,7 @@ async fn validator(repeat_proposal: bool) {
     );
     assert_eq!(
         shc.handle_event(&mut context, VALIDATE_PROPOSAL_EVENT.clone()).await,
-        Ok(ShcReturn::Tasks(vec![prevote_task(Some(BLOCK.id.0), 0)]))
+        Ok(ShcReturn::Tasks(vec![prevote_task(Some(BLOCK.id.0), 0, *VALIDATOR_ID_1)]))
     );
     if repeat_proposal {
         // Send the same proposal again, which should be ignored (no expectations).
@@ -215,7 +215,10 @@ async fn validator(repeat_proposal: bool) {
     assert_eq!(
         shc.handle_vote(&mut context, prevote(Some(BLOCK.id.0), HEIGHT.0, 0, *VALIDATOR_ID_2))
             .await,
-        Ok(ShcReturn::Tasks(vec![timeout_prevote_task(0), precommit_task(Some(BLOCK.id.0), 0)]))
+        Ok(ShcReturn::Tasks(vec![
+            timeout_prevote_task(0),
+            precommit_task(Some(BLOCK.id.0), 0, *VALIDATOR_ID_1)
+        ]))
     );
 
     let precommits = vec![
@@ -272,7 +275,7 @@ async fn vote_twice(same_vote: bool) {
     );
     assert_eq!(
         shc.handle_event(&mut context, VALIDATE_PROPOSAL_EVENT.clone()).await,
-        Ok(ShcReturn::Tasks(vec![prevote_task(Some(BLOCK.id.0), 0)]))
+        Ok(ShcReturn::Tasks(vec![prevote_task(Some(BLOCK.id.0), 0, *VALIDATOR_ID_1)]))
     );
 
     let res =
@@ -290,7 +293,10 @@ async fn vote_twice(same_vote: bool) {
     // The Node got a Prevote quorum.
     assert_eq!(
         res,
-        Ok(ShcReturn::Tasks(vec![timeout_prevote_task(0), precommit_task(Some(BLOCK.id.0), 0),]))
+        Ok(ShcReturn::Tasks(vec![
+            timeout_prevote_task(0),
+            precommit_task(Some(BLOCK.id.0), 0, *VALIDATOR_ID_1),
+        ]))
     );
 
     let first_vote = precommit(Some(BLOCK.id.0), HEIGHT.0, 0, *PROPOSER_ID);
@@ -345,7 +351,7 @@ async fn rebroadcast_votes() {
     assert_eq!(*shc_ret.as_tasks().unwrap()[0].as_build_proposal().unwrap().0, 0);
     assert_eq!(
         shc.handle_event(&mut context, StateMachineEvent::GetProposal(Some(BLOCK.id), 0)).await,
-        Ok(ShcReturn::Tasks(vec![prevote_task(Some(BLOCK.id.0), 0)]))
+        Ok(ShcReturn::Tasks(vec![prevote_task(Some(BLOCK.id.0), 0, *PROPOSER_ID)]))
     );
     assert_eq!(
         shc.handle_vote(&mut context, prevote(Some(BLOCK.id.0), 0, 0, *VALIDATOR_ID_1)).await,
@@ -362,12 +368,19 @@ async fn rebroadcast_votes() {
     // The Node got a Prevote quorum.
     assert_eq!(
         shc.handle_vote(&mut context, prevote(Some(BLOCK.id.0), 0, 0, *VALIDATOR_ID_2)).await,
-        Ok(ShcReturn::Tasks(vec![timeout_prevote_task(0), precommit_task(Some(BLOCK.id.0), 0),]))
+        Ok(ShcReturn::Tasks(vec![
+            timeout_prevote_task(0),
+            precommit_task(Some(BLOCK.id.0), 0, *PROPOSER_ID),
+        ]))
     );
     // Re-broadcast vote.
     assert_eq!(
-        shc.handle_event(&mut context, StateMachineEvent::Precommit(Some(BLOCK.id), 0)).await,
-        Ok(ShcReturn::Tasks(vec![precommit_task(Some(BLOCK.id.0), 0),]))
+        shc.handle_event(
+            &mut context,
+            StateMachineEvent::Precommit(precommit(Some(BLOCK.id.0), 0, 0, *PROPOSER_ID))
+        )
+        .await,
+        Ok(ShcReturn::Tasks(vec![precommit_task(Some(BLOCK.id.0), 0, *PROPOSER_ID),]))
     );
 }
 
@@ -410,7 +423,10 @@ async fn repropose() {
     // The Node got a Prevote quorum, and set valid proposal.
     assert_eq!(
         shc.handle_vote(&mut context, prevote(Some(BLOCK.id.0), 0, 0, *VALIDATOR_ID_2)).await,
-        Ok(ShcReturn::Tasks(vec![timeout_prevote_task(0), precommit_task(Some(BLOCK.id.0), 0),]))
+        Ok(ShcReturn::Tasks(vec![
+            timeout_prevote_task(0),
+            precommit_task(Some(BLOCK.id.0), 0, *PROPOSER_ID),
+        ]))
     );
     // Advance to the next round.
     let precommits = vec![