revamp closures.

frame/staking/src/lib.rs

@@ -1083,37 +1083,37 @@ decl_storage! {
 		/// Snapshot of validators at the beginning of the current election window. This should only
 		/// have a value when [`EraElectionStatus`] == `ElectionStatus::Open(_)`.
 		///
-		/// DEPRECATED.
+		/// TODO: TWO_PHASE: should be removed once we switch to two-phase.
 		pub SnapshotValidators get(fn snapshot_validators): Option<Vec<T::AccountId>>;
 
 		/// Snapshot of nominators at the beginning of the current election window. This should only
 		/// have a value when [`EraElectionStatus`] == `ElectionStatus::Open(_)`.
 		///
-		/// DEPRECATED.
+		/// TODO: TWO_PHASE: should be removed once we switch to two-phase.
 		pub SnapshotNominators get(fn snapshot_nominators): Option<Vec<T::AccountId>>;
 
 		/// The next validator set. At the end of an era, if this is available (potentially from the
 		/// result of an offchain worker), it is immediately used. Otherwise, the on-chain election
 		/// is executed.
 		///
-		/// DEPRECATED.
+		/// TODO: TWO_PHASE: should be removed once we switch to two-phase.
 		pub QueuedElected get(fn queued_elected): Option<ElectionResult<T::AccountId, BalanceOf<T>>>;
 
 		/// The score of the current [`QueuedElected`].
 		///
-		/// DEPRECATED.
+		/// TODO: TWO_PHASE: should be removed once we switch to two-phase.
 		pub QueuedScore get(fn queued_score): Option<ElectionScore>;
 
 		/// Flag to control the execution of the offchain election. When `Open(_)`, we accept
 		/// solutions to be submitted.
 		///
-		/// DEPRECATED.
+		/// TODO: TWO_PHASE: should be removed once we switch to two-phase.
 		pub EraElectionStatus get(fn era_election_status): ElectionStatus<T::BlockNumber>;
 
 		/// True if the current **planned** session is final. Note that this does not take era
 		/// forcing into account.
 		///
-		/// DEPRECATED.
+		/// TODO: TWO_PHASE: should be removed once we switch to two-phase.
 		pub IsCurrentSessionFinal get(fn is_current_session_final): bool = false;
 
 		/// True if network has been upgraded to this version.

frame/two-phase-election-provider/src/benchmarking.rs

@@ -127,11 +127,12 @@ where
 	// write the snapshot to staking or whoever is the data provider.
 	T::DataProvider::put_snapshot(all_voters.clone(), targets.clone());
 
-	crate::stake_of_fn!(let stake_of, all_voters, T);
-	crate::voter_index_fn!(let voter_index, all_voters, T);
-	crate::voter_at_fn!(let voter_at, all_voters, T);
-	crate::target_at_fn!(let target_at, targets, T);
-	crate::target_index_fn!(let target_index, targets, T);
+	let cache = helpers::generate_voter_cache::<T>(&all_voters);
+	let stake_of = helpers::stake_of_fn::<T>(&all_voters, &cache);
+	let voter_index = helpers::voter_index_fn::<T>(&cache);
+	let target_index = helpers::target_index_fn_linear::<T>(&targets);
+	let voter_at = helpers::voter_at_fn::<T>(&all_voters);
+	let target_at = helpers::target_at_fn::<T>(&targets);
 
 	let assignments = active_voters
 		.iter()

frame/two-phase-election-provider/src/helpers.rs

@@ -0,0 +1,208 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2020 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Some helper functions/macros for this crate.
+
+use super::{
+	Config, VoteWeight, CompactVoterIndexOf, CompactTargetIndexOf, CompactAccuracyOf,
+	OnChainAccuracyOf, ExtendedBalance,
+};
+use sp_runtime::InnerOf;
+use sp_std::{collections::btree_map::BTreeMap, convert::TryInto};
+
+#[macro_export]
+macro_rules! log {
+	($level:tt, $patter:expr $(, $values:expr)* $(,)?) => {
+		frame_support::debug::$level!(
+			target: $crate::LOG_TARGET,
+			concat!("🏦 ", $patter) $(, $values)*
+		)
+	};
+}
+
+/// Generate a btree-map cache of the voters and their indices.
+///
+/// This can be used to efficiently build index getter closures.
+pub fn generate_voter_cache<T: Config>(
+	snapshot: &Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)>,
+) -> BTreeMap<T::AccountId, usize>
+where
+	ExtendedBalance: From<InnerOf<CompactAccuracyOf<T>>>,
+	ExtendedBalance: From<InnerOf<OnChainAccuracyOf<T>>>,
+{
+	let mut cache: BTreeMap<T::AccountId, usize> = BTreeMap::new();
+	snapshot.iter().enumerate().for_each(|(i, (x, _, _))| {
+		let _existed = cache.insert(x.clone(), i);
+		// if a duplicate exists, we only consider the last one. Defensive only, should never happen.
+		debug_assert!(_existed.is_none());
+	});
+
+	cache
+}
+
+/// Create a function the returns the index a voter in the snapshot.
+///
+/// The returning index type is the same as the one defined in [`T::CompactSolution::Voter`].
+///
+/// ## Warning
+///
+/// The snapshot must be the same is the one used to create `cache`.
+pub fn voter_index_fn<T: Config>(
+	cache: &BTreeMap<T::AccountId, usize>,
+) -> Box<dyn Fn(&T::AccountId) -> Option<CompactVoterIndexOf<T>> + '_>
+where
+	ExtendedBalance: From<InnerOf<CompactAccuracyOf<T>>>,
+	ExtendedBalance: From<InnerOf<OnChainAccuracyOf<T>>>,
+{
+	Box::new(move |who| {
+		cache
+			.get(who)
+			.and_then(|i| <usize as TryInto<CompactVoterIndexOf<T>>>::try_into(*i).ok())
+	})
+}
+
+/// Same as [`voter_index_fn`], but the returning index is converted into usize, if possible.
+///
+/// ## Warning
+///
+/// The snapshot must be the same is the one used to create `cache`.
+pub fn voter_index_fn_usize<T: Config>(
+	cache: &BTreeMap<T::AccountId, usize>,
+) -> Box<dyn Fn(&T::AccountId) -> Option<usize> + '_>
+where
+	ExtendedBalance: From<InnerOf<CompactAccuracyOf<T>>>,
+	ExtendedBalance: From<InnerOf<OnChainAccuracyOf<T>>>,
+{
+	Box::new(move |who| cache.get(who).cloned())
+}
+
+/// A non-optimized, linear version of [`voter_index_fn`] that does not need a cache and does a
+/// linear search.
+///
+/// ## Warning
+///
+/// Not meant to be used in production.
+pub fn voter_index_fn_linear<T: Config>(
+	snapshot: &Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)>,
+) -> Box<dyn Fn(&T::AccountId) -> Option<CompactVoterIndexOf<T>> + '_>
+where
+	ExtendedBalance: From<InnerOf<CompactAccuracyOf<T>>>,
+	ExtendedBalance: From<InnerOf<OnChainAccuracyOf<T>>>,
+{
+	Box::new(move |who| {
+		snapshot
+			.iter()
+			.position(|(x, _, _)| x == who)
+			.and_then(|i| <usize as TryInto<CompactVoterIndexOf<T>>>::try_into(i).ok())
+	})
+}
+
+/// Create a function the returns the index a targets in the snapshot.
+///
+/// The returning index type is the same as the one defined in [`T::CompactSolution::Target`].
+pub fn target_index_fn_linear<T: Config>(
+	snapshot: &Vec<T::AccountId>,
+) -> Box<dyn Fn(&T::AccountId) -> Option<CompactTargetIndexOf<T>> + '_>
+where
+	ExtendedBalance: From<InnerOf<CompactAccuracyOf<T>>>,
+	ExtendedBalance: From<InnerOf<OnChainAccuracyOf<T>>>,
+{
+	Box::new(move |who| {
+		snapshot
+			.iter()
+			.position(|x| x == who)
+			.and_then(|i| <usize as TryInto<CompactTargetIndexOf<T>>>::try_into(i).ok())
+	})
+}
+
+/// Create a function that can map a voter index ([`CompactVoterIndexOf`]) to the actual voter
+/// account using a linearly indexible snapshot.
+pub fn voter_at_fn<T: Config>(
+	snapshot: &Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)>,
+) -> Box<dyn Fn(CompactVoterIndexOf<T>) -> Option<T::AccountId> + '_>
+where
+	ExtendedBalance: From<InnerOf<CompactAccuracyOf<T>>>,
+	ExtendedBalance: From<InnerOf<OnChainAccuracyOf<T>>>,
+{
+	Box::new(move |i| {
+		<CompactVoterIndexOf<T> as TryInto<usize>>::try_into(i)
+			.ok()
+			.and_then(|i| snapshot.get(i).map(|(x, _, _)| x).cloned())
+	})
+}
+
+/// Create a function that can map a target index ([`CompactTargetIndexOf`]) to the actual target
+/// account using a linearly indexible snapshot.
+pub fn target_at_fn<T: Config>(
+	snapshot: &Vec<T::AccountId>,
+) -> Box<dyn Fn(CompactTargetIndexOf<T>) -> Option<T::AccountId> + '_>
+where
+	ExtendedBalance: From<InnerOf<CompactAccuracyOf<T>>>,
+	ExtendedBalance: From<InnerOf<OnChainAccuracyOf<T>>>,
+{
+	Box::new(move |i| {
+		<CompactTargetIndexOf<T> as TryInto<usize>>::try_into(i)
+			.ok()
+			.and_then(|i| snapshot.get(i).cloned())
+	})
+}
+
+/// Create a function to get the stake of a voter.
+///
+/// This is not optimized and uses a linear search.
+pub fn stake_of_fn_linear<T: Config>(
+	snapshot: &Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)>,
+) -> Box<dyn Fn(&T::AccountId) -> VoteWeight + '_>
+where
+	ExtendedBalance: From<InnerOf<CompactAccuracyOf<T>>>,
+	ExtendedBalance: From<InnerOf<OnChainAccuracyOf<T>>>,
+{
+	Box::new(move |who| {
+		snapshot
+			.iter()
+			.find(|(x, _, _)| x == who)
+			.map(|(_, x, _)| *x)
+			.unwrap_or_default()
+	})
+}
+
+/// Create a function to get the stake of a voter.
+///
+/// ## Warning
+///
+/// The cache need must be derived from the same snapshot. Zero is returned if a voter is
+/// non-existent.
+pub fn stake_of_fn<'a, T: Config>(
+	snapshot: &'a Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)>,
+	cache: &'a BTreeMap<T::AccountId, usize>,
+) -> Box<dyn Fn(&T::AccountId) -> VoteWeight + 'a>
+where
+	ExtendedBalance: From<InnerOf<CompactAccuracyOf<T>>>,
+	ExtendedBalance: From<InnerOf<OnChainAccuracyOf<T>>>,
+{
+	Box::new(move |who| {
+		if let Some(index) = cache.get(who) {
+			snapshot
+				.get(*index)
+				.map(|(_, x, _)| x)
+				.cloned()
+				.unwrap_or_default()
+		} else {
+			0
+		}
+	})
+}

