From 44367a15a592537c53379e724fd85282440f1dfa Mon Sep 17 00:00:00 2001
From: Daniel Sanchez <sanchez.quiros.daniel@gmail.com>
Date: Thu, 10 Aug 2023 13:18:17 +0200
Subject: [PATCH] Sim update overlay (#298)

* Update sim engine on timeout qc

* Remove rebuild method
---
 consensus-engine/src/lib.rs                   |  1 -
 consensus-engine/src/overlay/flat_overlay.rs  |  4 --
 consensus-engine/src/overlay/mod.rs           |  1 -
 .../src/overlay/tree_overlay/overlay.rs       |  4 --
 simulations/src/node/carnot/mod.rs            | 52 +++++++++++++------
 5 files changed, 37 insertions(+), 25 deletions(-)

diff --git a/consensus-engine/src/lib.rs b/consensus-engine/src/lib.rs
index 0b5574a9..abf35cfa 100644
--- a/consensus-engine/src/lib.rs
+++ b/consensus-engine/src/lib.rs
@@ -105,7 +105,6 @@ impl<O: Overlay> Carnot<O> {
         new_state.update_timeout_qc(timeout_qc.clone());
 
         new_state.current_view = timeout_qc.view().next();
-        new_state.overlay.rebuild(timeout_qc);
 
         new_state
     }
diff --git a/consensus-engine/src/overlay/flat_overlay.rs b/consensus-engine/src/overlay/flat_overlay.rs
index 10c53556..0f8d608a 100644
--- a/consensus-engine/src/overlay/flat_overlay.rs
+++ b/consensus-engine/src/overlay/flat_overlay.rs
@@ -50,10 +50,6 @@ where
         self.nodes.clone().into_iter().collect()
     }
 
-    fn rebuild(&mut self, _timeout_qc: crate::TimeoutQc) {
-        // do nothing for now
-    }
-
     fn is_member_of_child_committee(&self, _parent: NodeId, _child: NodeId) -> bool {
         false
     }
diff --git a/consensus-engine/src/overlay/mod.rs b/consensus-engine/src/overlay/mod.rs
index 55ea22f3..9033f430 100644
--- a/consensus-engine/src/overlay/mod.rs
+++ b/consensus-engine/src/overlay/mod.rs
@@ -21,7 +21,6 @@ pub trait Overlay: Clone {
 
     fn new(settings: Self::Settings) -> Self;
     fn root_committee(&self) -> Committee;
-    fn rebuild(&mut self, timeout_qc: TimeoutQc);
     fn is_member_of_child_committee(&self, parent: NodeId, child: NodeId) -> bool;
     fn is_member_of_root_committee(&self, id: NodeId) -> bool;
     fn is_member_of_leaf_committee(&self, id: NodeId) -> bool;
diff --git a/consensus-engine/src/overlay/tree_overlay/overlay.rs b/consensus-engine/src/overlay/tree_overlay/overlay.rs
index cb90fdc0..b3fd3625 100644
--- a/consensus-engine/src/overlay/tree_overlay/overlay.rs
+++ b/consensus-engine/src/overlay/tree_overlay/overlay.rs
@@ -57,10 +57,6 @@ where
         self.carnot_tree.root_committee().clone()
     }
 
-    fn rebuild(&mut self, _timeout_qc: crate::TimeoutQc) {
-        // do nothing for now
-    }
-
     fn is_member_of_child_committee(&self, parent: NodeId, child: NodeId) -> bool {
         let child_parent = self.parent_committee(child);
         let parent = self.carnot_tree.committee_by_member_id(&parent);
diff --git a/simulations/src/node/carnot/mod.rs b/simulations/src/node/carnot/mod.rs
index 61b64d36..3bd7032b 100644
--- a/simulations/src/node/carnot/mod.rs
+++ b/simulations/src/node/carnot/mod.rs
@@ -344,21 +344,9 @@ impl<
                     "receive block proposal",
                 );
                 match self.engine.receive_block(block.header().clone()) {
-                    Ok(new) => {
+                    Ok(mut new) => {
                         if self.engine.current_view() != new.current_view() {
-                            // TODO: Refactor this into a method, use for timeout qc as well
-                            // new = new
-                            //     .update_overlay(|overlay| {
-                            //         let overlay = overlay
-                            //             .update_leader_selection(|leader_selection| {
-                            //                 leader_selection.on_new_block_received(&block)
-                            //             })
-                            //             .expect("Leader selection update should succeed");
-                            //         overlay.update_committees(|committee_membership| {
-                            //             committee_membership.on_new_block_received(&block)
-                            //         })
-                            //     })
-                            //     .unwrap_or(new);
+                            new = Self::update_overlay_with_block(new, &block);
                             self.engine = new;
                         }
                     }
@@ -455,7 +443,8 @@ impl<
                     timeout_view = %timeout_qc.view(),
                     "receive timeout qc message"
                 );
-                self.engine = self.engine.receive_timeout_qc(timeout_qc);
+                let new = self.engine.receive_timeout_qc(timeout_qc.clone());
+                self.engine = Self::update_overlay_with_timeout_qc(new, &timeout_qc);
             }
             Event::RootTimeout { timeouts } => {
                 tracing::debug!("root timeout {:?}", timeouts);
@@ -494,6 +483,39 @@ impl<
             self.handle_output(event);
         }
     }
+
+    fn update_overlay_with_block<Tx: Clone + Eq + Hash>(
+        state: Carnot<O>,
+        block: &nomos_core::block::Block<Tx>,
+    ) -> Carnot<O> {
+        state
+            .update_overlay(|overlay| {
+                overlay
+                    .update_leader_selection(|leader_selection| {
+                        leader_selection.on_new_block_received(block)
+                    })
+                    .expect("Leader selection update should succeed")
+                    .update_committees(|committee_membership| {
+                        committee_membership.on_new_block_received(block)
+                    })
+            })
+            .unwrap_or(state)
+    }
+
+    fn update_overlay_with_timeout_qc(state: Carnot<O>, qc: &TimeoutQc) -> Carnot<O> {
+        state
+            .update_overlay(|overlay| {
+                overlay
+                    .update_leader_selection(|leader_selection| {
+                        leader_selection.on_timeout_qc_received(qc)
+                    })
+                    .expect("Leader selection update should succeed")
+                    .update_committees(|committee_membership| {
+                        committee_membership.on_timeout_qc_received(qc)
+                    })
+            })
+            .unwrap_or(state)
+    }
 }
 
 impl<