[dag][bugfix] fix bitmask for incomplete first round

aptos-labs · Aug 24, 2023 · 2bcdbd4 · 2bcdbd4
1 parent f2e9918
commit 2bcdbd4
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 12 deletions.
diff --git 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
@@ -209,15 +209,14 @@ 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
             .get_strong_links_for_round(round, &epoch_state.verifier)
             .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![]));