Fast CompactAssignment search (#8385)

* use more efficient search through candidates in offchain-election

* mark linear accessors as test-only in election-provider-multi-phase

This prevents production code which uses them from compiling.

Also write an efficient helper for getting the target index.

* doc grammar

* use faster target_index_fn in benchmarks

* unbox helper functions

* remove unnecessary import

* write lifetime after primary trait

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

substrate/frame/election-provider-multi-phase/src/benchmarking.rs

@@ -115,7 +115,7 @@ fn solution_with_size<T: Config>(
 	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 target_index = helpers::target_index_fn::<T>(&targets);
 	let voter_at = helpers::voter_at_fn::<T>(&all_voters);
 	let target_at = helpers::target_at_fn::<T>(&targets);
 

substrate/frame/election-provider-multi-phase/src/helpers.rs

@@ -18,7 +18,7 @@
 //! Some helper functions/macros for this crate.
 
 use super::{Config, VoteWeight, CompactVoterIndexOf, CompactTargetIndexOf};
-use sp_std::{collections::btree_map::BTreeMap, convert::TryInto, boxed::Box, prelude::*};
+use sp_std::{collections::btree_map::BTreeMap, convert::TryInto, prelude::*};
 
 #[macro_export]
 macro_rules! log {
@@ -56,10 +56,10 @@ pub fn generate_voter_cache<T: Config>(
 /// 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>> + '_> {
-	Box::new(move |who| {
+) -> impl Fn(&T::AccountId) -> Option<CompactVoterIndexOf<T>> + '_ {
+	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.
@@ -69,8 +69,8 @@ pub fn voter_index_fn<T: Config>(
 /// 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> + '_> {
-	Box::new(move |who| cache.get(who).cloned())
+) -> impl Fn(&T::AccountId) -> Option<usize> + '_ {
+	move |who| cache.get(who).cloned()
 }
 
 /// A non-optimized, linear version of [`voter_index_fn`] that does not need a cache and does a
@@ -79,64 +79,90 @@ pub fn voter_index_fn_usize<T: Config>(
 /// ## Warning
 ///
 /// Not meant to be used in production.
+#[cfg(test)]
 pub fn voter_index_fn_linear<T: Config>(
 	snapshot: &Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)>,
-) -> Box<dyn Fn(&T::AccountId) -> Option<CompactVoterIndexOf<T>> + '_> {
-	Box::new(move |who| {
+) -> impl Fn(&T::AccountId) -> Option<CompactVoterIndexOf<T>> + '_ {
+	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.
+/// Create a function the returns the index to a target in the snapshot.
 ///
-/// The returning index type is the same as the one defined in `T::CompactSolution::Target`.
+/// The returned index type is the same as the one defined in `T::CompactSolution::Target`.
+///
+/// Note: to the extent possible, the returned function should be cached and reused. Producing that
+/// function requires a `O(n log n)` data transform. Each invocation of that function completes
+/// in `O(log n)`.
+pub fn target_index_fn<T: Config>(
+	snapshot: &Vec<T::AccountId>,
+) -> impl Fn(&T::AccountId) -> Option<CompactTargetIndexOf<T>> + '_ {
+	let cache: BTreeMap<_, _> =
+		snapshot.iter().enumerate().map(|(idx, account_id)| (account_id, idx)).collect();
+	move |who| {
+		cache
+			.get(who)
+			.and_then(|i| <usize as TryInto<CompactTargetIndexOf<T>>>::try_into(*i).ok())
+	}
+}
+
+/// Create a function the returns the index to a target in the snapshot.
+///
+/// The returned index type is the same as the one defined in `T::CompactSolution::Target`.
+///
+/// ## Warning
+///
+/// Not meant to be used in production.
+#[cfg(test)]
 pub fn target_index_fn_linear<T: Config>(
 	snapshot: &Vec<T::AccountId>,
-) -> Box<dyn Fn(&T::AccountId) -> Option<CompactTargetIndexOf<T>> + '_> {
-	Box::new(move |who| {
+) -> impl Fn(&T::AccountId) -> Option<CompactTargetIndexOf<T>> + '_ {
+	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> + '_> {
-	Box::new(move |i| {
+) -> impl Fn(CompactVoterIndexOf<T>) -> Option<T::AccountId> + '_ {
+	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> + '_> {
-	Box::new(move |i| {
+) -> impl Fn(CompactTargetIndexOf<T>) -> Option<T::AccountId> + '_ {
+	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.
+#[cfg(test)]
 pub fn stake_of_fn_linear<T: Config>(
 	snapshot: &Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)>,
-) -> Box<dyn Fn(&T::AccountId) -> VoteWeight + '_> {
-	Box::new(move |who| {
+) -> impl Fn(&T::AccountId) -> VoteWeight + '_ {
+	move |who| {
 		snapshot.iter().find(|(x, _, _)| x == who).map(|(_, x, _)| *x).unwrap_or_default()
-	})
+	}
 }
 
 /// Create a function to get the stake of a voter.
@@ -148,12 +174,12 @@ pub fn stake_of_fn_linear<T: Config>(
 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> {
-	Box::new(move |who| {
+) -> impl Fn(&T::AccountId) -> VoteWeight + 'a {
+	move |who| {
 		if let Some(index) = cache.get(who) {
 			snapshot.get(*index).map(|(_, x, _)| x).cloned().unwrap_or_default()
 		} else {
 			0
 		}
-	})
+	}
 }

substrate/frame/election-provider-multi-phase/src/unsigned.rs

@@ -145,7 +145,7 @@ impl<T: Config> Pallet<T> {
 		// closures.
 		let cache = helpers::generate_voter_cache::<T>(&voters);
 		let voter_index = helpers::voter_index_fn::<T>(&cache);
-		let target_index = helpers::target_index_fn_linear::<T>(&targets);
+		let target_index = helpers::target_index_fn::<T>(&targets);
 		let voter_at = helpers::voter_at_fn::<T>(&voters);
 		let target_at = helpers::target_at_fn::<T>(&targets);
 		let stake_of = helpers::stake_of_fn::<T>(&voters, &cache);