From 70dc4f3ae6f061deffe86b16cf476625b60d5804 Mon Sep 17 00:00:00 2001
From: gusto <bacv@users.noreply.github.com>
Date: Fri, 15 Sep 2023 09:10:46 +0300
Subject: [PATCH] Simulation branch with one committee fix (#402)

* Tree and branch overlay comparison tests

* Update leader super majority fn in branch overlay

* Cleanup tests
---
 .../src/overlay/branch_overlay.rs             |  5 +-
 consensus-engine/src/overlay/mod.rs           | 61 +++++++++++++++++++
 2 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/consensus-engine/src/overlay/branch_overlay.rs b/consensus-engine/src/overlay/branch_overlay.rs
index 69661c49..18eac14e 100644
--- a/consensus-engine/src/overlay/branch_overlay.rs
+++ b/consensus-engine/src/overlay/branch_overlay.rs
@@ -123,8 +123,9 @@ where
         (committee_size * 2 / 3) + 1
     }
 
-    fn leader_super_majority_threshold(&self, id: NodeId) -> usize {
-        self.super_majority_threshold(id)
+    fn leader_super_majority_threshold(&self, _id: NodeId) -> usize {
+        let committee_size = self.root_committee().len();
+        (committee_size * 2 / 3) + 1
     }
 
     fn update_leader_selection<F, E>(&self, f: F) -> Result<Self, E>
diff --git a/consensus-engine/src/overlay/mod.rs b/consensus-engine/src/overlay/mod.rs
index ab98205d..60473a8e 100644
--- a/consensus-engine/src/overlay/mod.rs
+++ b/consensus-engine/src/overlay/mod.rs
@@ -49,3 +49,64 @@ pub trait LeaderSelection: Clone {
 pub trait CommitteeMembership: Clone {
     fn reshape_committees(&self, nodes: &mut [NodeId]);
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::overlay::{FisherYatesShuffle, RoundRobin};
+
+    const ENTROPY: [u8; 32] = [0; 32];
+
+    fn overlay_fns_match(a: &impl Overlay, b: &impl Overlay, nodes: &[NodeId]) {
+        assert_eq!(a.root_committee(), b.root_committee());
+        assert_eq!(
+            a.is_member_of_child_committee(nodes[0], nodes[1]),
+            b.is_member_of_child_committee(nodes[0], nodes[1]),
+        );
+        assert_eq!(
+            a.is_member_of_root_committee(nodes[0]),
+            b.is_member_of_root_committee(nodes[0]),
+        );
+        assert_eq!(
+            a.is_member_of_leaf_committee(nodes[0]),
+            b.is_member_of_leaf_committee(nodes[0]),
+        );
+        assert_eq!(
+            a.is_child_of_root_committee(nodes[0]),
+            b.is_child_of_root_committee(nodes[0])
+        );
+        assert_eq!(a.parent_committee(nodes[0]), b.parent_committee(nodes[0]));
+        assert_eq!(a.child_committees(nodes[0]), b.child_committees(nodes[0]));
+        assert_eq!(a.leaf_committees(nodes[0]), b.leaf_committees(nodes[0]));
+        assert_eq!(a.node_committee(nodes[0]), b.node_committee(nodes[0]));
+        assert_eq!(
+            a.super_majority_threshold(nodes[0]),
+            b.super_majority_threshold(nodes[0])
+        );
+        assert_eq!(
+            a.leader_super_majority_threshold(nodes[0]),
+            b.leader_super_majority_threshold(nodes[0])
+        );
+    }
+
+    #[test]
+    fn compare_tree_branch_one_committee() {
+        let nodes: Vec<_> = (0..10).map(|i| NodeId::new([i as u8; 32])).collect();
+        let tree_overlay = TreeOverlay::new(TreeOverlaySettings {
+            nodes: nodes.clone(),
+            current_leader: nodes[0],
+            number_of_committees: 1,
+            leader: RoundRobin::new(),
+            committee_membership: FisherYatesShuffle::new(ENTROPY),
+        });
+        let branch_overlay = BranchOverlay::new(BranchOverlaySettings {
+            current_leader: nodes[0],
+            nodes: nodes.clone(),
+            branch_depth: 1,
+            leader: RoundRobin::new(),
+            committee_membership: FisherYatesShuffle::new(ENTROPY),
+        });
+
+        overlay_fns_match(&tree_overlay, &branch_overlay, &nodes);
+    }
+}