From bb25baf3bceaac39ba8305706891be6a998496cb Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Fri, 13 Oct 2023 01:04:41 -0400
Subject: [PATCH] Add logic to amortize excess key shares, correcting is_bft

---
 substrate/validator-sets/pallet/src/lib.rs    | 10 ++++---
 .../validator-sets/primitives/src/lib.rs      | 26 +++++++++++++++++++
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/substrate/validator-sets/pallet/src/lib.rs b/substrate/validator-sets/pallet/src/lib.rs
index 5fb254fc..81611390 100644
--- a/substrate/validator-sets/pallet/src/lib.rs
+++ b/substrate/validator-sets/pallet/src/lib.rs
@@ -187,9 +187,12 @@ pub mod pallet {
     fn is_bft(network: NetworkId) -> bool {
       let allocation_per_key_share = AllocationPerKeyShare::<T>::get(network).unwrap().0;
 
+      let mut validators_len = 0;
       let mut top = None;
       let mut key_shares = 0;
       for (_, amount) in SortedAllocationsIter::<T>::new(network) {
+        validators_len += 1;
+
         key_shares += amount.0 / allocation_per_key_share;
         if top.is_none() {
           top = Some(key_shares);
@@ -204,9 +207,10 @@ pub mod pallet {
 
       // key_shares may be over MAX_KEY_SHARES_PER_SET, which will cause an off-chain reduction of
       // each validator's key shares until their sum is MAX_KEY_SHARES_PER_SET
-      // That isn't modeled here, allowing an inaccuracy in an extreme edge case
-      // This may cause mis-reporting as BFT when not BFT (TODO: Investigate impact of this)
-      (top * 3) < key_shares
+      // post_amortization_key_shares_for_top_validator yields what the top validator's key shares
+      // would be after such a reduction, letting us evaluate this correctly
+      let top = post_amortization_key_shares_for_top_validator(validators_len, top, key_shares);
+      (top * 3) < key_shares.min(MAX_KEY_SHARES_PER_SET.into())
     }
   }
 
diff --git a/substrate/validator-sets/primitives/src/lib.rs b/substrate/validator-sets/primitives/src/lib.rs
index 592a5d0f..5942d92f 100644
--- a/substrate/validator-sets/primitives/src/lib.rs
+++ b/substrate/validator-sets/primitives/src/lib.rs
@@ -92,3 +92,29 @@ pub fn musig_key(set: ValidatorSet, set_keys: &[Public]) -> Public {
 pub fn set_keys_message(set: &ValidatorSet, key_pair: &KeyPair) -> Vec<u8> {
   [b"ValidatorSets-key_pair".as_ref(), &(set, key_pair).encode()].concat()
 }
+
+/// For a set of validators whose key shares may exceed the maximum, reduce until they equal the
+/// maximum.
+///
+/// Reduction occurs by reducing each validator in a reverse round-robin.
+pub fn amortize_excess_key_shares(validators: &mut [(Public, u64)]) {
+  let total_key_shares = validators.iter().map(|(_, shares)| shares).sum::<u64>();
+  for i in
+    0 .. usize::try_from(total_key_shares.saturating_sub(MAX_KEY_SHARES_PER_SET.into())).unwrap()
+  {
+    validators[validators.len() - ((i % validators.len()) + 1)].1 -= 1;
+  }
+}
+
+/// Returns the post-amortization key shares for the top validator.
+///
+/// Panics when `validators == 0`.
+pub fn post_amortization_key_shares_for_top_validator(
+  validators: usize,
+  top: u64,
+  key_shares: u64,
+) -> u64 {
+  top -
+    (key_shares.saturating_sub(MAX_KEY_SHARES_PER_SET.into()) /
+      u64::try_from(validators).unwrap())
+}