From 52bb918cc9c29d030083f27f0eee995e804e6a6f Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Thu, 4 Jul 2024 02:11:04 -0400
Subject: [PATCH] Ensure TotalAllocatedStake is set for the first set

---
 substrate/validator-sets/pallet/src/lib.rs | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/substrate/validator-sets/pallet/src/lib.rs b/substrate/validator-sets/pallet/src/lib.rs
index 8a002bab7..d660f93e5 100644
--- a/substrate/validator-sets/pallet/src/lib.rs
+++ b/substrate/validator-sets/pallet/src/lib.rs
@@ -691,6 +691,15 @@ pub mod pallet {
       }
     }
 
+    fn set_total_allocated_stake(network: NetworkId) {
+      let participants = Participants::<T>::get(network)
+        .expect("setting TotalAllocatedStake for a network without participants");
+      let total_stake = participants.iter().fold(0, |acc, (addr, _)| {
+        acc + Allocations::<T>::get((network, addr)).unwrap_or(Amount(0)).0
+      });
+      TotalAllocatedStake::<T>::set(network, Some(Amount(total_stake)));
+    }
+
     // TODO: This is called retire_set, yet just starts retiring the set
     // Update the nomenclature within this function
     pub fn retire_set(set: ValidatorSet) {
@@ -715,12 +724,7 @@ pub mod pallet {
       });
 
       // Update the total allocated stake to be for the current set
-      let participants =
-        Participants::<T>::get(set.network).expect("set retired without a new set");
-      let total_stake = participants.iter().fold(0, |acc, (addr, _)| {
-        acc + Allocations::<T>::get((set.network, addr)).unwrap_or(Amount(0)).0
-      });
-      TotalAllocatedStake::<T>::set(set.network, Some(Amount(total_stake)));
+      self.set_total_allocated_stake(set.network);
     }
 
     /// Take the amount deallocatable.
@@ -907,6 +911,11 @@ pub mod pallet {
 
       Keys::<T>::set(set, Some(key_pair.clone()));
 
+      // If this is the first ever set for this network, set TotalAllocatedStake now
+      if session == Session(0) {
+        self.set_total_allocated_stake(network);
+      }
+
       // This does not remove from TotalAllocatedStake or InSet in order to:
       // 1) Not decrease the stake present in this set. This means removed participants are
       //    still liable for the economic security of the external network. This prevents