Add update committees structure for overlay (#295)

* Move roundrobin to leadership module

* Use references in leader selection

* Add membership traits to overlay and create membership module

* Implement committee membership for random beacon state

* Update flat overlay

* Create updateable membership traits and impls

* Update tree overlay

* Update overlay on consensus service

* Update overlay on simulations nomos node

* Update types on tests and modules

* Use chacha for shuffling

* Change to mut slice instead of inner cloning

* Use fisher yates shuffle from scratch

* Stylish and clippy happy

consensus-engine/src/lib.rs

@@ -411,11 +411,11 @@ impl<O: Overlay> Carnot<O> {
 mod test {
     use std::convert::Infallible;
 
-    use crate::overlay::{FlatOverlay, FlatOverlaySettings, RoundRobin};
+    use crate::overlay::{FlatOverlay, FlatOverlaySettings, FreezeMembership, RoundRobin};
 
     use super::*;
 
-    fn init(nodes: Vec<NodeId>) -> Carnot<FlatOverlay<RoundRobin>> {
+    fn init(nodes: Vec<NodeId>) -> Carnot<FlatOverlay<RoundRobin, FreezeMembership>> {
         assert!(!nodes.is_empty());
 
         Carnot::from_genesis(
@@ -436,7 +436,10 @@ mod test {
         )
     }
 
-    fn next_block(engine: &Carnot<FlatOverlay<RoundRobin>>, block: &Block) -> Block {
+    fn next_block(
+        engine: &Carnot<FlatOverlay<RoundRobin, FreezeMembership>>,
+        block: &Block,
+    ) -> Block {
         let mut next_id = block.id;
         next_id.0[0] += 1;
 
@@ -454,8 +457,8 @@ mod test {
     }
 
     fn update_leader_selection(
-        engine: &Carnot<FlatOverlay<RoundRobin>>,
-    ) -> Carnot<FlatOverlay<RoundRobin>> {
+        engine: &Carnot<FlatOverlay<RoundRobin, FreezeMembership>>,
+    ) -> Carnot<FlatOverlay<RoundRobin, FreezeMembership>> {
         engine
             .update_overlay(|overlay| {
                 overlay.update_leader_selection(

consensus-engine/src/overlay/flat_overlay.rs

@@ -1,24 +1,30 @@
 use super::LeaderSelection;
+use crate::overlay::CommitteeMembership;
 use crate::{NodeId, Overlay};
 use fraction::{Fraction, ToPrimitive};
 use serde::{Deserialize, Serialize};
+use std::marker::PhantomData;
+
 const LEADER_SUPER_MAJORITY_THRESHOLD_NUM: u64 = 2;
 const LEADER_SUPER_MAJORITY_THRESHOLD_DEN: u64 = 3;
 
 #[derive(Clone, Debug, PartialEq)]
 /// Flat overlay with a single committee and round robin leader selection.
-pub struct FlatOverlay<L: LeaderSelection> {
+pub struct FlatOverlay<L: LeaderSelection, M: CommitteeMembership> {
     nodes: Vec<NodeId>,
     leader: L,
     leader_threshold: Fraction,
+    _committee_membership: PhantomData<M>,
 }
 
-impl<L> Overlay for FlatOverlay<L>
+impl<L, M> Overlay for FlatOverlay<L, M>
 where
     L: LeaderSelection + Send + Sync + 'static,
+    M: CommitteeMembership + Send + Sync + 'static,
 {
     type Settings = FlatOverlaySettings<L>;
     type LeaderSelection = L;
+    type CommitteeMembership = M;
 
     fn new(
         FlatOverlaySettings {
@@ -36,6 +42,7 @@ where
                     LEADER_SUPER_MAJORITY_THRESHOLD_DEN,
                 )
             }),
+            _committee_membership: Default::default(),
         }
     }
 
@@ -107,29 +114,12 @@ where
             Err(e) => Err(e),
         }
     }
-}
 
-#[derive(Clone, Debug, Default, PartialEq)]
-#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-pub struct RoundRobin {
-    cur: usize,
-}
-
-impl RoundRobin {
-    pub fn new() -> Self {
-        Self { cur: 0 }
-    }
-
-    pub fn advance(&self) -> Self {
-        Self {
-            cur: (self.cur + 1),
-        }
-    }
-}
-
-impl LeaderSelection for RoundRobin {
-    fn next_leader(&self, nodes: &[NodeId]) -> NodeId {
-        nodes[self.cur % nodes.len()]
+    fn update_committees<F, E>(&self, _f: F) -> Result<Self, E>
+    where
+        F: FnOnce(Self::CommitteeMembership) -> Result<Self::CommitteeMembership, E>,
+    {
+        Ok(self.clone())
     }
 }
 

consensus-engine/src/overlay/leadership.rs

@@ -0,0 +1,32 @@
+// std
+
+// crates
+use serde::{Deserialize, Serialize};
+
+// internal
+use crate::overlay::LeaderSelection;
+use crate::NodeId;
+
+#[derive(Clone, Debug, Default, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct RoundRobin {
+    cur: usize,
+}
+
+impl RoundRobin {
+    pub fn new() -> Self {
+        Self { cur: 0 }
+    }
+
+    pub fn advance(&self) -> Self {
+        Self {
+            cur: (self.cur + 1),
+        }
+    }
+}
+
+impl LeaderSelection for RoundRobin {
+    fn next_leader(&self, nodes: &[NodeId]) -> NodeId {
+        nodes[self.cur % nodes.len()]
+    }
+}

consensus-engine/src/overlay/membership.rs

@@ -0,0 +1,43 @@
+// std
+
+// crates
+use rand::{Rng, SeedableRng};
+use rand_chacha::ChaCha20Rng;
+
+// internal
+use crate::overlay::CommitteeMembership;
+use crate::NodeId;
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub struct FisherYatesShuffle {
+    entropy: [u8; 32],
+}
+
+impl FisherYatesShuffle {
+    pub fn new(entropy: [u8; 32]) -> Self {
+        Self { entropy }
+    }
+
+    pub fn shuffle<T: Clone>(elements: &mut [T], entropy: [u8; 32]) {
+        let mut rng = ChaCha20Rng::from_seed(entropy);
+        // Implementation of fisher yates shuffling
+        // https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle
+        for i in (1..elements.len()).rev() {
+            let j = rng.gen_range(0..=i);
+            elements.swap(i, j);
+        }
+    }
+}
+
+impl CommitteeMembership for FisherYatesShuffle {
+    fn reshape_committees(&self, nodes: &mut [NodeId]) {
+        FisherYatesShuffle::shuffle(nodes, self.entropy);
+    }
+}
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub struct FreezeMembership;
+
+impl CommitteeMembership for FreezeMembership {
+    fn reshape_committees(&self, _nodes: &mut [NodeId]) {}
+}

consensus-engine/src/overlay/mod.rs

@@ -1,10 +1,14 @@
 use super::types::*;
 
 mod flat_overlay;
+mod leadership;
+mod membership;
 mod random_beacon;
 mod tree_overlay;
 
 pub use flat_overlay::*;
+pub use leadership::*;
+pub use membership::*;
 pub use random_beacon::*;
 pub use tree_overlay::*;
 
@@ -13,6 +17,7 @@ use std::marker::Send;
 pub trait Overlay: Clone {
     type Settings: Clone + Send + Sync + 'static;
     type LeaderSelection: LeaderSelection + Clone + Send + Sync + 'static;
+    type CommitteeMembership: CommitteeMembership + Clone + Send + Sync + 'static;
 
     fn new(settings: Self::Settings) -> Self;
     fn root_committee(&self) -> Committee;
@@ -31,8 +36,15 @@ pub trait Overlay: Clone {
     fn update_leader_selection<F, E>(&self, f: F) -> Result<Self, E>
     where
         F: FnOnce(Self::LeaderSelection) -> Result<Self::LeaderSelection, E>;
+    fn update_committees<F, E>(&self, f: F) -> Result<Self, E>
+    where
+        F: FnOnce(Self::CommitteeMembership) -> Result<Self::CommitteeMembership, E>;
 }
 
 pub trait LeaderSelection: Clone {
     fn next_leader(&self, nodes: &[NodeId]) -> NodeId;
 }
+
+pub trait CommitteeMembership: Clone {
+    fn reshape_committees(&self, nodes: &mut [NodeId]);
+}

consensus-engine/src/overlay/random_beacon.rs

@@ -1,3 +1,4 @@
+use crate::overlay::{CommitteeMembership, FisherYatesShuffle};
 use crate::types::*;
 use bls_signatures::{PrivateKey, PublicKey, Serialize, Signature};
 use rand::{seq::SliceRandom, SeedableRng};
@@ -60,6 +61,15 @@ impl RandomBeaconState {
         Self::Sad { entropy }
     }
 
+    pub fn initial_sad_from_entropy(entropy: [u8; 32]) -> Self {
+        Self::generate_sad(
+            View::new(-1),
+            &Self::Sad {
+                entropy: Box::new(entropy),
+            },
+        )
+    }
+
     pub fn check_advance_happy(&self, rb: RandomBeaconState, view: View) -> Result<Self, Error> {
         let context = view_to_bytes(view);
         match rb {
@@ -97,6 +107,14 @@ impl LeaderSelection for RandomBeaconState {
     }
 }
 
+impl CommitteeMembership for RandomBeaconState {
+    fn reshape_committees(&self, nodes: &mut [NodeId]) {
+        let mut seed = [0; 32];
+        seed.copy_from_slice(&self.entropy().deref()[..32]);
+        FisherYatesShuffle::shuffle(nodes, seed);
+    }
+}
+
 mod serialize_bls {
     use super::*;
     use serde::{Deserializer, Serializer};

consensus-engine/src/overlay/tree_overlay/overlay.rs

@@ -1,57 +1,55 @@
 use super::tree::Tree;
+use crate::overlay::CommitteeMembership;
 use crate::{overlay::LeaderSelection, Committee, NodeId, Overlay};
-use rand::rngs::StdRng;
-use rand::seq::SliceRandom;
-use rand::SeedableRng;
 
 #[derive(Debug, Clone)]
-pub struct TreeOverlaySettings<L: LeaderSelection> {
+pub struct TreeOverlaySettings<L: LeaderSelection, M: CommitteeMembership> {
     pub nodes: Vec<NodeId>,
     pub current_leader: NodeId,
-    pub entropy: [u8; 32],
     pub number_of_committees: usize,
     pub leader: L,
+    pub committee_membership: M,
 }
 
 #[derive(Debug, Clone)]
-pub struct TreeOverlay<L> {
-    pub(super) entropy: [u8; 32],
+pub struct TreeOverlay<L, M> {
     pub(super) number_of_committees: usize,
     pub(super) nodes: Vec<NodeId>,
     pub(super) current_leader: NodeId,
     pub(super) carnot_tree: Tree,
     pub(super) leader: L,
+    pub(super) committee_membership: M,
 }
 
-impl<L> Overlay for TreeOverlay<L>
+impl<L, M> Overlay for TreeOverlay<L, M>
 where
     L: LeaderSelection + Send + Sync + 'static,
+    M: CommitteeMembership + Send + Sync + 'static,
 {
-    type Settings = TreeOverlaySettings<L>;
+    type Settings = TreeOverlaySettings<L, M>;
 
     type LeaderSelection = L;
+    type CommitteeMembership = M;
 
     fn new(settings: Self::Settings) -> Self {
         let TreeOverlaySettings {
             mut nodes,
             current_leader,
-            entropy,
             number_of_committees,
             leader,
+            committee_membership,
         } = settings;
-        let mut rng = StdRng::from_seed(entropy);
-        // TODO: support custom shuffling algorithm
-        nodes.shuffle(&mut rng);
 
+        committee_membership.reshape_committees(&mut nodes);
         let carnot_tree = Tree::new(&nodes, number_of_committees);
 
         Self {
-            entropy,
             number_of_committees,
             nodes,
             current_leader,
             carnot_tree,
             leader,
+            committee_membership,
         }
     }
 
@@ -131,8 +129,7 @@ where
     }
 
     fn next_leader(&self) -> NodeId {
-        let mut rng = StdRng::from_seed(self.entropy);
-        *self.nodes.choose(&mut rng).unwrap()
+        self.leader.next_leader(&self.nodes)
     }
 
     fn super_majority_threshold(&self, id: NodeId) -> usize {
@@ -179,19 +176,36 @@ where
             Err(e) => Err(e),
         }
     }
+
+    fn update_committees<F, E>(&self, f: F) -> Result<Self, E>
+    where
+        F: FnOnce(Self::CommitteeMembership) -> Result<Self::CommitteeMembership, E>,
+    {
+        f(self.committee_membership.clone()).map(|committee_membership| {
+            let settings = TreeOverlaySettings {
+                nodes: self.nodes.clone(),
+                current_leader: self.current_leader,
+                number_of_committees: self.number_of_committees,
+                leader: self.leader.clone(),
+                committee_membership,
+            };
+            Self::new(settings)
+        })
+    }
 }
 
-impl<L> TreeOverlay<L>
+impl<L, M> TreeOverlay<L, M>
 where
     L: LeaderSelection + Send + Sync + 'static,
+    M: CommitteeMembership + Send + Sync + 'static,
 {
-    pub fn advance(&self, entropy: [u8; 32], leader: L) -> Self {
+    pub fn advance(&self, leader: L, committee_membership: M) -> Self {
         Self::new(TreeOverlaySettings {
             nodes: self.nodes.clone(),
             current_leader: self.next_leader(),
-            entropy,
             number_of_committees: self.number_of_committees,
             leader,
+            committee_membership,
         })
     }
 
@@ -206,10 +220,12 @@ where
 
 #[cfg(test)]
 mod tests {
-    use crate::overlay::RoundRobin;
+    use crate::overlay::leadership::RoundRobin;
+    use crate::overlay::membership::FisherYatesShuffle;
     use crate::Overlay;
 
     use super::*;
+    const ENTROPY: [u8; 32] = [0; 32];
 
     #[test]
     fn test_carnot_overlay_leader() {
@@ -217,9 +233,9 @@ mod tests {
         let overlay = TreeOverlay::new(TreeOverlaySettings {
             nodes: nodes.clone(),
             current_leader: nodes[0],
-            entropy: [0; 32],
             number_of_committees: 3,
             leader: RoundRobin::new(),
+            committee_membership: FisherYatesShuffle::new(ENTROPY),
         });
 
         assert_eq!(*overlay.leader(), nodes[0]);
@@ -231,13 +247,13 @@ mod tests {
         let mut overlay = TreeOverlay::new(TreeOverlaySettings {
             nodes: nodes.clone(),
             current_leader: nodes[0],
-            entropy: [0; 32],
             number_of_committees: 3,
             leader: RoundRobin::new(),
+            committee_membership: FisherYatesShuffle::new(ENTROPY),
         });
 
         let leader = overlay.next_leader();
-        overlay = overlay.advance([1; 32], RoundRobin::new());
+        overlay = overlay.advance(RoundRobin::new(), FisherYatesShuffle::new(ENTROPY));
 
         assert_eq!(leader, *overlay.leader());
     }
@@ -248,9 +264,9 @@ mod tests {
         let overlay = TreeOverlay::new(TreeOverlaySettings {
             current_leader: nodes[0],
             nodes,
-            entropy: [0; 32],
             number_of_committees: 3,
             leader: RoundRobin::new(),
+            committee_membership: FisherYatesShuffle::new(ENTROPY),
         });
 
         let mut expected_root = Committee::new();
@@ -266,9 +282,9 @@ mod tests {
         let overlay = TreeOverlay::new(TreeOverlaySettings {
             current_leader: nodes[0],
             nodes,
-            entropy: [0; 32],
             number_of_committees: 3,
             leader: RoundRobin::new(),
+            committee_membership: FisherYatesShuffle::new(ENTROPY),
         });
 
         let mut leaf_committees = overlay
@@ -296,9 +312,9 @@ mod tests {
         let overlay = TreeOverlay::new(TreeOverlaySettings {
             current_leader: nodes[0],
             nodes,
-            entropy: [0; 32],
             number_of_committees: 3,
             leader: RoundRobin::new(),
+            committee_membership: FisherYatesShuffle::new(ENTROPY),
         });
 
         assert_eq!(overlay.super_majority_threshold(overlay.nodes[8]), 0);
@@ -310,9 +326,9 @@ mod tests {
         let overlay = TreeOverlay::new(TreeOverlaySettings {
             current_leader: nodes[0],
             nodes,
-            entropy: [0; 32],
             number_of_committees: 3,
             leader: RoundRobin::new(),
+            committee_membership: FisherYatesShuffle::new(ENTROPY),
         });
 
         assert_eq!(overlay.super_majority_threshold(overlay.nodes[0]), 3);
@@ -324,9 +340,9 @@ mod tests {
         let overlay = TreeOverlay::new(TreeOverlaySettings {
             nodes: nodes.clone(),
             current_leader: nodes[0],
-            entropy: [0; 32],
             number_of_committees: 3,
             leader: RoundRobin::new(),
+            committee_membership: FisherYatesShuffle::new(ENTROPY),
         });
 
         assert_eq!(

consensus-engine/tests/fuzz/sut.rs

@@ -1,5 +1,6 @@
 use std::{collections::HashSet, panic};
 
+use consensus_engine::overlay::FreezeMembership;
 use consensus_engine::{
     overlay::{FlatOverlay, FlatOverlaySettings, RoundRobin},
     *,
@@ -13,7 +14,7 @@ use crate::fuzz::transition::Transition;
 // This is called as SUT (System Under Test).
 #[derive(Clone, Debug)]
 pub struct ConsensusEngineTest {
-    pub engine: Carnot<FlatOverlay<RoundRobin>>,
+    pub engine: Carnot<FlatOverlay<RoundRobin, FreezeMembership>>,
 }
 
 impl ConsensusEngineTest {

nodes/nomos-node/src/lib.rs

@@ -1,7 +1,7 @@
 mod tx;
 
 use color_eyre::eyre::Result;
-use consensus_engine::overlay::{FlatOverlay, RoundRobin};
+use consensus_engine::overlay::{FlatOverlay, RandomBeaconState, RoundRobin};
 #[cfg(feature = "metrics")]
 use metrics::{backend::map::MapMetricsBackend, types::MetricsData, MetricsService};
 use nomos_consensus::{
@@ -38,7 +38,7 @@ pub type Carnot = CarnotConsensus<
     MockPool<Tx>,
     MempoolWakuAdapter<Tx>,
     MockFountain,
-    FlatOverlay<RoundRobin>,
+    FlatOverlay<RoundRobin, RandomBeaconState>,
 >;
 
 #[derive(Services)]

nomos-services/consensus/src/committee_membership/mod.rs

@@ -0,0 +1,53 @@
+use std::convert::Infallible;
+// std
+use std::error::Error;
+use std::hash::Hash;
+
+// crates
+
+// internal
+use consensus_engine::overlay::{
+    CommitteeMembership, Error as RandomBeaconError, FreezeMembership, RandomBeaconState,
+};
+use consensus_engine::TimeoutQc;
+use nomos_core::block::Block;
+
+pub trait UpdateableCommitteeMembership: CommitteeMembership {
+    type Error: Error;
+
+    fn on_new_block_received<Tx: Hash + Clone + Eq>(
+        &self,
+        block: &Block<Tx>,
+    ) -> Result<Self, Self::Error>;
+    fn on_timeout_qc_received(&self, qc: &TimeoutQc) -> Result<Self, Self::Error>;
+}
+
+impl UpdateableCommitteeMembership for FreezeMembership {
+    type Error = Infallible;
+
+    fn on_new_block_received<Tx: Hash + Clone + Eq>(
+        &self,
+        _block: &Block<Tx>,
+    ) -> Result<Self, Self::Error> {
+        Ok(Self)
+    }
+
+    fn on_timeout_qc_received(&self, _qc: &TimeoutQc) -> Result<Self, Self::Error> {
+        Ok(Self)
+    }
+}
+
+impl UpdateableCommitteeMembership for RandomBeaconState {
+    type Error = RandomBeaconError;
+
+    fn on_new_block_received<Tx: Hash + Clone + Eq>(
+        &self,
+        block: &Block<Tx>,
+    ) -> Result<Self, Self::Error> {
+        self.check_advance_happy(block.beacon().clone(), block.header().parent_qc.view())
+    }
+
+    fn on_timeout_qc_received(&self, qc: &TimeoutQc) -> Result<Self, Self::Error> {
+        Ok(Self::generate_sad(qc.view(), self))
+    }
+}

nomos-services/consensus/src/leader_selection/mod.rs

@@ -1,5 +1,6 @@
+use consensus_engine::overlay::RoundRobin;
 use consensus_engine::{
-    overlay::{Error as RandomBeaconError, LeaderSelection, RandomBeaconState, RoundRobin},
+    overlay::{Error as RandomBeaconError, LeaderSelection, RandomBeaconState},
     TimeoutQc,
 };
 use nomos_core::block::Block;
@@ -10,9 +11,9 @@ pub trait UpdateableLeaderSelection: LeaderSelection {
 
     fn on_new_block_received<Tx: Hash + Clone + Eq>(
         &self,
-        block: Block<Tx>,
+        block: &Block<Tx>,
     ) -> Result<Self, Self::Error>;
-    fn on_timeout_qc_received(&self, qc: TimeoutQc) -> Result<Self, Self::Error>;
+    fn on_timeout_qc_received(&self, qc: &TimeoutQc) -> Result<Self, Self::Error>;
 }
 
 impl UpdateableLeaderSelection for RoundRobin {
@@ -20,12 +21,12 @@ impl UpdateableLeaderSelection for RoundRobin {
 
     fn on_new_block_received<Tx: Hash + Clone + Eq>(
         &self,
-        _block: Block<Tx>,
+        _block: &Block<Tx>,
     ) -> Result<Self, Self::Error> {
         Ok(self.advance())
     }
 
-    fn on_timeout_qc_received(&self, _qc: TimeoutQc) -> Result<Self, Self::Error> {
+    fn on_timeout_qc_received(&self, _qc: &TimeoutQc) -> Result<Self, Self::Error> {
         Ok(self.advance())
     }
 }
@@ -35,13 +36,13 @@ impl UpdateableLeaderSelection for RandomBeaconState {
 
     fn on_new_block_received<Tx: Hash + Clone + Eq>(
         &self,
-        block: Block<Tx>,
+        block: &Block<Tx>,
     ) -> Result<Self, Self::Error> {
         self.check_advance_happy(block.beacon().clone(), block.header().parent_qc.view())
         // TODO: check random beacon public keys is leader id
     }
 
-    fn on_timeout_qc_received(&self, qc: TimeoutQc) -> Result<Self, Self::Error> {
+    fn on_timeout_qc_received(&self, qc: &TimeoutQc) -> Result<Self, Self::Error> {
         Ok(Self::generate_sad(qc.view(), self))
     }
 }

nomos-services/consensus/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod committee_membership;
 pub mod leader_selection;
 pub mod network;
 mod tally;
@@ -33,6 +34,7 @@ use consensus_engine::{
 };
 use task_manager::TaskManager;
 
+use crate::committee_membership::UpdateableCommitteeMembership;
 use nomos_core::block::Block;
 use nomos_core::fountain::FountainCode;
 use nomos_core::tx::Transaction;
@@ -144,6 +146,7 @@ where
     M: MempoolAdapter<Tx = P::Tx> + Send + Sync + 'static,
     O: Overlay + Debug + Send + Sync + 'static,
     O::LeaderSelection: UpdateableLeaderSelection,
+    O::CommitteeMembership: UpdateableCommitteeMembership,
 {
     fn init(service_state: ServiceStateHandle<Self>) -> Result<Self, overwatch_rs::DynError> {
         let network_relay = service_state.overwatch_handle.relay();
@@ -283,6 +286,7 @@ where
     M: MempoolAdapter<Tx = P::Tx> + Send + Sync + 'static,
     O: Overlay + Debug + Send + Sync + 'static,
     O::LeaderSelection: UpdateableLeaderSelection,
+    O::CommitteeMembership: UpdateableCommitteeMembership,
 {
     fn process_message(carnot: &Carnot<O>, msg: ConsensusMsg) {
         match msg {
@@ -426,9 +430,13 @@ where
                             tally_settings,
                         ),
                     );
-                    new_state = Self::update_leader_selection(new_state, |leader_selection| {
-                        leader_selection.on_new_block_received(original_block)
-                    });
+                    new_state = Self::update_overlay(
+                        new_state,
+                        |leader_selection| leader_selection.on_new_block_received(&original_block),
+                        |committee_membership| {
+                            committee_membership.on_new_block_received(&original_block)
+                        },
+                    );
                 } else {
                     task_manager.push(block.view, async move {
                         if let Some(block) = stream.next().await {
@@ -514,9 +522,11 @@ where
             Self::gather_new_views(adapter, self_committee, timeout_qc.clone(), tally_settings),
         );
         if carnot.current_view() != new_state.current_view() {
-            new_state = Self::update_leader_selection(new_state, |leader_selection| {
-                leader_selection.on_timeout_qc_received(timeout_qc)
-            });
+            new_state = Self::update_overlay(
+                new_state,
+                |leader_selection| leader_selection.on_timeout_qc_received(&timeout_qc),
+                |committee_membership| committee_membership.on_timeout_qc_received(&timeout_qc),
+            );
         }
         (new_state, None)
     }
@@ -723,16 +733,41 @@ where
 
     fn update_leader_selection<
         E: std::error::Error,
-        U: FnOnce(O::LeaderSelection) -> Result<O::LeaderSelection, E>,
+        Fl: FnOnce(O::LeaderSelection) -> Result<O::LeaderSelection, E>,
     >(
         carnot: Carnot<O>,
-        f: U,
+        leader_selection_f: Fl,
     ) -> Carnot<O> {
         carnot
-            .update_overlay(|overlay| overlay.update_leader_selection(f))
-            // TODO: remove unwrap
+            .update_overlay(|overlay| overlay.update_leader_selection(leader_selection_f))
             .unwrap()
     }
+
+    fn update_committee_membership<
+        E: std::error::Error,
+        Fm: FnOnce(O::CommitteeMembership) -> Result<O::CommitteeMembership, E>,
+    >(
+        carnot: Carnot<O>,
+        committee_membership_f: Fm,
+    ) -> Carnot<O> {
+        carnot
+            .update_overlay(|overlay| overlay.update_committees(committee_membership_f))
+            .unwrap()
+    }
+
+    fn update_overlay<
+        El: std::error::Error,
+        Em: std::error::Error,
+        Fl: FnOnce(O::LeaderSelection) -> Result<O::LeaderSelection, El>,
+        Fm: FnOnce(O::CommitteeMembership) -> Result<O::CommitteeMembership, Em>,
+    >(
+        carnot: Carnot<O>,
+        leader_selection_f: Fl,
+        committee_membership_f: Fm,
+    ) -> Carnot<O> {
+        let carnot = Self::update_leader_selection(carnot, leader_selection_f);
+        Self::update_committee_membership(carnot, committee_membership_f)
+    }
 }
 
 async fn handle_output<A, F, Tx>(adapter: &A, fountain: &F, node_id: NodeId, output: Output<Tx>)

simulations/src/bin/app/overlay_node.rs

@@ -1,5 +1,6 @@
+use consensus_engine::overlay::RandomBeaconState;
 use consensus_engine::{
-    overlay::{FlatOverlay, RoundRobin, TreeOverlay},
+    overlay::{FlatOverlay, FreezeMembership, RoundRobin, TreeOverlay},
     NodeId,
 };
 use rand::Rng;
@@ -26,37 +27,41 @@ pub fn to_overlay_node<R: Rng>(
                 leader: RoundRobin::new(),
                 leader_super_majority_threshold: None,
             };
-            Box::new(CarnotNode::<FlatOverlay<RoundRobin>>::new(
-                node_id,
-                CarnotSettings::new(
-                    settings.node_settings.timeout,
-                    settings.record_settings.clone(),
+            Box::new(
+                CarnotNode::<FlatOverlay<RoundRobin, FreezeMembership>>::new(
+                    node_id,
+                    CarnotSettings::new(
+                        settings.node_settings.timeout,
+                        settings.record_settings.clone(),
+                    ),
+                    overlay_settings,
+                    genesis,
+                    network_interface,
+                    &mut rng,
                 ),
-                overlay_settings,
-                genesis,
-                network_interface,
-                &mut rng,
-            ))
+            )
         }
         simulations::settings::OverlaySettings::Tree(tree_settings) => {
             let overlay_settings = consensus_engine::overlay::TreeOverlaySettings {
                 nodes,
                 current_leader: leader,
-                entropy: [0; 32],
                 number_of_committees: tree_settings.number_of_committees,
                 leader: RoundRobin::new(),
+                committee_membership: RandomBeaconState::initial_sad_from_entropy([0; 32]),
             };
-            Box::new(CarnotNode::<TreeOverlay<RoundRobin>>::new(
-                node_id,
-                CarnotSettings::new(
-                    settings.node_settings.timeout,
-                    settings.record_settings.clone(),
+            Box::new(
+                CarnotNode::<TreeOverlay<RoundRobin, RandomBeaconState>>::new(
+                    node_id,
+                    CarnotSettings::new(
+                        settings.node_settings.timeout,
+                        settings.record_settings.clone(),
+                    ),
+                    overlay_settings,
+                    genesis,
+                    network_interface,
+                    &mut rng,
                 ),
-                overlay_settings,
-                genesis,
-                network_interface,
-                &mut rng,
-            ))
+            )
         }
     }
 }

simulations/src/node/carnot/mod.rs

@@ -24,6 +24,7 @@ use consensus_engine::overlay::RandomBeaconState;
 use consensus_engine::{
     Block, BlockId, Carnot, Committee, Overlay, Payload, Qc, StandardQc, TimeoutQc, View, Vote,
 };
+use nomos_consensus::committee_membership::UpdateableCommitteeMembership;
 use nomos_consensus::network::messages::{ProposalChunkMsg, TimeoutQcMsg};
 use nomos_consensus::{
     leader_selection::UpdateableLeaderSelection,
@@ -224,7 +225,12 @@ pub struct CarnotNode<O: Overlay> {
     step_duration: Duration,
 }
 
-impl<L: UpdateableLeaderSelection, O: Overlay<LeaderSelection = L>> CarnotNode<O> {
+impl<
+        L: UpdateableLeaderSelection,
+        M: UpdateableCommitteeMembership,
+        O: Overlay<LeaderSelection = L, CommitteeMembership = M>,
+    > CarnotNode<O>
+{
     pub fn new<R: Rng>(
         id: consensus_engine::NodeId,
         settings: CarnotSettings,
@@ -342,8 +348,13 @@ impl<L: UpdateableLeaderSelection, O: Overlay<LeaderSelection = L>> CarnotNode<O
                         if self.engine.current_view() != new.current_view() {
                             new = new
                                 .update_overlay(|overlay| {
-                                    overlay.update_leader_selection(|leader_selection| {
-                                        leader_selection.on_new_block_received(block.clone())
+                                    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);
@@ -481,7 +492,12 @@ impl<L: UpdateableLeaderSelection, O: Overlay<LeaderSelection = L>> CarnotNode<O
     }
 }
 
-impl<L: UpdateableLeaderSelection, O: Overlay<LeaderSelection = L>> Node for CarnotNode<O> {
+impl<
+        L: UpdateableLeaderSelection,
+        M: UpdateableCommitteeMembership,
+        O: Overlay<LeaderSelection = L, CommitteeMembership = M>,
+    > Node for CarnotNode<O>
+{
     type Settings = CarnotSettings;
     type State = CarnotState;