Introduce KeyShares struct to represent the amount of key shares

Improvements, bug fixes associated.

substrate/primitives/src/address.rs

@@ -24,7 +24,10 @@ const HUMAN_READABLE_PART: bech32::Hrp = bech32::Hrp::parse_unchecked("sri");
 
 /// The address for an account on Serai.
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, Zeroize, BorshSerialize, BorshDeserialize)]
-#[cfg_attr(feature = "non_canonical_scale_derivations", derive(scale::Encode, scale::Decode))]
+#[cfg_attr(
+  feature = "non_canonical_scale_derivations",
+  derive(scale::Encode, scale::Decode, scale::MaxEncodedLen)
+)]
 pub struct SeraiAddress(pub [u8; 32]);
 
 // These share encodings as 32-byte arrays

substrate/primitives/src/constants.rs

@@ -6,8 +6,3 @@ pub const TARGET_BLOCK_TIME: Duration = Duration::from_secs(6);
 /// The intended duration for a session.
 // 1 week
 pub const SESSION_LENGTH: Duration = Duration::from_secs(7 * 24 * 60 * 60);
-
-/// The maximum amount of key shares per set.
-pub const MAX_KEY_SHARES_PER_SET: u16 = 150;
-/// The maximum amount of key shares per set, as an u32.
-pub const MAX_KEY_SHARES_PER_SET_U32: u32 = MAX_KEY_SHARES_PER_SET as u32;

substrate/primitives/src/signals.rs

@@ -1,16 +1,46 @@
 use zeroize::Zeroize;
 use borsh::{BorshSerialize, BorshDeserialize};
 
-use crate::network_id::ExternalNetworkId;
+use crate::{network_id::ExternalNetworkId, address::SeraiAddress};
+
+/// The ID of an protocol.
+pub type ProtocolId = [u8; 32];
 
 /// A signal.
 #[derive(Clone, Copy, PartialEq, Eq, Debug, Zeroize, BorshSerialize, BorshDeserialize)]
+#[cfg_attr(
+  feature = "non_canonical_scale_derivations",
+  allow(clippy::cast_possible_truncation),
+  derive(scale::Encode, scale::Decode, scale::MaxEncodedLen)
+)]
 pub enum Signal {
   /// A signal to retire the current protocol.
   Retire {
     /// The protocol to retire in favor of.
-    in_favor_of: [u8; 32],
+    in_favor_of: ProtocolId,
   },
   /// A signal to halt an external network.
   Halt(ExternalNetworkId),
 }
+
+/// A retirement signal, registered on chain.
+#[derive(Clone, Copy, PartialEq, Eq, Debug, Zeroize, BorshSerialize, BorshDeserialize)]
+#[cfg_attr(
+  feature = "non_canonical_scale_derivations",
+  derive(scale::Encode, scale::Decode, scale::MaxEncodedLen)
+)]
+pub struct RegisteredRetirementSignal {
+  /// The protocol to retire in favor of.
+  pub in_favor_of: ProtocolId,
+  /// The registrant of this signal.
+  pub registrant: SeraiAddress,
+  /// The block number this was registered at.
+  pub registered_at: u64,
+}
+
+impl RegisteredRetirementSignal {
+  /// The ID of this signal.
+  pub fn id(&self) -> ProtocolId {
+    sp_core::blake2_256(&borsh::to_vec(self).unwrap())
+  }
+}

substrate/primitives/src/validator_sets/mod.rs

@@ -7,9 +7,9 @@ use ciphersuite::{group::GroupEncoding, GroupIo};
 use dalek_ff_group::Ristretto;
 
 use crate::{
-  constants::MAX_KEY_SHARES_PER_SET,
   crypto::{Public, KeyPair},
   network_id::{ExternalNetworkId, NetworkId},
+  balance::Amount,
 };
 
 mod slashes;
@@ -103,19 +103,84 @@ impl ExternalValidatorSet {
   }
 }
 
