From 30376d0f30f02aab1ff92cb9bb6033f956adf150 Mon Sep 17 00:00:00 2001
From: Brendon Fish <bfish713@gmail.com>
Date: Tue, 23 Jan 2024 11:15:55 -0500
Subject: [PATCH] index payloads by view

---
 crates/hotshot/src/lib.rs          |  7 +++----
 crates/task-impls/src/consensus.rs |  2 +-
 crates/task-impls/src/da.rs        |  2 +-
 crates/types/src/consensus.rs      | 13 +++----------
 4 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/crates/hotshot/src/lib.rs b/crates/hotshot/src/lib.rs
index 69a10e582a..9612df7a78 100644
--- a/crates/hotshot/src/lib.rs
+++ b/crates/hotshot/src/lib.rs
@@ -39,7 +39,7 @@ use hotshot_task::{
 use hotshot_task_impls::{events::HotShotEvent, network::NetworkTaskKind};
 
 use hotshot_types::{
-    consensus::{Consensus, ConsensusMetricsValue, PayloadStore, View, ViewInner, ViewQueue},
+    consensus::{Consensus, ConsensusMetricsValue, View, ViewInner, ViewQueue},
     data::Leaf,
     error::StorageSnafu,
     event::EventType,
@@ -217,9 +217,8 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>> SystemContext<TYPES, I> {
         );
 
         let mut saved_leaves = HashMap::new();
-        let mut saved_payloads = PayloadStore::default();
+        let mut saved_payloads = BTreeMap::new();
         saved_leaves.insert(anchored_leaf.commit(), anchored_leaf.clone());
-        let payload_commitment = anchored_leaf.get_payload_commitment();
         if let Some(payload) = anchored_leaf.get_block_payload() {
             let encoded_txns = match payload.encode() {
                 // TODO (Keyao) [VALIDATED_STATE] - Avoid collect/copy on the encoded transaction bytes.
@@ -229,7 +228,7 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>> SystemContext<TYPES, I> {
                     return Err(HotShotError::BlockError { source: e });
                 }
             };
-            saved_payloads.insert(payload_commitment, encoded_txns);
+            saved_payloads.insert(anchored_leaf.get_view_number(), encoded_txns);
         }
 
         let start_view = anchored_leaf.get_view_number();
diff --git a/crates/task-impls/src/consensus.rs b/crates/task-impls/src/consensus.rs
index e2848584c3..7acc4254cf 100644
--- a/crates/task-impls/src/consensus.rs
+++ b/crates/task-impls/src/consensus.rs
@@ -702,7 +702,7 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>, A: ConsensusApi<TYPES, I> +
                                 // If the block payload is available for this leaf, include it in
                                 // the leaf chain that we send to the client.
                                 if let Some(encoded_txns) =
-                                    consensus.saved_payloads.get(leaf.get_payload_commitment())
+                                    consensus.saved_payloads.get(&leaf.get_view_number())
                                 {
                                     let payload = BlockPayload::from_bytes(
                                         encoded_txns.clone().into_iter(),
diff --git a/crates/task-impls/src/da.rs b/crates/task-impls/src/da.rs
index a420a15946..bab588d1dd 100644
--- a/crates/task-impls/src/da.rs
+++ b/crates/task-impls/src/da.rs
@@ -192,7 +192,7 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>, A: ConsensusApi<TYPES, I> +
                 // Record the payload we have promised to make available.
                 consensus
                     .saved_payloads
-                    .insert(payload_commitment, proposal.data.encoded_transactions);
+                    .insert(view, proposal.data.encoded_transactions);
             }
             HotShotEvent::DAVoteRecv(ref vote) => {
                 debug!("DA vote recv, Main Task {:?}", vote.get_view_number());
diff --git a/crates/types/src/consensus.rs b/crates/types/src/consensus.rs
index e86939804e..d18622ae03 100644
--- a/crates/types/src/consensus.rs
+++ b/crates/types/src/consensus.rs
@@ -55,7 +55,7 @@ pub struct Consensus<TYPES: NodeType> {
     ///
     /// Contains the block payload commitment and encoded transactions for every leaf in
     /// `saved_leaves` if that payload is available.
-    pub saved_payloads: PayloadStore,
+    pub saved_payloads: BTreeMap<TYPES::Time, Vec<u8>>,
 
     /// The `locked_qc` view number
     pub locked_view: TYPES::Time,
@@ -316,21 +316,14 @@ impl<TYPES: NodeType> Consensus<TYPES> {
         // perform gc
         self.saved_da_certs
             .retain(|view_number, _| *view_number >= old_anchor_view);
-        self.state_map
-            .range(old_anchor_view..new_anchor_view)
-            .filter_map(|(_view_number, view)| view.get_payload_commitment())
-            .for_each(|payload_commitment| {
-                self.saved_payloads.remove(payload_commitment);
-            });
         self.state_map
             .range(old_anchor_view..new_anchor_view)
             .filter_map(|(_view_number, view)| view.get_leaf_commitment())
             .for_each(|leaf| {
-                if let Some(removed) = self.saved_leaves.remove(&leaf) {
-                    self.saved_payloads.remove(removed.get_payload_commitment());
-                }
+                self.saved_leaves.remove(&leaf);
             });
         self.state_map = self.state_map.split_off(&new_anchor_view);
+        self.saved_payloads = self.saved_payloads.split_off(&new_anchor_view);
     }
 
     /// Gets the last decided state