From f3f7683a63e81c7c8a9d22914c29333e394a4876 Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Sun, 28 Apr 2024 03:36:18 -0400
Subject: [PATCH] Cache the proposal for a round

---
 coordinator/tributary/tendermint/src/lib.rs   | 26 +++++++++++--------
 .../tributary/tendermint/src/message_log.rs   |  9 -------
 2 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/coordinator/tributary/tendermint/src/lib.rs b/coordinator/tributary/tendermint/src/lib.rs
index 748720da5..145c51884 100644
--- a/coordinator/tributary/tendermint/src/lib.rs
+++ b/coordinator/tributary/tendermint/src/lib.rs
@@ -3,7 +3,7 @@ use core::fmt::Debug;
 use std::{
   sync::Arc,
   time::{SystemTime, Instant, Duration},
-  collections::VecDeque,
+  collections::{VecDeque, HashMap},
 };
 
 use parity_scale_codec::{Encode, Decode, IoReader};
@@ -263,6 +263,9 @@ pub struct TendermintMachine<N: Network> {
   synced_block_result_send: mpsc::UnboundedSender<bool>,
 
   block: BlockData<N>,
+  // TODO: Move this into the Block struct
+  round_proposals: HashMap<RoundNumber, (Option<RoundNumber>, N::Block)>,
+  // TODO: Move this into the Round struct
   upons: Upons,
 }
 
@@ -372,6 +375,9 @@ impl<N: Network + 'static> TendermintMachine<N> {
       proposal,
     );
 
+    // Reset the round proposals
+    self.round_proposals = HashMap::new();
+
     // Start the first round
     self.round(RoundNumber(0), Some(round_end));
   }
@@ -410,16 +416,7 @@ impl<N: Network + 'static> TendermintMachine<N> {
   }
 
   fn proposal_for_round(&self, round: RoundNumber) -> Option<(Option<RoundNumber>, &N::Block)> {
-    let proposer = self.weights.proposer(self.block.number, round);
-    if let Some(proposal_signed) = self.block.log.get(round, proposer, Step::Propose) {
-      if let Data::Proposal(vr, block) = &proposal_signed.msg.data {
-        Some((*vr, block))
-      } else {
-        panic!("message for Step::Propose didn't have Data::Proposal");
-      }
-    } else {
-      None?
-    }
+    self.round_proposals.get(&round).map(|(round, block)| (*round, block))
   }
 
   // L22-27
@@ -745,6 +742,11 @@ impl<N: Network + 'static> TendermintMachine<N> {
       msg.data.step(),
     );
 
+    // If this is a proposal, insert it
+    if let Data::Proposal(vr, block) = &msg.data {
+      self.round_proposals.insert(msg.round, (*vr, block.clone()));
+    }
+
     // L55-56
     // Jump ahead if we should
     if (msg.round.0 > self.block.round().number.0) &&
@@ -885,6 +887,8 @@ impl<N: Network + 'static> TendermintMachine<N> {
             Some(proposal),
           ),
 
+          round_proposals: HashMap::new(),
+
           upons: Upons {
             upon_prevotes: false,
             upon_successful_current_round_prevotes: false,
diff --git a/coordinator/tributary/tendermint/src/message_log.rs b/coordinator/tributary/tendermint/src/message_log.rs
index e65568cac..716f6d64e 100644
--- a/coordinator/tributary/tendermint/src/message_log.rs
+++ b/coordinator/tributary/tendermint/src/message_log.rs
@@ -85,13 +85,4 @@ impl<N: Network> MessageLog<N> {
     let (_, weight) = self.message_instances(round, data);
     weight >= self.weights.threshold()
   }
-
-  pub(crate) fn get(
-    &self,
-    round: RoundNumber,
-    sender: N::ValidatorId,
-    step: Step,
-  ) -> Option<&SignedMessageFor<N>> {
-    self.log.get(&round).and_then(|round| round.get(&sender).and_then(|msgs| msgs.get(&step)))
-  }
 }