From 8429ba47bf29a6a9d4941acefe4587f96e494399 Mon Sep 17 00:00:00 2001
From: Asmaa Magdoub <asmaa@starkware.co>
Date: Thu, 5 Sep 2024 13:34:31 +0300
Subject: [PATCH] feat(consensus): add locked_value to sm

---
 .../papyrus_consensus/src/state_machine.rs    | 25 ++++++++++++---
 .../src/state_machine_test.rs                 | 31 +++++++++++++++++++
 .../sequencing/papyrus_consensus/src/types.rs |  1 +
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/crates/sequencing/papyrus_consensus/src/state_machine.rs b/crates/sequencing/papyrus_consensus/src/state_machine.rs
index 4189c30ed6..c6551c12d2 100644
--- a/crates/sequencing/papyrus_consensus/src/state_machine.rs
+++ b/crates/sequencing/papyrus_consensus/src/state_machine.rs
@@ -12,7 +12,7 @@ use std::collections::{HashMap, VecDeque};
 use starknet_api::block::BlockHash;
 use tracing::trace;
 
-use crate::types::{Round, ValidatorId};
+use crate::types::{ProposalContentId, Round, ValidatorId};
 
 /// Events which the state machine sends/receives.
 #[derive(Debug, Clone, PartialEq)]
@@ -64,6 +64,7 @@ pub struct StateMachine {
     // events in `events_queue`.
     awaiting_get_proposal: bool,
     events_queue: VecDeque<StateMachineEvent>,
+    locked_value: Option<(ProposalContentId, Round)>,
 }
 
 impl StateMachine {
@@ -79,6 +80,7 @@ impl StateMachine {
             precommits: HashMap::new(),
             awaiting_get_proposal: false,
             events_queue: VecDeque::new(),
+            locked_value: None,
         }
     }
 
@@ -374,6 +376,11 @@ impl StateMachine {
             panic!("Proposal does not match quorum.");
         }
 
+        self.locked_value = match self.locked_value {
+            Some((_, locked_round)) if locked_round >= round => self.locked_value,
+            _ => block_hash.map(|hash| (hash, round)),
+        };
+
         output.append(&mut self.send_precommit(*block_hash, round, leader_fn));
         output
     }
@@ -447,9 +454,19 @@ impl StateMachine {
         self.round = round;
         self.step = Step::Propose;
         if self.id == leader_fn(self.round) {
-            self.awaiting_get_proposal = true;
-            // TODO(matan): Support re-proposing validValue.
-            return VecDeque::from([StateMachineEvent::GetProposal(None, self.round)]);
+            match self.locked_value {
+                Some((proposal_content_id, valid_round)) => {
+                    return VecDeque::from([StateMachineEvent::Proposal(
+                        Some(proposal_content_id),
+                        self.round,
+                        Some(valid_round),
+                    )]);
+                }
+                None => {
+                    self.awaiting_get_proposal = true;
+                    return VecDeque::from([StateMachineEvent::GetProposal(None, self.round)]);
+                }
+            }
         }
         let Some(proposal) = self.proposals.get(&round) else {
             return VecDeque::from([StateMachineEvent::TimeoutPropose(round)]);
diff --git a/crates/sequencing/papyrus_consensus/src/state_machine_test.rs b/crates/sequencing/papyrus_consensus/src/state_machine_test.rs
index c97746bc31..58ce24c3a8 100644
--- a/crates/sequencing/papyrus_consensus/src/state_machine_test.rs
+++ b/crates/sequencing/papyrus_consensus/src/state_machine_test.rs
@@ -347,3 +347,34 @@ fn dont_handle_enqueued_while_awaiting_get_proposal() {
     assert_eq!(wrapper.next_event().unwrap(), StateMachineEvent::GetProposal(None, ROUND + 2));
     assert!(wrapper.next_event().is_none());
 }
+
+#[test]
+fn return_proposal_if_locked_value_is_set() {
+    let mut wrapper = TestWrapper::new(*PROPOSER_ID, 4, |_: Round| *PROPOSER_ID);
+
+    wrapper.start();
+    assert_eq!(wrapper.next_event().unwrap(), StateMachineEvent::GetProposal(None, ROUND));
+    assert!(wrapper.next_event().is_none());
+
+    wrapper.send_get_proposal(BLOCK_HASH, ROUND);
+    assert_eq!(wrapper.next_event().unwrap(), StateMachineEvent::Proposal(BLOCK_HASH, ROUND, None));
+    assert_eq!(wrapper.next_event().unwrap(), StateMachineEvent::Prevote(BLOCK_HASH, ROUND));
+    // locked_value is set after receiving a Prevote quorum.
+    wrapper.send_prevote(BLOCK_HASH, ROUND);
+    wrapper.send_prevote(BLOCK_HASH, ROUND);
+    assert_eq!(wrapper.next_event().unwrap(), StateMachineEvent::TimeoutPrevote(ROUND));
+    assert_eq!(wrapper.next_event().unwrap(), StateMachineEvent::Precommit(BLOCK_HASH, ROUND));
+
+    wrapper.send_precommit(None, ROUND);
+    wrapper.send_precommit(None, ROUND);
+    assert_eq!(wrapper.next_event().unwrap(), StateMachineEvent::TimeoutPrecommit(ROUND));
+
+    wrapper.send_timeout_precommit(ROUND);
+
+    // no need to GetProposal since we already have a locked value.
+    assert_eq!(
+        wrapper.next_event().unwrap(),
+        StateMachineEvent::Proposal(BLOCK_HASH, ROUND + 1, Some(ROUND))
+    );
+    assert_eq!(wrapper.next_event().unwrap(), StateMachineEvent::Prevote(BLOCK_HASH, ROUND + 1));
+}
diff --git a/crates/sequencing/papyrus_consensus/src/types.rs b/crates/sequencing/papyrus_consensus/src/types.rs
index dfbd36619b..35554f2ddb 100644
--- a/crates/sequencing/papyrus_consensus/src/types.rs
+++ b/crates/sequencing/papyrus_consensus/src/types.rs
@@ -14,6 +14,7 @@ use starknet_api::core::ContractAddress;
 // TODO(matan): Determine the actual type of NodeId.
 pub type ValidatorId = ContractAddress;
 pub type Round = u32;
+pub type ProposalContentId = BlockHash;
 
 /// Interface that any concrete block type must implement to be used by consensus.
 ///