frame/two-phase-election-provider/src/lib.rs

@@ -202,7 +202,7 @@ mod benchmarking;
 #[cfg(test)]
 mod mock;
 #[macro_use]
-pub(crate) mod macros;
+pub mod helpers;
 
 const LOG_TARGET: &'static str = "election-provider";
 
@@ -634,8 +634,15 @@ pub mod pallet {
 		}
 
 		fn integrity_test() {
-			let max_vote: usize = <CompactOf<T> as CompactSolution>::LIMIT;
+			use sp_std::mem::size_of;
+			// The index type of both voters and targets need to be smaller than that of usize (very
+			// unlikely to be the case, but anyhow).
+			assert!(size_of::<CompactVoterIndexOf<T>>() <= size_of::<usize>());
+			assert!(size_of::<CompactTargetIndexOf<T>>() <= size_of::<usize>());
+
+			// ----------------------------
 			// based on the requirements of [`sp_npos_elections::Assignment::try_normalize`].
+			let max_vote: usize = <CompactOf<T> as CompactSolution>::LIMIT;
 
 			// 1. Maximum sum of [ChainAccuracy; 16] must fit into `UpperOf<ChainAccuracy>`..
 			let maximum_chain_accuracy: Vec<UpperOf<OnChainAccuracyOf<T>>> = (0..max_vote)
@@ -1064,9 +1071,10 @@ where
 		} = Self::snapshot().ok_or(FeasibilityError::SnapshotUnavailable)?;
 
 		// ----- Start building. First, we need some closures.
-		crate::voter_at_fn!(let voter_at , snapshot_voters, T);
-		crate::target_at_fn!(let target_at, snapshot_targets, T);
-		crate::voter_index_fn_usize!(let voter_index, snapshot_voters, T);
+		let cache = helpers::generate_voter_cache::<T>(&snapshot_voters);
+		let voter_at = helpers::voter_at_fn::<T>(&snapshot_voters);
+		let target_at = helpers::target_at_fn::<T>(&snapshot_targets);
+		let voter_index = helpers::voter_index_fn_usize::<T>(&cache);
 
 		// first, make sure that all the winners are sane.
 		let winners = winners
@@ -1106,8 +1114,8 @@ where
 			})
 			.collect::<Result<(), FeasibilityError>>()?;
 
-		// ----- Start building support. First, we need some more closures.
-		stake_of_fn!(let stake_of, snapshot_voters, T);
+		// ----- Start building support. First, we need one more closure.
+		let stake_of = helpers::stake_of_fn::<T>(&snapshot_voters, &cache);
 
 		// This might fail if the normalization fails. Very unlikely. See `integrity_test`.
 		let staked_assignments = assignment_ratio_to_staked_normalized(assignments, stake_of)