-/// For a set of validators whose key shares may exceed the maximum, reduce until they are less
-/// than or equal to the maximum.
-///
-/// This runs in time linear to the exceed key shares and assumes the excess fits within a usize,
-/// panicking otherwise.
-///
-/// Reduction occurs by reducing each validator in a reverse round-robin. This means the worst
-/// validators lose their key shares first.
-pub fn amortize_excess_key_shares(validators: &mut [(sp_core::sr25519::Public, u64)]) {
-  let total_key_shares = validators.iter().map(|(_key, shares)| shares).sum::<u64>();
-  for i in 0 .. usize::try_from(total_key_shares.saturating_sub(u64::from(MAX_KEY_SHARES_PER_SET)))
-    .unwrap()
-  {
-    validators[validators.len() - ((i % validators.len()) + 1)].1 -= 1;
+/// The representation for an amount of key shares.
+#[derive(Clone, Copy, PartialEq, Eq, Debug, Zeroize, BorshSerialize, BorshDeserialize)]
+#[cfg_attr(
+  feature = "non_canonical_scale_derivations",
+  derive(scale::Encode, scale::Decode, scale::MaxEncodedLen)
+)]
+pub struct KeyShares(pub u16);
+
+impl KeyShares {
+  /// One key share.
+  pub const ONE: KeyShares = KeyShares(1);
+  /// The maximum amount of key shares per set.
+  pub const MAX_PER_SET: u16 = 150;
+  /// The maximum amount of key shares per set, represented as a `u32`.
+  pub const MAX_PER_SET_U32: u32 = 150;
+
+  /// Create key shares from a `u16`.
+  ///
+  /// This will saturate the value if the `u16` exceeds the maximum amount of key shares.
+  pub fn saturating_from(key_shares: u16) -> KeyShares {
+    KeyShares(key_shares.min(Self::MAX_PER_SET))
+  }
+
+  /// Create key shares from an allocation.
+  ///
+  /// Presumably panics if `allocation_per_key_share` is zero.
+  pub fn from_allocation(allocation: Amount, allocation_per_key_share: Amount) -> Self {
+    Self::saturating_from(
+      u16::try_from(allocation.0 / allocation_per_key_share.0).unwrap_or(u16::MAX),
+    )
+  }
+
+  /// For a set of validators whose key shares may exceed the maximum, reduce until they are less
+  /// than or equal to the maximum.
+  ///
+  /// Returns the new amount of validators with a non-zero amount of key shares.
+  ///
+  /// This runs in time linear to the exceeded key shares and may panic if:
+  ///   - The total amount of key shares exceeds `u16::MAX`.
+  ///   - The list of validators is absurdly long
+  ///   - The list of validators includes validators without key shares
+  ///
+  /// Reduction occurs by reducing each validator in a reverse round-robin. This means the
+  /// validators with the least key shares are evicted first.
+  #[must_use]
+  pub fn amortize_excess(validators: &mut [(sp_core::sr25519::Public, KeyShares)]) -> usize {
+    let total_key_shares = validators.iter().map(|(_key, shares)| shares.0).sum::<u16>();
+    let mut actual_len = validators.len();
+    let mut offset = 1;
+    for _ in 0 .. usize::from(total_key_shares.saturating_sub(Self::MAX_PER_SET)) {
+      // If the offset exceeds the new length, reset it
+      if offset > actual_len {
+        offset = 1;
+      }
+
+      // Take one key share from this validator
+      let index = actual_len - offset;
+      validators[index].1 .0 -= 1;
+      // If they now have zero key shares, shrink the length and continue
+      if validators[index].1 .0 == 0 {
+        actual_len -= 1;
+        continue;
+      }
+
+      // Increment the offset to take from the next validator on the next iteration
+      offset += 1;
+    }
+    actual_len
+  }
+}
+
+impl TryFrom<u16> for KeyShares {
+  type Error = ();
+  fn try_from(value: u16) -> Result<Self, ()> {
+    if value > Self::MAX_PER_SET {
+      Err(())
+    } else {
+      Ok(Self(value))
+    }
   }
 }

substrate/primitives/src/validator_sets/slashes.rs

