From bc4d9ee513b62632f1eff27be89371d652118531 Mon Sep 17 00:00:00 2001
From: Balaji Arun <balajia@vt.edu>
Date: Tue, 22 Aug 2023 16:49:14 -0700
Subject: [PATCH] [dag][bugfix] fix bitmask for incomplete first round

---
 consensus/src/dag/dag_store.rs      | 18 +++++++++---------
 consensus/src/dag/tests/dag_test.rs |  5 ++---
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/consensus/src/dag/dag_store.rs b/consensus/src/dag/dag_store.rs
index 52612995c8128..67fb60ba95228 100644
--- a/consensus/src/dag/dag_store.rs
+++ b/consensus/src/dag/dag_store.rs
@@ -271,22 +271,22 @@ impl Dag {
         }
     }
 
-    pub fn lowest_incomplete_round(&self) -> Option<Round> {
+    pub fn lowest_incomplete_round(&self) -> Round {
+        if self.nodes_by_round.is_empty() {
+            return self.lowest_round();
+        }
+
         for (round, round_nodes) in &self.nodes_by_round {
             if round_nodes.iter().any(|node| node.is_none()) {
-                return Some(*round);
+                return *round;
             }
         }
-        None
+
+        self.highest_round() + 1
     }
 
     pub fn bitmask(&self, target_round: Round) -> DagSnapshotBitmask {
-        let lowest_round = match self.lowest_incomplete_round() {
-            Some(round) => round,
-            None => {
-                return DagSnapshotBitmask::new(self.highest_round() + 1, vec![]);
-            },
-        };
+        let lowest_round = self.lowest_incomplete_round();
 
         let bitmask = self
             .nodes_by_round
diff --git a/consensus/src/dag/tests/dag_test.rs b/consensus/src/dag/tests/dag_test.rs
index 185fcc05f639d..098992fa9a08f 100644
--- a/consensus/src/dag/tests/dag_test.rs
+++ b/consensus/src/dag/tests/dag_test.rs
@@ -209,7 +209,7 @@ fn test_dag_recover_from_storage() {
 fn test_dag_bitmask() {
     let (signers, epoch_state, mut dag, _) = setup();
 
-    let mut metadatas = vec![];
+    assert_eq!(dag.bitmask(15), DagSnapshotBitmask::new(1, vec![]));
 
     for round in 1..5 {
         let parents = dag
@@ -217,7 +217,6 @@ fn test_dag_bitmask() {
             .unwrap_or_default();
         for signer in &signers[0..3] {
             let node = new_certified_node(round, signer.author(), parents.clone());
-            metadatas.push(node.metadata().clone());
             assert!(dag.add_node(node).is_ok());
         }
     }
@@ -226,12 +225,12 @@ fn test_dag_bitmask() {
         DagSnapshotBitmask::new(1, vec![vec![true, true, true, false]; 4])
     );
 
+    // Populate the fourth author for all rounds
     for round in 1..5 {
         let parents = dag
             .get_strong_links_for_round(round, &epoch_state.verifier)
             .unwrap_or_default();
         let node = new_certified_node(round, signers[3].author(), parents.clone());
-        metadatas.push(node.metadata().clone());
         assert!(dag.add_node(node).is_ok());
     }
     assert_eq!(dag.bitmask(15), DagSnapshotBitmask::new(5, vec![]));