Start on the task to manage the swarm

coordinator/src/p2p/validators.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashSet, HashMap};
 
 use serai_client::{primitives::NetworkId, validator_sets::primitives::Session, Serai};
 
@@ -12,13 +12,18 @@ pub(crate) struct Validators {
   // A cache for which session we're populated with the validators of
   sessions: HashMap<NetworkId, Session>,
   // The validators by network
-  by_network: HashMap<NetworkId, Vec<PeerId>>,
-  // The set of all validators (as a HashMap<PeerId, usize> to represent the amount of inclusions)
-  set: HashMap<PeerId, usize>,
+  by_network: HashMap<NetworkId, HashSet<PeerId>>,
+  // The validators and their networks
+  validators: HashMap<PeerId, HashSet<NetworkId>>,
 }
 
 impl Validators {
-  pub(crate) async fn update(&mut self) -> Result<(), String> {
+  /// Update the view of the validators.
+  ///
+  /// Returns all validators removed from the active validator set.
+  pub(crate) async fn update(&mut self) -> Result<HashSet<PeerId>, String> {
+    let mut removed = HashSet::new();
+
     let temporal_serai =
       self.serai.as_of_latest_finalized_block().await.map_err(|e| format!("{e:?}"))?;
     let temporal_serai = temporal_serai.validator_sets();
@@ -35,20 +40,22 @@ impl Validators {
         let new_validators =
           temporal_serai.active_network_validators(network).await.map_err(|e| format!("{e:?}"))?;
         let new_validators =
-          new_validators.into_iter().map(peer_id_from_public).collect::<Vec<_>>();
+          new_validators.into_iter().map(peer_id_from_public).collect::<HashSet<_>>();
 
         // Remove the existing validators
-        for validator in self.by_network.remove(&network).unwrap_or(vec![]) {
-          let mut inclusions = self.set.remove(&validator).unwrap();
-          inclusions -= 1;
-          if inclusions != 0 {
-            self.set.insert(validator, inclusions);
+        for validator in self.by_network.remove(&network).unwrap_or_else(HashSet::new) {
+          let mut networks = self.validators.remove(&validator).unwrap();
+          networks.remove(&network);
+          if networks.is_empty() {
+            removed.insert(validator);
+          } else {
+            self.validators.insert(validator, networks);
           }
         }
 
         // Add the new validators
         for validator in new_validators.iter().copied() {
-          *self.set.entry(validator).or_insert(0) += 1;
+          self.validators.entry(validator).or_insert_with(HashSet::new).insert(network);
         }
         self.by_network.insert(network, new_validators);
 
@@ -56,14 +63,19 @@ impl Validators {
         self.sessions.insert(network, session);
       }
     }
-    Ok(())
+
+    Ok(removed)
   }
 
-  pub(crate) fn validators(&self) -> &HashMap<NetworkId, Vec<PeerId>> {
+  pub(crate) fn by_network(&self) -> &HashMap<NetworkId, HashSet<PeerId>> {
     &self.by_network
   }
 
   pub(crate) fn contains(&self, peer_id: &PeerId) -> bool {
-    self.set.contains_key(peer_id)
+    self.validators.contains_key(peer_id)
+  }
+
+  pub(crate) fn networks(&self, peer_id: &PeerId) -> Option<&HashSet<NetworkId>> {
+    self.validators.get(peer_id)
   }
 }