Refactor election solution trimming for efficiency

frame/election-provider-multi-phase/Cargo.toml

@@ -20,7 +20,7 @@ log = { version = "0.4.14", default-features = false }
 frame-support = { version = "3.0.0", default-features = false, path = "../support" }
 frame-system = { version = "3.0.0", default-features = false, path = "../system" }
 
-sp-io ={ version = "3.0.0", default-features = false, path = "../../primitives/io" }
+sp-io = { version = "3.0.0", default-features = false, path = "../../primitives/io" }
 sp-std = { version = "3.0.0", default-features = false, path = "../../primitives/std" }
 sp-runtime = { version = "3.0.0", default-features = false, path = "../../primitives/runtime" }
 sp-npos-elections = { version = "3.0.0", default-features = false, path = "../../primitives/npos-elections" }
@@ -37,8 +37,8 @@ parking_lot = "0.11.0"
 rand = { version = "0.7.3" }
 hex-literal = "0.3.1"
 substrate-test-utils = { version = "3.0.0", path = "../../test-utils" }
+sp-core = { version = "3.0.0", default-features = false, path = "../../primitives/core" }
 sp-io = { version = "3.0.0", path = "../../primitives/io" }
-sp-core = { version = "3.0.0", path = "../../primitives/core" }
 sp-tracing = { version = "3.0.0", path = "../../primitives/tracing" }
 frame-election-provider-support = { version = "3.0.0", features = ["runtime-benchmarks"], path = "../election-provider-support" }
 pallet-balances = { version = "3.0.0", path = "../balances" }

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

@@ -135,7 +135,7 @@ fn solution_with_size<T: Config>(
 		.collect::<Vec<_>>();
 
 	let compact =
-		<CompactOf<T>>::from_assignment(assignments, &voter_index, &target_index).unwrap();
+		<CompactOf<T>>::from_assignment(&assignments, &voter_index, &target_index).unwrap();
 	let score = compact.clone().score(&winners, stake_of, voter_at, target_at).unwrap();
 	let round = <MultiPhase<T>>::round();
 

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

@@ -62,6 +62,18 @@ pub fn voter_index_fn<T: Config>(
 	}
 }
 
+/// Create a function that returns the index of a voter in the snapshot.
+///
+/// Same as [`voter_index_fn`] but the returned function owns all its necessary data; nothing is
+/// borrowed.
+pub fn voter_index_fn_owned<T: Config>(
+	cache: BTreeMap<T::AccountId, usize>,
+) -> 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.
 ///
 /// ## Warning

frame/election-provider-multi-phase/src/mock.rs

@@ -17,6 +17,7 @@
 
 use super::*;
 use crate as multi_phase;
+use multi_phase::unsigned::{IndexAssignmentOf, Voter};
 pub use frame_support::{assert_noop, assert_ok};
 use frame_support::{
 	parameter_types,
@@ -41,7 +42,7 @@ use sp_runtime::{
 	traits::{BlakeTwo256, IdentityLookup},
 	PerU16,
 };
-use std::sync::Arc;
+use std::{convert::TryFrom, sync::Arc};
 
 pub type Block = sp_runtime::generic::Block<Header, UncheckedExtrinsic>;
 pub type UncheckedExtrinsic = sp_runtime::generic::UncheckedExtrinsic<AccountId, Call, (), ()>;
@@ -95,19 +96,70 @@ pub fn roll_to_with_ocw(n: u64) {
 	}
 }
 
