Trim compact solution for length during preparation (#8317)

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>
Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

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

@@ -46,6 +46,8 @@ pub enum MinerError {
 	PreDispatchChecksFailed,
 	/// The solution generated from the miner is not feasible.
 	Feasibility(FeasibilityError),
+	/// There are no more voters to remove to trim the solution.
+	NoMoreVoters,
 }
 
 impl From<sp_npos_elections::Error> for MinerError {
@@ -177,8 +179,13 @@ impl<T: Config> Pallet<T> {
 			maximum_allowed_voters,
 		);
 
-		// trim weight.
-		let compact = Self::trim_compact(maximum_allowed_voters, compact, &voter_index)?;
+		// trim length and weight
+		let compact = Self::trim_compact_weight(maximum_allowed_voters, compact, &voter_index)?;
+		let compact = Self::trim_compact_length(
+			T::MinerMaxLength::get(),
+			compact,
+			&voter_index,
+		)?;
 
 		// re-calc score.
 		let winners = sp_npos_elections::to_without_backing(winners);
@@ -221,7 +228,7 @@ impl<T: Config> Pallet<T> {
 	///
 	/// Indeed, the score must be computed **after** this step. If this step reduces the score too
 	/// much or remove a winner, then the solution must be discarded **after** this step.
-	pub fn trim_compact<FN>(
+	pub fn trim_compact_weight<FN>(
 		maximum_allowed_voters: u32,
 		mut compact: CompactOf<T>,
 		voter_index: FN,
@@ -267,6 +274,50 @@ impl<T: Config> Pallet<T> {
 		}
 	}
 
+	/// Greedily reduce the size of the solution to fit into the block w.r.t length.
+	///
+	/// The length of the solution is largely a function of the number of voters. The number of
+	/// winners cannot be changed. Thus, to reduce the solution size, we need to strip voters.
+	///
+	/// Note that this solution is already computed, and winners are elected based on the merit of
+	/// the total stake in the system. Nevertheless, some of the voters may be removed here.
+	///
+	/// Sometimes, removing a voter can cause a validator to also be implicitly removed, if
+	/// that voter was the only backer of that winner. In such cases, this solution is invalid, which
+	/// will be caught prior to submission.
+	///
+	/// The score must be computed **after** this step. If this step reduces the score too much,
+	/// then the solution must be discarded.
+	pub fn trim_compact_length(
+		max_allowed_length: u32,
+		mut compact: CompactOf<T>,
+		voter_index: impl Fn(&T::AccountId) -> Option<CompactVoterIndexOf<T>>,
+	) -> Result<CompactOf<T>, MinerError> {
+		// short-circuit to avoid getting the voters if possible
+		// this involves a redundant encoding, but that should hopefully be relatively cheap
+		if (compact.encoded_size().saturated_into::<u32>()) <= max_allowed_length {
+			return Ok(compact);
+		}
+
+		// grab all voters and sort them by least stake.
+		let RoundSnapshot { voters, .. } =
+			Self::snapshot().ok_or(MinerError::SnapshotUnAvailable)?;
+		let mut voters_sorted = voters
+			.into_iter()
+			.map(|(who, stake, _)| (who.clone(), stake))
+			.collect::<Vec<_>>();
+		voters_sorted.sort_by_key(|(_, y)| *y);
+		voters_sorted.reverse();
+
+		while compact.encoded_size() > max_allowed_length.saturated_into() {
+			let (smallest_stake_voter, _) = voters_sorted.pop().ok_or(MinerError::NoMoreVoters)?;
+			let index = voter_index(&smallest_stake_voter).ok_or(MinerError::SnapshotUnAvailable)?;
+			compact.remove_voter(index);
+		}
+
+		Ok(compact)
+	}
+
 	/// Find the maximum `len` that a compact can have in order to fit into the block weight.
 	///
 	/// This only returns a value between zero and `size.nominators`.
@@ -506,6 +557,7 @@ mod tests {
 		Call, *,
 	};
 	use frame_support::{dispatch::Dispatchable, traits::OffchainWorker};
+	use helpers::voter_index_fn_linear;
 	use mock::Call as OuterCall;
 	use frame_election_provider_support::Assignment;
 	use sp_runtime::{traits::ValidateUnsigned, PerU16};
@@ -889,4 +941,116 @@ mod tests {
 			assert!(matches!(call, OuterCall::MultiPhase(Call::submit_unsigned(_, _))));
 		})
 	}
+
+	#[test]
+	fn trim_compact_length_does_not_modify_when_short_enough() {
+		let mut ext = ExtBuilder::default().build();
+		ext.execute_with(|| {
+			roll_to(25);
+
+			// given
+			let RoundSnapshot { voters, ..} = MultiPhase::snapshot().unwrap();
+			let RawSolution { mut compact, .. } = raw_solution();
+			let encoded_len = compact.encode().len() as u32;
+			let compact_clone = compact.clone();
+
+			// when
+			assert!(encoded_len < <Runtime as Config>::MinerMaxLength::get());
+
+			// then
+			compact = MultiPhase::trim_compact_length(
+				encoded_len,
+				compact,
+				voter_index_fn_linear::<Runtime>(&voters),
+			).unwrap();
+			assert_eq!(compact, compact_clone);
+		});
+	}
+
+	#[test]
+	fn trim_compact_length_modifies_when_too_long() {
+		let mut ext = ExtBuilder::default().build();
+		ext.execute_with(|| {
+			roll_to(25);
+
+			let RoundSnapshot { voters, ..} =
+				MultiPhase::snapshot().unwrap();
+
+			let RawSolution { mut compact, .. } = raw_solution();
+			let encoded_len = compact.encoded_size() as u32;
+			let compact_clone = compact.clone();
+
+			compact = MultiPhase::trim_compact_length(
+				encoded_len - 1,
+				compact,
+				voter_index_fn_linear::<Runtime>(&voters),
+			).unwrap();
+
+			assert_ne!(compact, compact_clone);
+			assert!((compact.encoded_size() as u32) < encoded_len);
+		});
+	}
+
+	#[test]
+	fn trim_compact_length_trims_lowest_stake() {
+		let mut ext = ExtBuilder::default().build();
+		ext.execute_with(|| {
+			roll_to(25);
+
+			let RoundSnapshot { voters, ..} =
+				MultiPhase::snapshot().unwrap();
+
+			let RawSolution { mut compact, .. } = raw_solution();
+			let encoded_len = compact.encoded_size() as u32;
+			let voter_count = compact.voter_count();
+			let min_stake_voter = voters.iter()
+				.map(|(id, weight, _)| (weight, id))
+				.min()
+				.map(|(_, id)| id)
+				.unwrap();
+
+
+			compact = MultiPhase::trim_compact_length(
+				encoded_len - 1,
+				compact,
+				voter_index_fn_linear::<Runtime>(&voters),
+			).unwrap();
+
+			assert_eq!(compact.voter_count(), voter_count - 1, "we must have removed exactly 1 voter");
+
+			let assignments = compact.into_assignment(
+				|voter| Some(voter as AccountId),
+				|target| Some(target as AccountId),
+			).unwrap();
+			assert!(
+				assignments.iter()
+					.all(|Assignment{ who, ..}| who != min_stake_voter),
+				"min_stake_voter must no longer be in the set of voters",
+			);
+		});
+	}
+
+	// all the other solution-generation functions end up delegating to `mine_solution`, so if we
+	// demonstrate that `mine_solution` solutions are all trimmed to an acceptable length, then
+	// we know that higher-level functions will all also have short-enough solutions.
+	#[test]
+	fn mine_solution_solutions_always_within_acceptable_length() {
+		let mut ext = ExtBuilder::default().build();
+		ext.execute_with(|| {
+			roll_to(25);
+
+			// how long would the default solution be?
+			let solution = MultiPhase::mine_solution(0).unwrap();
+			let max_length = <Runtime as Config>::MinerMaxLength::get();
+			let solution_size = solution.0.compact.encoded_size();
+			assert!(solution_size <= max_length as usize);
+
+			// now set the max size to less than the actual size and regenerate
+			<Runtime as Config>::MinerMaxLength::set(solution_size as u32 - 1);
+			let solution = MultiPhase::mine_solution(0).unwrap();
+			let max_length = <Runtime as Config>::MinerMaxLength::get();
+			let solution_size = solution.0.compact.encoded_size();
+			assert!(solution_size <= max_length as usize);
+		});
+	}
 }