Store voters in unsorted bags to get good stake iteration properties

frame/staking/src/lib.rs

@@ -276,8 +276,9 @@ pub mod testing_utils;
 #[cfg(any(feature = "runtime-benchmarks", test))]
 pub mod benchmarking;
 
-pub mod slashing;
 pub mod inflation;
+pub mod slashing;
+pub mod voter_bags;
 pub mod weights;
 
 use sp_std::{
@@ -353,6 +354,9 @@ type NegativeImbalanceOf<T> = <<T as Config>::Currency as Currency<
 	<T as frame_system::Config>::AccountId,
 >>::NegativeImbalance;
 
+type AccountIdOf<T> = <T as frame_system::Config>::AccountId;
+type VotingDataOf<T> = (AccountIdOf<T>, VoteWeight, Vec<AccountIdOf<T>>);
+
 /// Information regarding the active era (era in used in session).
 #[derive(Encode, Decode, RuntimeDebug)]
 pub struct ActiveEraInfo {
@@ -994,6 +998,28 @@ decl_storage! {
 		///
 		/// This is set to v6.0.0 for new networks.
 		StorageVersion build(|_: &GenesisConfig<T>| Releases::V6_0_0): Releases;
+
+		// The next storage items collectively comprise the voter bags: a composite data structure
+		// designed to allow efficient iteration of the top N voters by stake, mostly. See
+		// `mod voter_bags` for details.
+
+		/// How many voters are registered.
+		VoterCount: u32;
+
+		/// Which bag currently contains a particular voter.
+		///
+		/// This may not be the appropriate bag for the voter's weight if they have been rewarded or
+		/// slashed.
+		VoterBagFor: map hasher(twox_64_concat) T::AccountId => voter_bags::BagIdx;
+
+		/// The head and tail of each bag of voters.
+		VoterBags: map hasher(twox_64_concat) voter_bags::BagIdx => Option<voter_bags::Bag<T>>;
+
+		/// The nodes comprising each bag.
+		VoterNodes: double_map
+			hasher(twox_64_concat) voter_bags::BagIdx,
+			hasher(twox_64_concat) T::AccountId =>
+			Option<voter_bags::Node<T>>;
 	}
 	add_extra_genesis {
 		config(stakers):
@@ -2495,37 +2521,20 @@ impl<T: Config> Module<T> {
 	/// auto-chilled.
 	///
 	/// Note that this is VERY expensive. Use with care.
-	pub fn get_npos_voters() -> Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)> {
-		let weight_of = Self::slashable_balance_of_fn();
-		let mut all_voters = Vec::new();
-
-		for (validator, _) in <Validators<T>>::iter() {
-			// append self vote
-			let self_vote = (validator.clone(), weight_of(&validator), vec![validator.clone()]);
-			all_voters.push(self_vote);
-		}
+	pub fn get_npos_voters(
+		maybe_max_len: Option<usize>,
+		voter_count: usize,
+	) -> Vec<VotingDataOf<T>> {
+		let wanted_voters = maybe_max_len.unwrap_or(voter_count).min(voter_count);
 
+		let weight_of = Self::slashable_balance_of_fn();
 		// collect all slashing spans into a BTreeMap for further queries.
 		let slashing_spans = <SlashingSpans<T>>::iter().collect::<BTreeMap<_, _>>();
 
-		for (nominator, nominations) in <Nominators<T>>::iter() {
-			let Nominations { submitted_in, mut targets, suppressed: _ } = nominations;
-
-			// Filter out nomination targets which were nominated before the most recent
-			// slashing span.
-			targets.retain(|stash| {
-				slashing_spans
-					.get(stash)
-					.map_or(true, |spans| submitted_in >= spans.last_nonzero_slash())
-			});
-
-			if !targets.is_empty() {
-				let vote_weight = weight_of(&nominator);
-				all_voters.push((nominator, vote_weight, targets))
-			}
-		}
-
-		all_voters
+		voter_bags::VoterList::<T>::iter()
+			.filter_map(|node| node.voting_data(&weight_of, &slashing_spans))
+			.take(wanted_voters)
+			.collect()
 	}
 
 	pub fn get_npos_targets() -> Vec<T::AccountId> {
@@ -2544,24 +2553,16 @@ impl<T: Config> frame_election_provider_support::ElectionDataProvider<T::Account
 	fn voters(
 		maybe_max_len: Option<usize>,
 	) -> data_provider::Result<(Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)>, Weight)> {
-		// NOTE: reading these counts already needs to iterate a lot of storage keys, but they get
-		// cached. This is okay for the case of `Ok(_)`, but bad for `Err(_)`, as the trait does not
-		// report weight in failures.
-		let nominator_count = <Nominators<T>>::iter().count();
-		let validator_count = <Validators<T>>::iter().count();
-		let voter_count = nominator_count.saturating_add(validator_count);
-
-		if maybe_max_len.map_or(false, |max_len| voter_count > max_len) {
-			return Err("Voter snapshot too big");
-		}
+		let voter_count = voter_bags::VoterList::<T>::decode_len().unwrap_or_default();
 
 		let slashing_span_count = <SlashingSpans<T>>::iter().count();
 		let weight = T::WeightInfo::get_npos_voters(
-			nominator_count as u32,
-			validator_count as u32,
+			// TODO: fix the weight calculation here
+			0 as u32,
+			voter_count as u32,
 			slashing_span_count as u32,
 		);
-		Ok((Self::get_npos_voters(), weight))
+		Ok((Self::get_npos_voters(maybe_max_len, voter_count), weight))
 	}
 
 	fn targets(maybe_max_len: Option<usize>) -> data_provider::Result<(Vec<T::AccountId>, Weight)> {