Refactor election solution trimming for efficiency

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

@@ -39,6 +39,7 @@ 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-npos-elections = { version = "3.0.0", default-features = false, features = [ "mocks" ], path = "../../primitives/npos-elections" }
 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" }
@@ -64,5 +65,6 @@ std = [
 runtime-benchmarks = [
 	"frame-benchmarking",
 	"rand",
+	"sp-npos-elections/mocks",
 ]
 try-runtime = ["frame-support/try-runtime"]

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

@@ -18,15 +18,16 @@
 //! Two phase election pallet benchmarking.
 
 use super::*;
-use crate::Pallet as MultiPhase;
+use crate::{Pallet as MultiPhase, unsigned::IndexAssignmentOf};
 use frame_benchmarking::impl_benchmark_test_suite;
 use frame_support::{assert_ok, traits::OnInitialize};
 use frame_system::RawOrigin;
 use rand::{prelude::SliceRandom, rngs::SmallRng, SeedableRng};
 use frame_election_provider_support::Assignment;
-use sp_arithmetic::traits::One;
+use sp_arithmetic::{per_things::Percent, traits::One};
+use sp_npos_elections::IndexAssignment;
 use sp_runtime::InnerOf;
-use sp_std::convert::TryInto;
+use sp_std::convert::{TryFrom, TryInto};
 
 const SEED: u32 = 999;
 
@@ -254,6 +255,69 @@ frame_benchmarking::benchmarks! {
 		assert!(<MultiPhase<T>>::queued_solution().is_some());
 	}
 
+	#[extra]
+	trim_assignments_length {
+		// number of votes in snapshot.
+		let v in (T::BenchmarkingConfig::VOTERS[0]) .. T::BenchmarkingConfig::VOTERS[1];
+		// number of targets in snapshot.
+		let t in (T::BenchmarkingConfig::TARGETS[0]) .. T::BenchmarkingConfig::TARGETS[1];
+		// number of assignments, i.e. compact.len(). This means the active nominators, thus must be
+		// a subset of `v` component.
+		let a in (T::BenchmarkingConfig::ACTIVE_VOTERS[0]) .. T::BenchmarkingConfig::ACTIVE_VOTERS[1];
+		// number of desired targets. Must be a subset of `t` component.
+		let d in (T::BenchmarkingConfig::DESIRED_TARGETS[0]) .. T::BenchmarkingConfig::DESIRED_TARGETS[1];
+		// Subtract this percentage from the actual encoded size
+		let f in 0 .. 95;
+
+		// Compute a random solution, then work backwards to get the lists of voters, targets, and assignments
+		let witness = SolutionOrSnapshotSize { voters: v, targets: t };
+		let RawSolution { compact, .. } = solution_with_size::<T>(witness, a, d);
+		let RoundSnapshot { voters, targets } = MultiPhase::<T>::snapshot().unwrap();
+		let voter_at = helpers::voter_at_fn::<T>(&voters);
+		let target_at = helpers::target_at_fn::<T>(&targets);
+		let mut assignments = compact.into_assignment(voter_at, target_at).unwrap();
+
+		// make a voter cache and some helper functions for access
+		let cache = helpers::generate_voter_cache::<T>(&voters);
+		let voter_index = helpers::voter_index_fn::<T>(&cache);
+		let target_index = helpers::target_index_fn::<T>(&targets);
+
+		// sort assignments by decreasing voter stake
+		assignments.sort_by_key(|crate::unsigned::Assignment::<T> { who, .. }| {
+			let stake = cache.get(&who).map(|idx| {
+				let (_, stake, _) = voters[*idx];
+				stake
+			}).unwrap_or_default();
+			sp_std::cmp::Reverse(stake)
+		});
+
+		let mut index_assignments = assignments
+			.into_iter()
+			.map(|assignment| IndexAssignment::new(&assignment, &voter_index, &target_index))
+			.collect::<Result<Vec<_>, _>>()
+			.unwrap();
+
+		let encoded_size_of = |assignments: &[IndexAssignmentOf<T>]| {
+			CompactOf::<T>::try_from(assignments).map(|compact| compact.encoded_size())
+		};
+
+		let desired_size = Percent::from_percent(100 - f.saturated_into::<u8>())
+			.mul_ceil(encoded_size_of(index_assignments.as_slice()).unwrap());
+		log!(trace, "desired_size = {}", desired_size);
+	}: {
+		MultiPhase::<T>::trim_assignments_length(
+			desired_size.saturated_into(),
+			&mut index_assignments,
+			&encoded_size_of,
+		).unwrap();
+	} verify {
+		let compact = CompactOf::<T>::try_from(index_assignments.as_slice()).unwrap();
+		let encoding = compact.encode();
+		log!(trace, "encoded size prediction = {}", encoded_size_of(index_assignments.as_slice()).unwrap());
+		log!(trace, "actual encoded size = {}", encoding.len());
+		assert!(encoding.len() <= desired_size);
+	}
+
 	// This is checking a valid solution. The worse case is indeed a valid solution.
 	feasibility_check {
 		// number of votes in snapshot.

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

@@ -189,16 +189,21 @@ impl<T: Config> Pallet<T> {
 
 			// Sort the assignments by reversed voter stake. This ensures that we can efficiently
 			// truncate the list.
-			staked.sort_by_key(|sp_npos_elections::StakedAssignment::<T::AccountId> { who, .. }| {
-				// though staked assignments are expressed in terms of absolute stake, we'd still
-				// need to iterate over all votes in order to actually compute the total stake.
-				// it should be faster to look it up from the cache.
-				let stake = cache.get(who).map(|idx| {
-					let (_, stake, _) = voters[*idx];
-					stake
-				}).unwrap_or_default();
-				sp_std::cmp::Reverse(stake)
-			});
+			staked.sort_by_key(
+				|sp_npos_elections::StakedAssignment::<T::AccountId> { who, .. }| {
+					// though staked assignments are expressed in terms of absolute stake, we'd
+					// still need to iterate over all votes in order to actually compute the total
+					// stake. it should be faster to look it up from the cache.
+					let stake = cache
+						.get(who)
+						.map(|idx| {
+							let (_, stake, _) = voters[*idx];
+							stake
+						})
+						.unwrap_or_default();
+					sp_std::cmp::Reverse(stake)
+				},
+			);
 
 			// convert back.
 			assignment_staked_to_ratio_normalized(staked)?
@@ -303,7 +308,7 @@ impl<T: Config> Pallet<T> {
 	///
 	/// The score must be computed **after** this step. If this step reduces the score too much,
 	/// then the solution must be discarded.
-	fn trim_assignments_length(
+	pub(crate) fn trim_assignments_length(
 		max_allowed_length: u32,
 		assignments: &mut Vec<IndexAssignmentOf<T>>,
 		encoded_size_of: impl Fn(&[IndexAssignmentOf<T>]) -> Result<usize, sp_npos_elections::Error>,
@@ -313,23 +318,21 @@ impl<T: Config> Pallet<T> {
 		let max_allowed_length: usize = max_allowed_length.saturated_into();
 		let mut high = assignments.len();
 		let mut low = 0;
-		let mut maximum_allowed_voters = 0;
-		let mut size;
 
 		while high - low > 1 {
-			maximum_allowed_voters = (high + low) / 2;
-			size = encoded_size_of(&assignments[..maximum_allowed_voters])?;
-			if size > max_allowed_length {
-				high = maximum_allowed_voters;
+			let test = (high + low) / 2;
+			if encoded_size_of(&assignments[..test])? <= max_allowed_length {
+				low = test;
 			} else {
-				low = maximum_allowed_voters;
+				high = test;
 			}
 		}
-		if high - low == 1
-			&& encoded_size_of(&assignments[..maximum_allowed_voters + 1])? <= max_allowed_length
-		{
-			maximum_allowed_voters += 1;
-		}
+		let maximum_allowed_voters =
+			if encoded_size_of(&assignments[..low + 1])? <= max_allowed_length {
+				low + 1
+			} else {
+				low
+			};
 
 		// ensure our postconditions are correct
 		debug_assert!(

primitives/npos-elections/src/helpers.rs

@@ -72,10 +72,9 @@ pub fn assignment_staked_to_ratio_normalized<A: IdentifierT, P: PerThing128>(
 	staked: Vec<StakedAssignment<A>>,
 ) -> Result<Vec<Assignment<A, P>>, Error> {
 	let mut ratio = staked.into_iter().map(|a| a.into_assignment()).collect::<Vec<_>>();
-	ratio
-		.iter_mut()
-		.map(|a| a.try_normalize().map_err(|err| Error::ArithmeticError(err)))
-		.collect::<Result<_, _>>()?;
+	for assignment in ratio.iter_mut() {
+		assignment.try_normalize().map_err(|err| Error::ArithmeticError(err))?;
+	}
 	Ok(ratio)
 }