+pub struct TrimHelpers {
+	pub voters: Vec<Voter<Runtime>>,
+	pub assignments: Vec<IndexAssignmentOf<Runtime>>,
+	pub encoded_size_of:
+		Box<dyn Fn(&[IndexAssignmentOf<Runtime>]) -> Result<usize, sp_npos_elections::Error>>,
+	pub voter_index: Box<
+		dyn Fn(
+			&<Runtime as frame_system::Config>::AccountId,
+		) -> Option<CompactVoterIndexOf<Runtime>>,
+	>,
+}
+
+/// Helpers for setting up trimming tests.
+///
+/// Assignments are pre-sorted in reverse order of stake.
+pub fn trim_helpers() -> TrimHelpers {
+	let RoundSnapshot { voters, targets } = MultiPhase::snapshot().unwrap();
+	let stakes: std::collections::HashMap<_, _> =
+		voters.iter().map(|(id, stake, _)| (*id, *stake)).collect();
+
+	// Compute the size of a compact solution comprised of the selected arguments.
+	//
+	// This function completes in `O(edges)`; it's expensive, but linear.
+	let encoded_size_of = Box::new(|assignments: &[IndexAssignmentOf<Runtime>]| {
+		CompactOf::<Runtime>::try_from(assignments).map(|compact| compact.encoded_size())
+	});
+	let cache = helpers::generate_voter_cache::<Runtime>(&voters);
+	let voter_index = helpers::voter_index_fn_owned::<Runtime>(cache);
+	let target_index = helpers::target_index_fn::<Runtime>(&targets);
+
+	let desired_targets = MultiPhase::desired_targets().unwrap();
+
+	let ElectionResult { mut assignments, .. } = seq_phragmen::<_, CompactAccuracyOf<Runtime>>(
+		desired_targets as usize,
+		targets.clone(),
+		voters.clone(),
+		None,
+	)
+	.unwrap();
+
+	// sort by decreasing order of stake
+	assignments.sort_unstable_by_key(|assignment| {
+		std::cmp::Reverse(stakes.get(&assignment.who).cloned().unwrap_or_default())
+	});
+
+	// convert to IndexAssignment
+	let assignments = assignments
+		.iter()
+		.map(|assignment| {
+			IndexAssignmentOf::<Runtime>::new(assignment, &voter_index, &target_index)
+		})
+		.collect::<Result<Vec<_>, _>>()
+		.expect("test assignments don't contain any voters with too many votes");
+
+	TrimHelpers { voters, assignments, encoded_size_of, voter_index: Box::new(voter_index) }
+}
+
 /// Spit out a verifiable raw solution.
 ///
 /// This is a good example of what an offchain miner would do.
 pub fn raw_solution() -> RawSolution<CompactOf<Runtime>> {
 	let RoundSnapshot { voters, targets } = MultiPhase::snapshot().unwrap();
 	let desired_targets = MultiPhase::desired_targets().unwrap();
 
-	// closures
-	let cache = helpers::generate_voter_cache::<Runtime>(&voters);
-	let voter_index = helpers::voter_index_fn_linear::<Runtime>(&voters);
-	let target_index = helpers::target_index_fn_linear::<Runtime>(&targets);
-	let stake_of = helpers::stake_of_fn::<Runtime>(&voters, &cache);
-
 	let ElectionResult { winners, assignments } = seq_phragmen::<_, CompactAccuracyOf<Runtime>>(
 		desired_targets as usize,
 		targets.clone(),
@@ -116,14 +168,20 @@ pub fn raw_solution() -> RawSolution<CompactOf<Runtime>> {
 	)
 	.unwrap();
 
+	// closures
+	let cache = helpers::generate_voter_cache::<Runtime>(&voters);
+	let voter_index = helpers::voter_index_fn_linear::<Runtime>(&voters);
+	let target_index = helpers::target_index_fn_linear::<Runtime>(&targets);
+	let stake_of = helpers::stake_of_fn::<Runtime>(&voters, &cache);
+
 	let winners = to_without_backing(winners);
 
 	let score = {
 		let staked = assignment_ratio_to_staked_normalized(assignments.clone(), &stake_of).unwrap();
 		to_supports(&winners, &staked).unwrap().evaluate()
 	};
 	let compact =
-		<CompactOf<Runtime>>::from_assignment(assignments, &voter_index, &target_index).unwrap();
+		<CompactOf<Runtime>>::from_assignment(&assignments, &voter_index, &target_index).unwrap();
 
 	let round = MultiPhase::round();
 	RawSolution { compact, score, round }