@@ -8,8 +8,9 @@ use borsh::{BorshSerialize, BorshDeserialize};
 use sp_core::{ConstU32, bounded::BoundedVec};
 
 use crate::{
-  constants::{TARGET_BLOCK_TIME, SESSION_LENGTH, MAX_KEY_SHARES_PER_SET_U32},
+  constants::{TARGET_BLOCK_TIME, SESSION_LENGTH},
   balance::Amount,
+  validator_sets::KeyShares,
 };
 
 /// Each slash point is equivalent to the downtime implied by missing a block proposal.
@@ -212,7 +213,7 @@ pub struct SlashReport(
     serialize_with = "crate::borsh_serialize_bounded_vec",
     deserialize_with = "crate::borsh_deserialize_bounded_vec"
   )]
-  pub BoundedVec<Slash, ConstU32<{ MAX_KEY_SHARES_PER_SET_U32 }>>,
+  pub BoundedVec<Slash, ConstU32<{ KeyShares::MAX_PER_SET_U32 }>>,
 );
 
 /// An error when converting from a `Vec`.
@@ -251,7 +252,7 @@ impl SlashReport {
 
 #[test]
 fn test_penalty() {
-  for validators in [1, 50, 100, crate::constants::MAX_KEY_SHARES_PER_SET] {
+  for validators in [1, 50, 100, KeyShares::MAX_PER_SET_U32] {
     let validators = NonZero::new(validators).unwrap();
     // 12 hours of slash points should only decrease the rewards proportionately
     let twelve_hours_of_slash_points =

substrate/validator-sets/src/allocations.rs

@@ -1,6 +1,6 @@
 use sp_core::{Encode, sr25519::Public};
 
-use serai_primitives::{constants::MAX_KEY_SHARES_PER_SET, network_id::NetworkId, balance::Amount};
+use serai_primitives::{network_id::NetworkId, balance::Amount, validator_sets::KeyShares};
 
 use frame_support::storage::{StorageMap, StoragePrefixedMap};
 
@@ -63,7 +63,7 @@ pub(crate) trait Allocations {
   ) -> impl Iterator<Item = (Public, Amount)>;
 
   /// Calculate the expected key shares for a network, per the current allocations.
-  fn expected_key_shares(network: NetworkId, allocation_per_key_share: Amount) -> u64;
+  fn expected_key_shares(network: NetworkId, allocation_per_key_share: Amount) -> KeyShares;
 }
 
 /// Reverses the lexicographic order of a given byte array.
@@ -149,17 +149,16 @@ impl<Storage: AllocationsStorage> Allocations for Storage {
     .filter(move |(_key, allocation)| *allocation >= minimum_allocation)
   }
 
-  fn expected_key_shares(network: NetworkId, allocation_per_key_share: Amount) -> u64 {
+  fn expected_key_shares(network: NetworkId, allocation_per_key_share: Amount) -> KeyShares {
     let mut total_key_shares = 0;
     for (_, amount) in Self::iter_allocations(network, allocation_per_key_share) {
-      let key_shares = amount.0 / allocation_per_key_share.0;
-      total_key_shares += key_shares;
+      total_key_shares += KeyShares::from_allocation(amount, allocation_per_key_share).0;
 
-      if total_key_shares >= u64::from(MAX_KEY_SHARES_PER_SET) {
+      if total_key_shares >= KeyShares::MAX_PER_SET {
         break;
       }
     }
-    total_key_shares
+    KeyShares::saturating_from(total_key_shares)
   }
 }
 

substrate/validator-sets/src/lib.rs

@@ -81,7 +81,12 @@ mod pallet {
   use frame_support::pallet_prelude::*;
 
   use serai_primitives::{
-    crypto::KeyPair, network_id::*, coin::*, balance::*, validator_sets::*, address::SeraiAddress,
+    crypto::KeyPair,
+    network_id::*,
+    coin::*,
+    balance::*,
+    validator_sets::{Session, ExternalValidatorSet, ValidatorSet, KeyShares as KeySharesStruct},
+    address::SeraiAddress,
   };
 
   use coins_pallet::Pallet as Coins;
@@ -197,10 +202,12 @@ mod pallet {
   type CurrentSession<T: Config> = StorageMap<_, Identity, NetworkId, Session, OptionQuery>;
   #[pallet::storage]
   type LatestDecidedSession<T: Config> = StorageMap<_, Identity, NetworkId, Session, OptionQuery>;
+  #[pallet::storage]
+  type KeyShares<T: Config> = StorageMap<_, Identity, ValidatorSet, KeySharesStruct, OptionQuery>;
   // This has to use `Identity` per the documentation of `SessionsStorage`
   #[pallet::storage]
   type SelectedValidators<T: Config> =
-    StorageMap<_, Identity, SelectedValidatorsKey, u64, OptionQuery>;
+    StorageMap<_, Identity, SelectedValidatorsKey, KeySharesStruct, OptionQuery>;
   #[pallet::storage]
   type TotalAllocatedStake<T: Config> = StorageMap<_, Identity, NetworkId, Amount, OptionQuery>;
   #[pallet::storage]
@@ -212,6 +219,7 @@ mod pallet {
     type AllocationPerKeyShare = AllocationPerKeyShare<T>;
     type CurrentSession = CurrentSession<T>;
     type LatestDecidedSession = LatestDecidedSession<T>;
+    type KeyShares = KeyShares<T>;
     type SelectedValidators = SelectedValidators<T>;
     type TotalAllocatedStake = TotalAllocatedStake<T>;
     type DelayedDeallocations = DelayedDeallocations<T>;
@@ -341,6 +349,40 @@ mod pallet {
       SeraiAddress::system(b"ValidatorSets").into()
     }
 
+    /// The current session for a network.
+    pub fn current_session(network: NetworkId) -> Option<Session> {
+      Abstractions::<T>::current_session(network)
+    }
+
+    /// The latest decided session for a network.
+    pub fn latest_decided_session(network: NetworkId) -> Option<Session> {
+      Abstractions::<T>::latest_decided_session(network)
+    }
+
+    /// The amount of key shares a validator has.
+    ///
+    /// Returns `None` for historic sessions which we no longer have the data for.
+    pub fn key_shares(set: ValidatorSet) -> Option<KeySharesStruct> {
+      Abstractions::<T>::key_shares(set)
+    }
+
+    /// If a validator is present within the specified validator set.
+    ///
+    /// This MAY return `false` for _any_ historic session, even if the validator _was_ present,
+    pub fn in_validator_set(set: ValidatorSet, validator: Public) -> bool {
+      Abstractions::<T>::in_validator_set(set, validator)
+    }
+
+    /// The key shares possessed by a validator, within a validator set.
+    ///
+    /// This MAY return `None` for _any_ historic session, even if the validator _was_ present,
+    pub fn key_shares_possessed_by_validator(
+      set: ValidatorSet,
+      validator: Public,
+    ) -> Option<KeySharesStruct> {
+      Abstractions::<T>::key_shares_possessed_by_validator(set, validator)
+    }
+
     /*
     // is_bft returns if the network is able to survive any single node becoming byzantine.
     fn is_bft(network: NetworkId) -> bool {

substrate/validator-sets/src/sessions.rs

@@ -1,10 +1,9 @@
 use sp_core::{Encode, Decode, ConstU32, sr25519::Public, bounded::BoundedVec};
 
 use serai_primitives::{
-  constants::{MAX_KEY_SHARES_PER_SET, MAX_KEY_SHARES_PER_SET_U32},
   network_id::NetworkId,
   balance::Amount,
-  validator_sets::{Session, ValidatorSet, amortize_excess_key_shares},
+  validator_sets::{KeyShares as KeySharesStruct, Session, ValidatorSet},
 };
 
 use frame_support::storage::{StorageValue, StorageMap, StorageDoubleMap, StoragePrefixedMap};
@@ -12,7 +11,8 @@ use frame_support::storage::{StorageValue, StorageMap, StorageDoubleMap, Storage
 use crate::{embedded_elliptic_curve_keys::EmbeddedEllipticCurveKeys, allocations::Allocations};
 
 /// The list of genesis validators.
-pub(crate) type GenesisValidators = BoundedVec<Public, ConstU32<{ MAX_KEY_SHARES_PER_SET_U32 }>>;
+pub(crate) type GenesisValidators =
+  BoundedVec<Public, ConstU32<{ KeySharesStruct::MAX_PER_SET_U32 }>>;
 
 /// The key for the SelectedValidators map.
 pub(crate) type SelectedValidatorsKey = (ValidatorSet, [u8; 16], Public);
@@ -38,14 +38,23 @@ pub(crate) trait SessionsStorage: EmbeddedEllipticCurveKeys + Allocations {
   /// This is opaque and to be exclusively read/write by `Sessions`.
   type LatestDecidedSession: StorageMap<NetworkId, Session, Query = Option<Session>>;
 
+  /// The amount of key shares a validator set has.
+  ///
+  /// This is opaque and to be exclusively read/write by `Sessions`.
+  type KeyShares: StorageMap<ValidatorSet, KeySharesStruct, Query = Option<KeySharesStruct>>;
+
   /// The selected validators for a set.
   ///
   /// This MUST be instantiated with a map using `Identity` for its hasher.
   ///
   /// This is opaque and to be exclusively read/write by `Sessions`.
   // The value is how many key shares the validator has.
-  type SelectedValidators: StorageMap<SelectedValidatorsKey, u64, Query = Option<u64>>
-    + StoragePrefixedMap<u64>;
+  #[rustfmt::skip]
+  type SelectedValidators: StorageMap<
+    SelectedValidatorsKey,
+    KeySharesStruct,
+    Query = Option<KeySharesStruct>
+  > + StoragePrefixedMap<KeySharesStruct>;
 
   /// The total allocated stake for a network.
   ///
@@ -64,9 +73,9 @@ fn selected_validators_key(set: ValidatorSet, key: Public) -> SelectedValidators
   (set, hash, key)
 }
 
-fn selected_validators<Storage: StoragePrefixedMap<u64>>(
+fn selected_validators<Storage: StoragePrefixedMap<KeySharesStruct>>(
   set: ValidatorSet,
-) -> impl Iterator<Item = (Public, u64)> {
+) -> impl Iterator<Item = (Public, KeySharesStruct)> {
   let mut prefix = Storage::final_prefix().to_vec();
   prefix.extend(&set.encode());
   frame_support::storage::PrefixIterator::<_, ()>::new(
@@ -77,13 +86,13 @@ fn selected_validators<Storage: StoragePrefixedMap<u64>>(
         // Recover the validator's key from the storage key
         <[u8; 32]>::try_from(&key[(key.len() - 32) ..]).unwrap().into(),
         // Decode the key shares from the value
-        u64::decode(&mut key_shares).unwrap(),
+        KeySharesStruct::decode(&mut key_shares).unwrap(),
       ))
     },
   )
 }
 
-fn clear_selected_validators<Storage: StoragePrefixedMap<u64>>(set: ValidatorSet) {
+fn clear_selected_validators<Storage: StoragePrefixedMap<KeySharesStruct>>(set: ValidatorSet) {
   let mut prefix = Storage::final_prefix().to_vec();
   prefix.extend(&set.encode());
   assert!(matches!(
@@ -178,6 +187,30 @@ pub(crate) trait Sessions {
     network: NetworkId,
     session: Session,
   ) -> Result<Amount, DeallocationError>;
+
+  /// The currently active session for a network.
+  fn current_session(network: NetworkId) -> Option<Session>;
+
+  /// The latest decided session for a network.
+  fn latest_decided_session(network: NetworkId) -> Option<Session>;
+
+  /// The amount of key shares a validator has.
+  ///
+  /// Returns `None` for historic sessions which we no longer have the data for.
+  fn key_shares(set: ValidatorSet) -> Option<KeySharesStruct>;
+
+  /// If a validator is present within the specified validator set.
+  ///
+  /// This MAY return `false` for _any_ historic session, even if the validator _was_ present,
+  fn in_validator_set(set: ValidatorSet, validator: Public) -> bool;
+
+  /// The key shares possessed by a validator, within a validator set.
+  ///
+  /// This MAY return `None` for _any_ historic session, even if the validator _was_ present,
+  fn key_shares_possessed_by_validator(
+    set: ValidatorSet,
+    validator: Public,
+  ) -> Option<KeySharesStruct>;
 }
 
 impl<Storage: SessionsStorage> Sessions for Storage {
@@ -202,45 +235,40 @@ impl<Storage: SessionsStorage> Sessions for Storage {
       }
     }
 
-    let mut selected_validators = Vec::with_capacity(usize::from(MAX_KEY_SHARES_PER_SET / 2));
+    let mut selected_validators = Vec::with_capacity(usize::from(KeySharesStruct::MAX_PER_SET / 2));
     let mut total_key_shares = 0;
     if let Some(allocation_per_key_share) = Storage::AllocationPerKeyShare::get(network) {
       for (validator, amount) in Self::iter_allocations(network, allocation_per_key_share) {
-        // If this allocation is absurd, causing this to not fit within a u16, bound to the max
-        let key_shares = amount.0 / allocation_per_key_share.0;
+        let key_shares = KeySharesStruct::from_allocation(amount, allocation_per_key_share);
         selected_validators.push((validator, key_shares));
-        // We're tracking key shares as a u64 yet the max allowed is a u16, so this won't overflow
-        total_key_shares += key_shares;
-        if total_key_shares >= u64::from(MAX_KEY_SHARES_PER_SET) {
+        total_key_shares += key_shares.0;
+        if total_key_shares >= KeySharesStruct::MAX_PER_SET {
           break;
         }
       }
     }
 
     // Perform amortization if we've exceeded the maximum amount of key shares
-    // This is guaranteed not to cause any validators have zero key shares as we'd only be over if
-    // the last-added (worst) validator had multiple key shares, meaning everyone has more shares
-    // than we'll amortize here
-    amortize_excess_key_shares(selected_validators.as_mut_slice());
+    {
+      let new_len = KeySharesStruct::amortize_excess(selected_validators.as_mut_slice());
+      selected_validators.truncate(new_len);
+    }
 
     if include_genesis_validators {
       let mut genesis_validators = Storage::GenesisValidators::get()
         .expect("genesis validators wasn't set")
         .into_iter()
-        .map(|validator| (validator, 1))
+        .map(|validator| (validator, KeySharesStruct::ONE))
         .collect::<Vec<_>>();
-      let genesis_validator_key_shares = u64::try_from(genesis_validators.len()).unwrap();
-      while (total_key_shares + genesis_validator_key_shares) > u64::from(MAX_KEY_SHARES_PER_SET) {
+      let genesis_validator_key_shares = u16::try_from(genesis_validators.len()).unwrap();
+      total_key_shares += genesis_validator_key_shares;
+      while total_key_shares > KeySharesStruct::MAX_PER_SET {
         let (_key, key_shares) = selected_validators.pop().unwrap();
-        total_key_shares -= key_shares;
+        total_key_shares -= key_shares.0;
       }
       selected_validators.append(&mut genesis_validators);
-      total_key_shares += genesis_validator_key_shares;
     }
 
-    // We kept this accurate but don't actually further read from it
-    let _ = total_key_shares;
-
     let latest_decided_session = Storage::LatestDecidedSession::mutate(network, |session| {
       let next_session = session.map(|session| Session(session.0 + 1)).unwrap_or(Session(0));
       *session = Some(next_session);
@@ -248,6 +276,10 @@ impl<Storage: SessionsStorage> Sessions for Storage {
     });
 
     let latest_decided_set = ValidatorSet { network, session: latest_decided_session };
+    Storage::KeyShares::insert(
+      latest_decided_set,
+      KeySharesStruct::try_from(total_key_shares).expect("amortization failure"),
+    );
     for (key, key_shares) in selected_validators {
       Storage::SelectedValidators::insert(
         selected_validators_key(latest_decided_set, key),
@@ -285,10 +317,9 @@ impl<Storage: SessionsStorage> Sessions for Storage {
 
     // Clean-up the historic set's storage, if one exists
     if let Some(historic_session) = current.0.checked_sub(2).map(Session) {
-      clear_selected_validators::<Storage::SelectedValidators>(ValidatorSet {
-        network,
-        session: historic_session,
-      });
+      let historic_set = ValidatorSet { network, session: historic_session };
+      Storage::KeyShares::remove(historic_set);
+      clear_selected_validators::<Storage::SelectedValidators>(historic_set);
     }
   }
 
@@ -322,26 +353,28 @@ impl<Storage: SessionsStorage> Sessions for Storage {
     {
       // Check the validator set's current expected key shares
       let expected_key_shares = Self::expected_key_shares(network, allocation_per_key_share);
-      // Check if the top validator in this set may be faulty under this f
-      let top_validator_may_be_faulty = if let Some(top_validator) =
+      // Check if the top validator in this set may be faulty without causing a halt under this f
+      let currently_tolerates_single_point_of_failure = if let Some(top_validator) =
         Self::iter_allocations(network, allocation_per_key_share).next()
       {
         let (_key, amount) = top_validator;
-        let key_shares = amount.0 / allocation_per_key_share.0;
-        key_shares <= (expected_key_shares / 3)
+        let key_shares = KeySharesStruct::from_allocation(amount, allocation_per_key_share);
+        key_shares.0 <= (expected_key_shares.0 / 3)
       } else {
-        // If there are no validators, we claim the top validator may not be faulty so the
-        // following check doesn't run
         false
       };
-      if top_validator_may_be_faulty {
-        let old_key_shares = old_allocation.0 / allocation_per_key_share.0;
-        let new_key_shares = new_allocation.0 / allocation_per_key_share.0;
+      // If the set currently tolerates the fault of the top validator, don't let that change
+      if currently_tolerates_single_point_of_failure {
+        let old_key_shares =
+          KeySharesStruct::from_allocation(old_allocation, allocation_per_key_share);
+        let new_key_shares =
+          KeySharesStruct::from_allocation(new_allocation, allocation_per_key_share);
         // Update the amount of expected key shares per the key shares added
-        let expected_key_shares = (expected_key_shares + (new_key_shares - old_key_shares))
-          .min(u64::from(MAX_KEY_SHARES_PER_SET));
+        let expected_key_shares = KeySharesStruct::saturating_from(
+          expected_key_shares.0 + (new_key_shares.0 - old_key_shares.0),
+        );
         // If the new key shares exceeds the fault tolerance, don't allow the allocation
-        if new_key_shares > (expected_key_shares / 3) {
+        if new_key_shares.0 > (expected_key_shares.0 / 3) {
           Err(AllocationError::IntroducesSinglePointOfFailure)?
         }
       }
@@ -460,4 +493,27 @@ impl<Storage: SessionsStorage> Sessions for Storage {
     Storage::DelayedDeallocations::take(validator, session)
       .ok_or(DeallocationError::NoDelayedDeallocation)
   }
+
+  fn current_session(network: NetworkId) -> Option<Session> {
+    Storage::CurrentSession::get(network)
+  }
+
+  fn latest_decided_session(network: NetworkId) -> Option<Session> {
+    Storage::LatestDecidedSession::get(network)
+  }
+
+  fn key_shares(set: ValidatorSet) -> Option<KeySharesStruct> {
+    Storage::KeyShares::get(set)
+  }
+
+  fn in_validator_set(set: ValidatorSet, validator: Public) -> bool {
+    Storage::SelectedValidators::contains_key(selected_validators_key(set, validator))
+  }
+
+  fn key_shares_possessed_by_validator(
+    set: ValidatorSet,
+    validator: Public,
+  ) -> Option<KeySharesStruct> {
+    Storage::SelectedValidators::get(selected_validators_key(set, validator))
+  }
 }