Clean Phragmén Equlise API (#5452)

* Clean phragmen API and equalise()

* Stabilize new api

* Fix phragmen fuzzers

* More fixes

* Make fuzzers reproducible

* improvements

* Make equalize update assignments as well.

* total function for staked_assignment.

* Fix fuzzer build

* remvoe TODO

* Fix a bunch more.

* clean stray debug stuff

* Update primitives/phragmen/src/lib.rs

Co-Authored-By: thiolliere <gui.thiolliere@gmail.com>

* fix range function

* fix number generator

Co-authored-by: thiolliere <gui.thiolliere@gmail.com>

substrate/Cargo.lock

@@ -7515,6 +7515,17 @@ dependencies = [
  "syn 1.0.17",
 ]
 
+[[package]]
+name = "sp-phragmen-fuzzer"
+version = "2.0.0-alpha.5"
+dependencies = [
+ "honggfuzz",
+ "rand 0.7.3",
+ "sp-phragmen",
+ "sp-runtime",
+ "sp-std",
+]
+
 [[package]]
 name = "sp-rpc"
 version = "2.0.0-dev"

substrate/Cargo.toml

@@ -134,6 +134,7 @@ members = [
 	"primitives/offchain",
 	"primitives/panic-handler",
 	"primitives/phragmen",
+	"primitives/phragmen/fuzzer",
 	"primitives/phragmen/compact",
 	"primitives/rpc",
 	"primitives/runtime-interface",

substrate/frame/elections-phragmen/src/lib.rs

@@ -94,7 +94,7 @@ use frame_support::{
 		ChangeMembers, OnUnbalanced, WithdrawReason, Contains, BalanceStatus, InitializeMembers,
 	}
 };
-use sp_phragmen::{build_support_map, ExtendedBalance};
+use sp_phragmen::{build_support_map, ExtendedBalance, VoteWeight, PhragmenResult};
 use frame_system::{self as system, ensure_signed, ensure_root};
 
 const MODULE_ID: LockIdentifier = *b"phrelect";
@@ -123,7 +123,7 @@ pub trait Trait: frame_system::Trait {
 
 	/// Convert a balance into a number used for election calculation.
 	/// This must fit into a `u64` but is allowed to be sensibly lossy.
-	type CurrencyToVote: Convert<BalanceOf<Self>, u64> + Convert<u128, BalanceOf<Self>>;
+	type CurrencyToVote: Convert<BalanceOf<Self>, VoteWeight> + Convert<ExtendedBalance, BalanceOf<Self>>;
 
 	/// How much should be locked up in order to submit one's candidacy.
 	type CandidacyBond: Get<BalanceOf<Self>>;
@@ -703,17 +703,28 @@ impl<T: Trait> Module<T> {
 		// previous runners_up are also always candidates for the next round.
 		candidates.append(&mut Self::runners_up_ids());
 
+		// helper closures to deal with balance/stake.
+		let to_votes = |b: BalanceOf<T>| -> VoteWeight {
+			<T::CurrencyToVote as Convert<BalanceOf<T>, VoteWeight>>::convert(b)
+		};
+		let to_balance = |e: ExtendedBalance| -> BalanceOf<T> {
+			<T::CurrencyToVote as Convert<ExtendedBalance, BalanceOf<T>>>::convert(e)
+		};
+		let stake_of = |who: &T::AccountId| -> VoteWeight {
+			to_votes(Self::locked_stake_of(who))
+		};
+
 		let voters_and_votes = Voting::<T>::iter()
-			.map(|(voter, (stake, targets))| { (voter, stake, targets) })
+			.map(|(voter, (stake, targets))| { (voter, to_votes(stake), targets) })
 			.collect::<Vec<_>>();
-		let maybe_phragmen_result = sp_phragmen::elect::<_, _, T::CurrencyToVote, Perbill>(
+		let maybe_phragmen_result = sp_phragmen::elect::<T::AccountId, Perbill>(
 			num_to_elect,
 			0,
 			candidates,
 			voters_and_votes.clone(),
 		);
 
-		if let Some(phragmen_result) = maybe_phragmen_result {
+		if let Some(PhragmenResult { winners, assignments }) = maybe_phragmen_result {
 			let old_members_ids = <Members<T>>::take().into_iter()
 				.map(|(m, _)| m)
 				.collect::<Vec<T::AccountId>>();
@@ -726,26 +737,19 @@ impl<T: Trait> Module<T> {
 			// vote are still considered by phragmen and when good candidates are scarce, then these
 			// cheap ones might get elected. We might actually want to remove the filter and allow
 			// zero-voted candidates to also make it to the membership set.
-			let new_set_with_approval = phragmen_result.winners;
+			let new_set_with_approval = winners;
 			let new_set = new_set_with_approval
 				.into_iter()
 				.filter_map(|(m, a)| if a.is_zero() { None } else { Some(m) } )
 				.collect::<Vec<T::AccountId>>();
 
-			let stake_of = |who: &T::AccountId| -> ExtendedBalance {
-				<T::CurrencyToVote as Convert<BalanceOf<T>, u64>>::convert(
-					Self::locked_stake_of(who)
-				) as ExtendedBalance
-			};
 			let staked_assignments = sp_phragmen::assignment_ratio_to_staked(
-				phragmen_result.assignments,
+				assignments,
 				stake_of,
 			);
 
 			let (support_map, _) = build_support_map::<T::AccountId>(&new_set, &staked_assignments);
 
-			let to_balance = |e: ExtendedBalance|
-				<T::CurrencyToVote as Convert<ExtendedBalance, BalanceOf<T>>>::convert(e);
 			let new_set_with_stake = new_set
 				.into_iter()
 				.map(|ref m| {
@@ -766,14 +770,14 @@ impl<T: Trait> Module<T> {
 			// save the members, sorted based on account id.
 			new_members.sort_by(|i, j| i.0.cmp(&j.0));
 
-			let mut prime_votes: Vec<_> = new_members.iter().map(|c| (&c.0, BalanceOf::<T>::zero())).collect();
+			let mut prime_votes: Vec<_> = new_members.iter().map(|c| (&c.0, VoteWeight::zero())).collect();
 			for (_, stake, targets) in voters_and_votes.into_iter() {
 				for (votes, who) in targets.iter()
 					.enumerate()
 					.map(|(votes, who)| ((MAXIMUM_VOTE - votes) as u32, who))
 				{
 					if let Ok(i) = prime_votes.binary_search_by_key(&who, |k| k.0) {
-						prime_votes[i].1 += stake * votes.into();
+						prime_votes[i].1 += stake * votes as VoteWeight;
 					}
 				}
 			}

substrate/frame/staking/src/lib.rs

@@ -308,7 +308,7 @@ use frame_system::{
 };
 use sp_phragmen::{
 	ExtendedBalance, Assignment, PhragmenScore, PhragmenResult, build_support_map, evaluate_support,
-	elect, generate_compact_solution_type, is_score_better, VotingLimit, SupportMap,
+	elect, generate_compact_solution_type, is_score_better, VotingLimit, SupportMap, VoteWeight,
 };
 
 const DEFAULT_MINIMUM_VALIDATOR_COUNT: u32 = 4;
@@ -737,12 +737,11 @@ pub trait Trait: frame_system::Trait {
 	/// is not used.
 	type UnixTime: UnixTime;
 
-	/// Convert a balance into a number used for election calculation.
+	/// Convert a balance into a number used for election calculation. This must fit into a `u64`
-	/// This must fit into a `u64` but is allowed to be sensibly lossy.
+	/// but is allowed to be sensibly lossy. The `u64` is used to communicate with the
-	/// TODO: #1377
+	/// [`sp_phragmen`] crate which accepts u64 numbers and does operations in 128. Consequently,
-	/// The backward convert should be removed as the new Phragmen API returns ratio.
+	/// the backward convert is used convert the u128s from phragmen back to a [`BalanceOf`].
-	/// The post-processing needs it but will be moved to off-chain. TODO: #2908
+	type CurrencyToVote: Convert<BalanceOf<Self>, VoteWeight> + Convert<u128, BalanceOf<Self>>;
-	type CurrencyToVote: Convert<BalanceOf<Self>, u64> + Convert<u128, BalanceOf<Self>>;
 
 	/// Tokens have been minted and are unused for validator-reward.
 	type RewardRemainder: OnUnbalanced<NegativeImbalanceOf<Self>>;
@@ -1853,7 +1852,7 @@ decl_module! {
 		/// - Memory: O(n + m) reads to map index to `AccountId` for un-compact.
 		///
 		/// - Storage: O(e) accountid reads from `Nomination` to read correct nominations.
-		/// - Storage: O(e) calls into `slashable_balance_of_extended` to convert ratio to staked.
+		/// - Storage: O(e) calls into `slashable_balance_of_vote_weight` to convert ratio to staked.
 		///
 		/// - Memory: build_support_map. O(e).
 		/// - Memory: evaluate_support: O(E).
@@ -1953,11 +1952,11 @@ impl<T: Trait> Module<T> {
 		Self::bonded(stash).and_then(Self::ledger).map(|l| l.active).unwrap_or_default()
 	}
 
-	/// internal impl of [`slashable_balance_of`] that returns [`ExtendedBalance`].
+	/// internal impl of [`slashable_balance_of`] that returns [`VoteWeight`].
-	fn slashable_balance_of_extended(stash: &T::AccountId) -> ExtendedBalance {
+	fn slashable_balance_of_vote_weight(stash: &T::AccountId) -> VoteWeight {
-		<T::CurrencyToVote as Convert<BalanceOf<T>, u64>>::convert(
+		<T::CurrencyToVote as Convert<BalanceOf<T>, VoteWeight>>::convert(
 			Self::slashable_balance_of(stash)
-		) as ExtendedBalance
+		)
 	}
 
 	/// Dump the list of validators and nominators into vectors and keep them on-chain.
@@ -2456,7 +2455,7 @@ impl<T: Trait> Module<T> {
 		// convert into staked assignments.
 		let staked_assignments = sp_phragmen::assignment_ratio_to_staked(
 			assignments,
-			Self::slashable_balance_of_extended,
+			Self::slashable_balance_of_vote_weight,
 		);
 
 		// build the support map thereof in order to evaluate.
@@ -2711,7 +2710,7 @@ impl<T: Trait> Module<T> {
 
 			let staked_assignments = sp_phragmen::assignment_ratio_to_staked(
 				assignments,
-				Self::slashable_balance_of_extended,
+				Self::slashable_balance_of_vote_weight,
 			);
 
 			let (supports, _) = build_support_map::<T::AccountId>(
@@ -2747,11 +2746,11 @@ impl<T: Trait> Module<T> {
 	///
 	/// No storage item is updated.
 	fn do_phragmen<Accuracy: PerThing>() -> Option<PhragmenResult<T::AccountId, Accuracy>> {
-		let mut all_nominators: Vec<(T::AccountId, BalanceOf<T>, Vec<T::AccountId>)> = Vec::new();
+		let mut all_nominators: Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)> = Vec::new();
 		let mut all_validators = Vec::new();
 		for (validator, _) in <Validators<T>>::iter() {
 			// append self vote
-			let self_vote = (validator.clone(), Self::slashable_balance_of(&validator), vec![validator.clone()]);
+			let self_vote = (validator.clone(), Self::slashable_balance_of_vote_weight(&validator), vec![validator.clone()]);
 			all_nominators.push(self_vote);
 			all_validators.push(validator);
 		}
@@ -2771,11 +2770,11 @@ impl<T: Trait> Module<T> {
 			(nominator, targets)
 		});
 		all_nominators.extend(nominator_votes.map(|(n, ns)| {
-			let s = Self::slashable_balance_of(&n);
+			let s = Self::slashable_balance_of_vote_weight(&n);
 			(n, s, ns)
 		}));
 
-		elect::<_, _, T::CurrencyToVote, Accuracy>(
+		elect::<_, Accuracy>(
 			Self::validator_count() as usize,
 			Self::minimum_validator_count().max(1) as usize,
 			all_validators,

substrate/frame/staking/src/mock.rs

@@ -33,6 +33,7 @@ use frame_system::offchain::TransactionSubmitter;
 use sp_io;
 use sp_phragmen::{
 	build_support_map, evaluate_support, reduce, ExtendedBalance, StakedAssignment, PhragmenScore,
+	VoteWeight,
 };
 use crate::*;
 
@@ -846,10 +847,10 @@ pub(crate) fn prepare_submission_with(
 	} = Staking::do_phragmen::<OffchainAccuracy>().unwrap();
 	let winners = winners.into_iter().map(|(w, _)| w).collect::<Vec<AccountId>>();
 
-	let stake_of = |who: &AccountId| -> ExtendedBalance {
+	let stake_of = |who: &AccountId| -> VoteWeight {
-		<CurrencyToVoteHandler as Convert<Balance, u64>>::convert(
+		<CurrencyToVoteHandler as Convert<Balance, VoteWeight>>::convert(
 			Staking::slashable_balance_of(&who)
-		) as ExtendedBalance
+		)
 	};
 	let mut staked = sp_phragmen::assignment_ratio_to_staked(assignments, stake_of);
 
@@ -888,7 +889,7 @@ pub(crate) fn prepare_submission_with(
 	let score = {
 		let staked = sp_phragmen::assignment_ratio_to_staked(
 			assignments_reduced.clone(),
-			Staking::slashable_balance_of_extended,
+			Staking::slashable_balance_of_vote_weight,
 		);
 
 		let (support_map, _) = build_support_map::<AccountId>(

substrate/frame/staking/src/offchain_election.rs

@@ -167,7 +167,7 @@ pub fn prepare_submission<T: Trait>(
 	// convert into absolute value and to obtain the reduced version.
 	let mut staked = sp_phragmen::assignment_ratio_to_staked(
 		assignments,
-		<Module<T>>::slashable_balance_of_extended,
+		<Module<T>>::slashable_balance_of_vote_weight,
 	);
 
 	if do_reduce {
@@ -188,7 +188,7 @@ pub fn prepare_submission<T: Trait>(
 	let score = {
 		let staked = sp_phragmen::assignment_ratio_to_staked(
 			low_accuracy_assignment.clone(),
-			<Module<T>>::slashable_balance_of_extended,
+			<Module<T>>::slashable_balance_of_vote_weight,
 		);
 
 		let (support_map, _) = build_support_map::<T::AccountId>(&winners, &staked);

substrate/primitives/phragmen/Cargo.toml

@@ -21,7 +21,7 @@ sp-phragmen-compact = { version = "2.0.0-dev", path = "./compact" }
 [dev-dependencies]
 substrate-test-utils = { version = "2.0.0-dev", path = "../../test-utils" }
 rand = "0.7.3"
-sp-phragmen = { path = "." , version = "2.0.0-dev"}
+sp-phragmen = { version = "2.0.0-dev", path = "." }
 
 [features]
 default = ["std"]

substrate/primitives/phragmen/benches/phragmen.rs

@@ -24,133 +24,139 @@ extern crate test;
 use test::Bencher;
 
 use rand::{self, Rng};
-extern crate sp_phragmen as phragmen;
+use sp_phragmen::{PhragmenResult, VoteWeight};
-use phragmen::{Support, SupportMap, PhragmenStakedAssignment};
 
 use std::collections::BTreeMap;
-use sp_runtime::traits::{Convert, SaturatedConversion};
+use sp_runtime::{Perbill, traits::Zero};
 
-const VALIDATORS: u64 = 1000;
+// default params. Each will be scaled by the benchmarks individually.
-const NOMINATORS: u64 = 10_000;
+const VALIDATORS: u64 = 100;
+const NOMINATORS: u64 = 1_000;
 const EDGES: u64 = 2;
-const TO_ELECT: usize = 100;
+const TO_ELECT: usize = 10;
-const STAKE: Balance = 1000;
+const STAKE: VoteWeight = 1000;
+
+const PREFIX: AccountId = 1000_000;
 
-type Balance = u128;
 type AccountId = u64;
 
-pub struct TestCurrencyToVote;
+mod bench_closure_and_slice {
-impl Convert<Balance, u64> for TestCurrencyToVote {
+	use sp_phragmen::{
-	fn convert(x: Balance) -> u64 { x.saturated_into() }
+		VoteWeight, ExtendedBalance, Assignment, StakedAssignment, IdentifierT,
-}
+		assignment_ratio_to_staked,
-impl Convert<u128, Balance> for TestCurrencyToVote {
+	};
-	fn convert(x: u128) -> Balance { x.saturated_into() }
+	use sp_runtime::{Perbill, PerThing};
+	use rand::{self, Rng, RngCore};
+	use test::Bencher;
+
+	fn random_assignment() -> Assignment<u32, Perbill> {
+		let mut rng = rand::thread_rng();
+		let who = rng.next_u32();
+		let distribution = (0..5)
+			.map(|x| (x + rng.next_u32(), Perbill::from_percent(rng.next_u32() % 100)))
+			.collect::<Vec<_>>();
+		Assignment { who, distribution }
+	}
+
+	/// Converts a vector of ratio assignments into ones with absolute budget value.
+	pub fn assignment_ratio_to_staked_slice<A: IdentifierT, T: PerThing>(
+		ratio: Vec<Assignment<A, T>>,
+		stakes: &[VoteWeight],
+	) -> Vec<StakedAssignment<A>>
+	where
+		T: sp_std::ops::Mul<ExtendedBalance, Output = ExtendedBalance>,
+		ExtendedBalance: From<<T as PerThing>::Inner>,
+	{
+		ratio
+			.into_iter()
+			.zip(stakes.into_iter().map(|x| *x as ExtendedBalance))
+			.map(|(a, stake)| {
+				a.into_staked(stake.into(), true)
+			})
+			.collect()
+	}
+
+	#[bench]
+	fn closure(b: &mut Bencher) {
+		let assignments = (0..1000).map(|_| random_assignment()).collect::<Vec<Assignment<_ ,_>>>();
+		let stake_of = |x: &u32| -> VoteWeight { (x * 2 + 100).into() };
+
+		// each have one clone of assignments
+		b.iter(|| assignment_ratio_to_staked(assignments.clone(), stake_of));
+	}
+
+	#[bench]
+	fn slice(b: &mut Bencher) {
+		let assignments = (0..1000).map(|_| random_assignment()).collect::<Vec<Assignment<_ ,_>>>();
+		let stake_of = |x: &u32| -> VoteWeight { (x * 2 + 100).into() };
+
+		b.iter(|| {
+			let local = assignments.clone();
+			let stakes = local.iter().map(|x| stake_of(&x.who)).collect::<Vec<_>>();
+			assignment_ratio_to_staked_slice(local, stakes.as_ref());
+		});
+	}
 }
 
 fn do_phragmen(
 	b: &mut Bencher,
-	num_vals: u64,
+	num_validators: u64,
-	num_noms: u64,
+	num_nominators: u64,
-	count: usize,
+	to_elect: usize,
-	votes_per: u64,
+	edge_per_voter: u64,
 	eq_iters: usize,
-	_eq_tolerance: u128,
+	eq_tolerance: u128,
 ) {
-	assert!(num_vals > votes_per);
+	assert!(num_validators > edge_per_voter);
-	let rr = |a, b| rand::thread_rng().gen_range(a as usize, b as usize) as Balance;
+	let rr = |a, b| rand::thread_rng().gen_range(a as usize, b as usize) as VoteWeight;
 
-	// prefix to distinguish the validator and nominator account ranges.
+	let mut candidates = Vec::with_capacity(num_validators as usize);
-	let np = 10_000;
+	let mut stake_of_tree: BTreeMap<AccountId, VoteWeight> = BTreeMap::new();
 
-	let mut candidates = Vec::with_capacity(num_vals as usize);
+	(1 ..= num_validators).for_each(|acc| {
-	let mut slashable_balance_of: BTreeMap<AccountId, Balance> = BTreeMap::new();
-
-	(1 ..= num_vals)
-		.for_each(|acc| {
 		candidates.push(acc);
-			slashable_balance_of.insert(acc, STAKE + rr(10, 50));
+		stake_of_tree.insert(acc, STAKE + rr(10, 1000));
 	});
 
-	let mut voters = Vec::with_capacity(num_noms as usize);
+	let mut voters = Vec::with_capacity(num_nominators as usize);
-	(np ..= (np + num_noms))
+	(PREFIX ..= (PREFIX + num_nominators)).for_each(|acc| {
-		.for_each(|acc| {
+		// all possible targets
-			let mut stashes_to_vote = candidates.clone();
+		let mut all_targets = candidates.clone();
-			let votes = (0 .. votes_per)
+		// we remove and pop into `targets` `edge_per_voter` times.
-				.map(|_| {
+		let targets = (0 .. edge_per_voter).map(|_| {
-					stashes_to_vote.remove(rr(0, stashes_to_vote.len()) as usize)
+			all_targets.remove(rr(0, all_targets.len()) as usize)
 		})
 		.collect::<Vec<AccountId>>();
-			voters.push((acc, votes));
+
-			slashable_balance_of.insert(acc, STAKE + rr(10, 50));
+		let stake = STAKE + rr(10, 1000);
+		stake_of_tree.insert(acc, stake);
+		voters.push((acc, stake, targets));
 	});
 
-	let slashable_balance = |who: &AccountId| -> Balance {
-		*slashable_balance_of.get(who).unwrap()
-	};
-
 	b.iter(|| {
-		let r = phragmen::elect::<AccountId, Balance, _, TestCurrencyToVote>(
+		let PhragmenResult { winners, assignments } = sp_phragmen::elect::<AccountId, Perbill>(
-			count,
+			to_elect,
-			1_usize,
+			Zero::zero(),
 			candidates.clone(),
 			voters.clone(),
-			slashable_balance,
-			true,
 		).unwrap();
 
+		let stake_of = |who: &AccountId| -> VoteWeight {
+			*stake_of_tree.get(who).unwrap()
+		};
+
 		// Do the benchmarking with equalize.
 		if eq_iters > 0 {
-			let elected_stashes = r.winners;
+			use sp_phragmen::{equalize, assignment_ratio_to_staked, build_support_map, to_without_backing};
-			let assignments = r.assignments;
+			let staked = assignment_ratio_to_staked(assignments, &stake_of);
+			let winners = to_without_backing(winners);
+			let mut support = build_support_map(winners.as_ref(), staked.as_ref()).0;
 
-			let to_votes = |b: Balance|
+			equalize(
-				<TestCurrencyToVote as Convert<Balance, u128>>::convert(b) as u128;
+				staked.into_iter().map(|a| (a.clone(), stake_of(&a.who))).collect(),
-
+				&mut support,
-			// Initialize the support of each candidate.
+				eq_tolerance,
-			let mut supports = <SupportMap<u64>>::new();
+				eq_iters,
-			elected_stashes
-				.iter()
-				.map(|(e, _)| (e, to_votes(slashable_balance(e))))
-				.for_each(|(e, s)| {
-					let item = Support { own: s, total: s, ..Default::default() };
-					supports.insert(e.clone(), item);
-				});
-
-			// build support struct.
-			for (n, assignment) in assignments.iter() {
-				for (c, per_thing) in assignment.iter() {
-					let nominator_stake = to_votes(slashable_balance(n));
-					let other_stake = *per_thing * nominator_stake;
-					if let Some(support) = supports.get_mut(c) {
-						support.total = support.total.saturating_add(other_stake);
-						support.others.push((n.clone(), other_stake));
-					}
-				}
-			}
-
-			let mut staked_assignments
-				: Vec<(AccountId, Vec<PhragmenStakedAssignment<AccountId>>)>
-				= Vec::with_capacity(assignments.len());
-			for (n, assignment) in assignments.iter() {
-				let mut staked_assignment
-					: Vec<PhragmenStakedAssignment<AccountId>>
-					= Vec::with_capacity(assignment.len());
-				for (c, per_thing) in assignment.iter() {
-					let nominator_stake = to_votes(slashable_balance(n));
-					let other_stake = *per_thing * nominator_stake;
-					staked_assignment.push((c.clone(), other_stake));
-				}
-				staked_assignments.push((n.clone(), staked_assignment));
-			}
-
-			let tolerance = 0_u128;
-			let iterations = 2_usize;
-			phragmen::equalize::<_, _, TestCurrencyToVote, _>(
-				staked_assignments,
-				&mut supports,
-				tolerance,
-				iterations,
-				slashable_balance,
 			);
 		}
 	})
@@ -176,13 +182,13 @@ macro_rules! phragmen_benches {
 
 phragmen_benches! {
 	bench_1_1: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES, 0, 0),
-	bench_1_2: (VALIDATORS*2, NOMINATORS, TO_ELECT, EDGES, 0, 0),
+	bench_1_2: (VALIDATORS * 2, NOMINATORS, TO_ELECT, EDGES, 0, 0),
-	bench_1_3: (VALIDATORS*4, NOMINATORS, TO_ELECT, EDGES, 0, 0),
+	bench_1_3: (VALIDATORS * 4, NOMINATORS, TO_ELECT, EDGES, 0, 0),
-	bench_1_4: (VALIDATORS*8, NOMINATORS, TO_ELECT, EDGES, 0, 0),
+	bench_1_4: (VALIDATORS * 8, NOMINATORS, TO_ELECT, EDGES, 0, 0),
 	bench_1_1_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES, 2, 0),
-	bench_1_2_eq: (VALIDATORS*2, NOMINATORS, TO_ELECT, EDGES, 2, 0),
+	bench_1_2_eq: (VALIDATORS * 2, NOMINATORS, TO_ELECT, EDGES, 2, 0),
-	bench_1_3_eq: (VALIDATORS*4, NOMINATORS, TO_ELECT, EDGES, 2, 0),
+	bench_1_3_eq: (VALIDATORS * 4, NOMINATORS, TO_ELECT, EDGES, 2, 0),
-	bench_1_4_eq: (VALIDATORS*8, NOMINATORS, TO_ELECT, EDGES, 2, 0),
+	bench_1_4_eq: (VALIDATORS * 8, NOMINATORS, TO_ELECT, EDGES, 2, 0),
 
 	bench_0_1: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES, 0, 0),
 	bench_0_2: (VALIDATORS, NOMINATORS, TO_ELECT * 4, EDGES, 0, 0),
@@ -194,20 +200,20 @@ phragmen_benches! {
 	bench_0_4_eq: (VALIDATORS, NOMINATORS, TO_ELECT * 16, EDGES , 2, 0),
 
 	bench_2_1: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES, 0, 0),
-	bench_2_2: (VALIDATORS, NOMINATORS*2, TO_ELECT, EDGES, 0, 0),
+	bench_2_2: (VALIDATORS, NOMINATORS * 2, TO_ELECT, EDGES, 0, 0),
-	bench_2_3: (VALIDATORS, NOMINATORS*4, TO_ELECT, EDGES, 0, 0),
+	bench_2_3: (VALIDATORS, NOMINATORS * 4, TO_ELECT, EDGES, 0, 0),
-	bench_2_4: (VALIDATORS, NOMINATORS*8, TO_ELECT, EDGES, 0, 0),
+	bench_2_4: (VALIDATORS, NOMINATORS * 8, TO_ELECT, EDGES, 0, 0),
 	bench_2_1_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES, 2, 0),
-	bench_2_2_eq: (VALIDATORS, NOMINATORS*2, TO_ELECT, EDGES, 2, 0),
+	bench_2_2_eq: (VALIDATORS, NOMINATORS * 2, TO_ELECT, EDGES, 2, 0),
-	bench_2_3_eq: (VALIDATORS, NOMINATORS*4, TO_ELECT, EDGES, 2, 0),
+	bench_2_3_eq: (VALIDATORS, NOMINATORS * 4, TO_ELECT, EDGES, 2, 0),
-	bench_2_4_eq: (VALIDATORS, NOMINATORS*8, TO_ELECT, EDGES, 2, 0),
+	bench_2_4_eq: (VALIDATORS, NOMINATORS * 8, TO_ELECT, EDGES, 2, 0),
 
 	bench_3_1: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES, 0, 0 ),
-	bench_3_2: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES*2, 0, 0),
+	bench_3_2: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES * 2, 0, 0),
-	bench_3_3: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES*4, 0, 0),
+	bench_3_3: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES * 4, 0, 0),
-	bench_3_4: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES*8, 0, 0),
+	bench_3_4: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES * 8, 0, 0),
 	bench_3_1_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES, 2, 0),
-	bench_3_2_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES*2, 2, 0),
+	bench_3_2_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES * 2, 2, 0),
-	bench_3_3_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES*4, 2, 0),
+	bench_3_3_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES * 4, 2, 0),
-	bench_3_4_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES*8, 2, 0),
+	bench_3_4_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES * 8, 2, 0),
 }

substrate/primitives/phragmen/compact/src/staked.rs

@@ -73,7 +73,7 @@ fn into_impl(count: usize) -> TokenStream2 {
 		quote!(
 			for (voter_index, target_index) in self.#name {
 				let who = voter_at(voter_index).ok_or(_phragmen::Error::CompactInvalidIndex)?;
-				let all_stake = max_of(&who);
+				let all_stake: u128 = max_of(&who).into();
 				assignments.push(_phragmen::StakedAssignment {
 					who,
 					distribution: vec![(target_at(target_index).ok_or(_phragmen::Error::CompactInvalidIndex)?, all_stake)],
@@ -87,7 +87,7 @@ fn into_impl(count: usize) -> TokenStream2 {
 		quote!(
 			for (voter_index, (t1_idx, w1), t2_idx) in self.#name {
 				let who = voter_at(voter_index).ok_or(_phragmen::Error::CompactInvalidIndex)?;
-				let all_stake = max_of(&who);
+				let all_stake: u128 = max_of(&who).into();
 
 				if w1 >= all_stake {
 					return Err(_phragmen::Error::CompactStakeOverflow);
@@ -112,7 +112,7 @@ fn into_impl(count: usize) -> TokenStream2 {
 			for (voter_index, inners, t_last_idx) in self.#name {
 				let who = voter_at(voter_index).ok_or(_phragmen::Error::CompactInvalidIndex)?;
 				let mut sum = u128::min_value();
-				let all_stake = max_of(&who);
+				let all_stake: u128 = max_of(&who).into();
 
 				let mut inners_parsed = inners
 					.iter()
@@ -154,6 +154,7 @@ pub(crate) fn staked(
 
 	let from_impl = from_impl(count);
 	let into_impl = into_impl(count);
+
 	quote!(
 		impl<
 			#voter_type: _phragmen::codec::Codec + Default + Copy,
@@ -196,7 +197,7 @@ pub(crate) fn staked(
 			)
 				-> Result<Vec<_phragmen::StakedAssignment<A>>, _phragmen::Error>
 			where
-				for<'r> FM: Fn(&'r A) -> u128,
+				for<'r> FM: Fn(&'r A) -> u64,
 				A: _phragmen::IdentifierT,
 			{
 				let mut assignments: Vec<_phragmen::StakedAssignment<A>> = Default::default();

substrate/primitives/phragmen/fuzzer/Cargo.toml

@@ -1,19 +1,28 @@
 [package]
 name = "sp-phragmen-fuzzer"
-version = "2.0.0"
+version = "2.0.0-alpha.5"
 authors = ["Parity Technologies <admin@parity.io>"]
 edition = "2018"
+license = "GPL-3.0"
+homepage = "https://substrate.dev"
+repository = "https://github.com/paritytech/substrate/"
+description = "Fuzzer for phragmén implementation."
+documentation = "https://docs.rs/sp-phragmen-fuzzer"
 
 [dependencies]
-sp-phragmen = { version = "2.0.0-alpha.3", path = ".." }
+sp-phragmen = { version = "2.0.0-alpha.5", path = ".." }
+sp-std = { version = "2.0.0-alpha.5", path = "../../std" }
+sp-runtime = { version = "2.0.0-alpha.5", path = "../../runtime" }
 honggfuzz = "0.5"
-rand = "0.7.3"
+rand = { version = "0.7.3", features = ["std", "small_rng"] }
-
-[workspace]
 
 [[bin]]
 name = "reduce"
 path = "src/reduce.rs"
 
+[[bin]]
+name = "equalize"
+path = "src/equalize.rs"
+
 [package.metadata.docs.rs]
 targets = ["x86_64-unknown-linux-gnu"]

substrate/primitives/phragmen/fuzzer/src/common.rs

@@ -0,0 +1,29 @@
+// Copyright 2020 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Common fuzzing utils.
+
+/// converts x into the range [a, b] in a pseudo-fair way.
+pub fn to_range(x: usize, a: usize, b: usize) -> usize {
+	// does not work correctly if b < 2*a
+	assert!(b > 2 * a);
+	let collapsed = x % b;
+	if collapsed >= a {
+		collapsed
+	} else {
+		collapsed + a
+	}
+}

substrate/primitives/phragmen/fuzzer/src/equalize.rs

@@ -0,0 +1,146 @@
+// Copyright 2020 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Fuzzing fro the equalize algorithm
+//!
+//! It ensures that any solution which gets equalized will lead into a better or equally scored
+//! one.
+
+mod common;
+use common::to_range;
+use honggfuzz::fuzz;
+use sp_phragmen::{
+	equalize, assignment_ratio_to_staked, build_support_map, to_without_backing, elect,
+	PhragmenResult, VoteWeight, evaluate_support, is_score_better,
+};
+use sp_std::collections::btree_map::BTreeMap;
+use sp_runtime::Perbill;
+use rand::{self, Rng, SeedableRng, RngCore};
+
+type AccountId = u64;
+
+fn generate_random_phragmen_result(
+	voter_count: u64,
+	target_count: u64,
+	to_elect: usize,
+	edge_per_voter: u64,
+	mut rng: impl RngCore,
+) -> (PhragmenResult<AccountId, Perbill>, BTreeMap<AccountId, VoteWeight>) {
+	let prefix = 100_000;
+	// Note, it is important that stakes are always bigger than ed and
+	let base_stake: u64 = 1_000_000_000;
+	let ed: u64 = base_stake;
+
+	let mut candidates = Vec::with_capacity(target_count as usize);
+	let mut stake_of_tree: BTreeMap<AccountId, VoteWeight> = BTreeMap::new();
+
+	(1..=target_count).for_each(|acc| {
+		candidates.push(acc);
+		let stake_var = rng.gen_range(ed, 100 * ed);
+		stake_of_tree.insert(acc, base_stake + stake_var);
+	});
+
+	let mut voters = Vec::with_capacity(voter_count as usize);
+	(prefix ..= (prefix + voter_count)).for_each(|acc| {
+		// all possible targets
+		let mut all_targets = candidates.clone();
+		// we remove and pop into `targets` `edge_per_voter` times.
+		let targets = (0..edge_per_voter).map(|_| {
+			let upper = all_targets.len() - 1;
+			let idx = rng.gen_range(0, upper);
+			all_targets.remove(idx)
+		})
+		.collect::<Vec<AccountId>>();
+
+		let stake_var = rng.gen_range(ed, 100 * ed) ;
+		let stake = base_stake + stake_var;
+		stake_of_tree.insert(acc, stake);
+		voters.push((acc, stake, targets));
+	});
+
+	(
+		elect::<AccountId, sp_runtime::Perbill>(
+			to_elect,
+			0,
+			candidates,
+			voters,
+		).unwrap(),
+		stake_of_tree,
+	)
+}
+
+fn main() {
+	loop {
+		fuzz!(|data: (usize, usize, usize, usize, usize, u64)| {
+			let (mut target_count, mut voter_count, mut iterations, mut edge_per_voter, mut to_elect, seed) = data;
+			let rng = rand::rngs::SmallRng::seed_from_u64(seed);
+			target_count = to_range(target_count, 50, 2000);
+			voter_count = to_range(voter_count, 50, 1000);
+			iterations = to_range(iterations, 1, 20);
+			to_elect = to_range(to_elect, 25, target_count);
+			edge_per_voter = to_range(edge_per_voter, 1, target_count);
+
+			println!("++ [{} / {} / {} / {}]", voter_count, target_count, to_elect, iterations);
+			let (PhragmenResult { winners, assignments }, stake_of_tree) = generate_random_phragmen_result(
+				voter_count as u64,
+				target_count as u64,
+				to_elect,
+				edge_per_voter as u64,
+				rng,
+			);
+
+			let stake_of = |who: &AccountId| -> VoteWeight {
+				*stake_of_tree.get(who).unwrap()
+			};
+
+			let mut staked = assignment_ratio_to_staked(assignments.clone(), &stake_of);
+			let winners = to_without_backing(winners);
+			let mut support = build_support_map(winners.as_ref(), staked.as_ref()).0;
+
+			let initial_score = evaluate_support(&support);
+			if initial_score[0] == 0 {
+				// such cases cannot be improved by reduce.
+				return;
+			}
+
+			let i = equalize(
+				&mut staked,
+				&mut support,
+				10,
+				iterations,
+			);
+
+			let final_score = evaluate_support(&support);
+			if final_score[0] == initial_score[0] {
+				// such solutions can only be improved by such a tiny fiction that it is most often
+				// wrong due to rounding errors.
+				return;
+			}
+
+			let enhance = is_score_better(initial_score, final_score);
+
+			println!(
+				"iter = {} // {:?} -> {:?} [{}]",
+				i,
+				initial_score,
+				final_score,
+				enhance,
+			);
+			// if more than one iteration has been done, or they must be equal.
+			assert!(enhance || initial_score == final_score || i == 0)
+		});
+	}
+}

substrate/primitives/phragmen/fuzzer/src/reduce.rs

@@ -14,6 +14,12 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
+//! Fuzzing for the reduce algorithm.
+//!
+//! It that reduce always return a new set og edges in which the bound is kept (`edges_after <= m +
+//! n,`) and the result must effectively be the same, meaning that the same support map should be
+//! computable from both.
+//!
 //! # Running
 //!
 //! Run with `cargo hfuzz run reduce`. `honggfuzz`.
@@ -24,8 +30,11 @@
 //! `cargo hfuzz run-debug reduce hfuzz_workspace/reduce/*.fuzz`.
 
 use honggfuzz::fuzz;
+
+mod common;
+use common::to_range;
 use sp_phragmen::{StakedAssignment, ExtendedBalance, build_support_map, reduce};
-use rand::{self, Rng};
+use rand::{self, Rng, SeedableRng, RngCore};
 
 type Balance = u128;
 type AccountId = u64;
@@ -35,12 +44,17 @@ const KSM: Balance = 1_000_000_000_000;
 
 fn main() {
 	loop {
-		fuzz!(|_data: _| {
+		fuzz!(|data: (usize, usize, u64)| {
+			let (mut voter_count, mut target_count, seed) = data;
+			let rng = rand::rngs::SmallRng::seed_from_u64(seed);
+			target_count = to_range(target_count, 100, 1000);
+			voter_count = to_range(voter_count, 100, 2000);
 			let (assignments, winners) = generate_random_phragmen_assignment(
-	 			rr(100, 1000),
+				voter_count,
-	 			rr(100, 2000),
+				target_count,
 				8,
 				8,
+				rng
 			);
 			reduce_and_compare(&assignments, &winners);
 		});
@@ -52,13 +66,10 @@ fn generate_random_phragmen_assignment(
 	target_count: usize,
 	avg_edge_per_voter: usize,
 	edge_per_voter_var: usize,
+	mut rng: impl RngCore,
 ) -> (Vec<StakedAssignment<AccountId>>, Vec<AccountId>) {
-	// random in range of (a, b)
-	let rr_128 = |a: u128, b: u128| -> u128 { rand::thread_rng().gen_range(a, b) };
-
 	// prefix to distinguish the voter and target account ranges.
 	let target_prefix = 1_000_000;
-	// let target_prefix = 1000;
 	assert!(voter_count < target_prefix);
 
 	let mut assignments = Vec::with_capacity(voter_count as usize);
@@ -70,17 +81,17 @@ fn generate_random_phragmen_assignment(
 
 	(1..=voter_count).for_each(|acc| {
 		let mut targets_to_chose_from = all_targets.clone();
-		let targets_to_chose = if edge_per_voter_var > 0 { rr(
+		let targets_to_chose = if edge_per_voter_var > 0 { rng.gen_range(
 			avg_edge_per_voter - edge_per_voter_var,
 			avg_edge_per_voter + edge_per_voter_var,
 		) } else { avg_edge_per_voter };
 
 		let distribution = (0..targets_to_chose).map(|_| {
-			let target = targets_to_chose_from.remove(rr(0, targets_to_chose_from.len()));
+			let target = targets_to_chose_from.remove(rng.gen_range(0, targets_to_chose_from.len()));
 			if winners.iter().find(|w| **w == target).is_none() {
 				winners.push(target.clone());
 			}
-			(target, rr_128(1 * KSM, 100 * KSM))
+			(target, rng.gen_range(1 * KSM, 100 * KSM))
 		}).collect::<Vec<(AccountId, ExtendedBalance)>>();
 
 		assignments.push(StakedAssignment {
@@ -139,7 +150,3 @@ fn assignment_len(assignments: &[StakedAssignment<AccountId>]) -> u32 {
 	assignments.iter().for_each(|x| x.distribution.iter().for_each(|_| counter += 1));
 	counter
 }
-
-fn rr(a: usize, b: usize) -> usize {
-	rand::thread_rng().gen_range(a, b)
-}

substrate/primitives/phragmen/src/helpers.rs

@@ -16,7 +16,7 @@
 
 //! Helper methods for phragmen.
 
-use crate::{Assignment, ExtendedBalance, IdentifierT, StakedAssignment};
+use crate::{Assignment, ExtendedBalance, VoteWeight, IdentifierT, StakedAssignment, WithApprovalOf};
 use sp_runtime::PerThing;
 use sp_std::prelude::*;
 
@@ -26,7 +26,7 @@ pub fn assignment_ratio_to_staked<A: IdentifierT, T: PerThing, FS>(
 	stake_of: FS,
 ) -> Vec<StakedAssignment<A>>
 where
-	for<'r> FS: Fn(&'r A) -> ExtendedBalance,
+	for<'r> FS: Fn(&'r A) -> VoteWeight,
 	T: sp_std::ops::Mul<ExtendedBalance, Output = ExtendedBalance>,
 	ExtendedBalance: From<<T as PerThing>::Inner>,
 {
@@ -34,30 +34,34 @@ where
 		.into_iter()
 		.map(|a| {
 			let stake = stake_of(&a.who);
-			a.into_staked(stake, true)
+			a.into_staked(stake.into(), true)
 		})
 		.collect()
 }
 
 /// Converts a vector of staked assignments into ones with ratio values.
 pub fn assignment_staked_to_ratio<A: IdentifierT, T: PerThing>(
-	ratio: Vec<StakedAssignment<A>>,
+	staked: Vec<StakedAssignment<A>>,
 ) -> Vec<Assignment<A, T>>
 where
 	ExtendedBalance: From<<T as PerThing>::Inner>,
 {
-	ratio.into_iter().map(|a| a.into_assignment(true)).collect()
+	staked.into_iter().map(|a| a.into_assignment(true)).collect()
+}
+
+/// consumes a vector of winners with backing stake to just winners.
+pub fn to_without_backing<A: IdentifierT>(winners: Vec<WithApprovalOf<A>>) -> Vec<A> {
+	winners.into_iter().map(|(who, _)| who).collect::<Vec<A>>()
 }
 
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use crate::ExtendedBalance;
 	use sp_runtime::Perbill;
 
 	#[test]
 	fn into_staked_works() {
-		let ratio = vec![
+		let assignments = vec![
 			Assignment {
 				who: 1u32,
 				distribution: vec![
@@ -74,8 +78,8 @@ mod tests {
 			},
 		];
 
-		let stake_of = |_: &u32| -> ExtendedBalance { 100u128 };
+		let stake_of = |_: &u32| -> VoteWeight { 100 };
-		let staked = assignment_ratio_to_staked(ratio, stake_of);
+		let staked = assignment_ratio_to_staked(assignments, stake_of);
 
 		assert_eq!(
 			staked,

substrate/primitives/phragmen/src/lib.rs

@@ -35,7 +35,7 @@
 
 use sp_std::{prelude::*, collections::btree_map::BTreeMap, fmt::Debug, cmp::Ordering, convert::TryFrom};
 use sp_runtime::{helpers_128bit::multiply_by_rational, PerThing, Rational128, RuntimeDebug, SaturatedConversion};
-use sp_runtime::traits::{Zero, Convert, Member, AtLeast32Bit, Saturating, Bounded};
+use sp_runtime::traits::{Zero, Member, Saturating, Bounded};
 
 #[cfg(test)]
 mod mock;
@@ -88,17 +88,19 @@ pub enum Error {
 	CompactInvalidIndex,
 }
 
-/// A type in which performing operations on balances and stakes of candidates and voters are safe.
+/// A type which is used in the API of this crate as a numeric weight of a vote, most often the
-///
+/// stake of the voter. It is always converted to [`ExtendedBalance`] for computation.
-/// This module's functions expect a `Convert` type to convert all balances to u64. Hence, u128 is
+pub type VoteWeight = u64;
-/// a safe type for arithmetic operations over them.
+
-///
+/// A type in which performing operations on vote weights are safe.
-/// Balance types converted to `ExtendedBalance` are referred to as `Votes`.
 pub type ExtendedBalance = u128;
 
 /// The score of an assignment. This can be computed from the support map via [`evaluate_support`].
 pub type PhragmenScore = [ExtendedBalance; 3];
 
+/// A winner, with their respective approval stake.
+pub type WithApprovalOf<A> = (A, ExtendedBalance);
+
 /// The denominator used for loads. Since votes are collected as u64, the smallest ratio that we
 /// might collect is `1/approval_stake` where approval stake is the sum of votes. Hence, some number
 /// bigger than u64::max_value() is needed. For maximum accuracy we simply use u128;
@@ -146,7 +148,7 @@ struct Edge<AccountId> {
 pub struct PhragmenResult<AccountId, T: PerThing> {
 	/// Just winners zipped with their approval stake. Note that the approval stake is merely the
 	/// sub of their received stake and could be used for very basic sorting and approval voting.
-	pub winners: Vec<(AccountId, ExtendedBalance)>,
+	pub winners: Vec<WithApprovalOf<AccountId>>,
 	/// Individual assignments. for each tuple, the first elements is a voter and the second
 	/// is the list of candidates that it supports.
 	pub assignments: Vec<Assignment<AccountId, T>>,
@@ -285,6 +287,11 @@ impl<AccountId> StakedAssignment<AccountId> {
 			distribution,
 		}
 	}
+
+	/// Get the total stake of this assignment (aka voter budget).
+	pub fn total(&self) -> ExtendedBalance {
+		self.distribution.iter().fold(Zero::zero(), |a, b| a.saturating_add(b.1))
+	}
 }
 
 /// A structure to demonstrate the phragmen result from the perspective of the candidate, i.e. how
@@ -316,25 +323,20 @@ pub type SupportMap<A> = BTreeMap<A, Support<A>>;
 ///   `None` is returned.
 /// * `initial_candidates`: candidates list to be elected from.
 /// * `initial_voters`: voters list.
-/// * `stake_of`: something that can return the stake stake of a particular candidate or voter.
 ///
 /// This function does not strip out candidates who do not have any backing stake. It is the
 /// responsibility of the caller to make sure only those candidates who have a sensible economic
 /// value are passed in. From the perspective of this function, a candidate can easily be among the
 /// winner with no backing stake.
-pub fn elect<AccountId, Balance, C, R>(
+pub fn elect<AccountId, R>(
 	candidate_count: usize,
 	minimum_candidate_count: usize,
 	initial_candidates: Vec<AccountId>,
-	initial_voters: Vec<(AccountId, Balance, Vec<AccountId>)>,
+	initial_voters: Vec<(AccountId, VoteWeight, Vec<AccountId>)>,
 ) -> Option<PhragmenResult<AccountId, R>> where
 	AccountId: Default + Ord + Member,
-	Balance: Default + Copy + AtLeast32Bit,
-	C: Convert<Balance, u64> + Convert<u128, Balance>,
 	R: PerThing,
 {
-	let to_votes = |b: Balance| <C as Convert<Balance, u64>>::convert(b) as ExtendedBalance;
-
 	// return structures
 	let mut elected_candidates: Vec<(AccountId, ExtendedBalance)>;
 	let mut assigned: Vec<Assignment<AccountId, R>>;
@@ -368,14 +370,14 @@ pub fn elect<AccountId, Balance, C, R>(
 			if let Some(idx) = c_idx_cache.get(&v) {
 				// This candidate is valid + already cached.
 				candidates[*idx].approval_stake = candidates[*idx].approval_stake
-					.saturating_add(to_votes(voter_stake));
+					.saturating_add(voter_stake.into());
 				edges.push(Edge { who: v.clone(), candidate_index: *idx, ..Default::default() });
 			} // else {} would be wrong votes. We don't really care about it.
 		}
 		Voter {
 			who,
 			edges: edges,
-			budget: to_votes(voter_stake),
+			budget: voter_stake.into(),
 			load: Rational128::zero(),
 		}
 	}));
@@ -633,32 +635,27 @@ pub fn is_score_better(this: PhragmenScore, that: PhragmenScore) -> bool {
 /// rounds. The number of rounds and the maximum diff-per-round tolerance can be tuned through input
 /// parameters.
 ///
-/// No value is returned from the function and the `supports` parameter is updated.
+/// Returns the number of iterations that were preformed.
 ///
 /// - `assignments`: exactly the same is the output of phragmen.
 /// - `supports`: mutable reference to s `SupportMap`. This parameter is updated.
 /// - `tolerance`: maximum difference that can occur before an early quite happens.
 /// - `iterations`: maximum number of iterations that will be processed.
-/// - `stake_of`: something that can return the stake stake of a particular candidate or voter.
+pub fn equalize<AccountId>(
-pub fn equalize<Balance, AccountId, C, FS>(
+	assignments: &mut Vec<StakedAssignment<AccountId>>,
-	mut assignments: Vec<StakedAssignment<AccountId>>,
 	supports: &mut SupportMap<AccountId>,
 	tolerance: ExtendedBalance,
 	iterations: usize,
-	stake_of: FS,
+) -> usize where AccountId: Ord + Clone {
-) where
+	if iterations == 0 { return 0; }
-	C: Convert<Balance, u64> + Convert<u128, Balance>,
+
-	for<'r> FS: Fn(&'r AccountId) -> Balance,
+	let mut i = 0 ;
-	AccountId: Ord + Clone,
+	loop {
-{
-	// prepare the data for equalise
-	for _i in 0..iterations {
 		let mut max_diff = 0;
-
+		for assignment in assignments.iter_mut() {
-		for StakedAssignment { who, distribution } in assignments.iter_mut() {
+			let voter_budget = assignment.total();
-			let voter_budget = stake_of(&who);
+			let StakedAssignment { who, distribution } =  assignment;
-
+			let diff = do_equalize(
-			let diff = do_equalize::<_, _, C>(
 				who,
 				voter_budget,
 				distribution,
@@ -668,28 +665,22 @@ pub fn equalize<Balance, AccountId, C, FS>(
 			if diff > max_diff { max_diff = diff; }
 		}
 
-		if max_diff < tolerance {
+		i += 1;
-			break;
+		if max_diff <= tolerance || i >= iterations {
+			break i;
 		}
 	}
 }
 
 /// actually perform equalize. same interface is `equalize`. Just called in loops with a check for
 /// maximum difference.
-fn do_equalize<Balance, AccountId, C>(
+fn do_equalize<AccountId>(
 	voter: &AccountId,
-	budget_balance: Balance,
+	budget: ExtendedBalance,
 	elected_edges: &mut Vec<(AccountId, ExtendedBalance)>,
 	support_map: &mut SupportMap<AccountId>,
 	tolerance: ExtendedBalance
-) -> ExtendedBalance where
+) -> ExtendedBalance where AccountId: Ord + Clone {
-	C: Convert<Balance, u64> + Convert<u128, Balance>,
-	AccountId: Ord + Clone,
-{
-	let to_votes = |b: Balance|
-		<C as Convert<Balance, u64>>::convert(b) as ExtendedBalance;
-	let budget = to_votes(budget_balance);
-
 	// Nothing to do. This voter had nothing useful.
 	// Defensive only. Assignment list should always be populated. 1 might happen for self vote.
 	if elected_edges.is_empty() || elected_edges.len() == 1 { return 0; }

substrate/primitives/phragmen/src/mock.rs

@@ -18,21 +18,13 @@
 
 #![cfg(test)]
 
-use crate::{elect, PhragmenResult, Assignment};
+use crate::{elect, PhragmenResult, Assignment, VoteWeight, ExtendedBalance};
 use sp_runtime::{
 	assert_eq_error_rate, PerThing,
-	traits::{Convert, Member, SaturatedConversion, Zero, One}
+	traits::{Member, SaturatedConversion, Zero, One}
 };
 use sp_std::collections::btree_map::BTreeMap;
 
-pub(crate) struct TestCurrencyToVote;
-impl Convert<Balance, u64> for TestCurrencyToVote {
-	fn convert(x: Balance) -> u64 { x.saturated_into() }
-}
-impl Convert<u128, Balance> for TestCurrencyToVote {
-	fn convert(x: u128) -> Balance { x }
-}
-
 #[derive(Default, Debug)]
 pub(crate) struct _Candidate<A> {
 	who: A,
@@ -66,12 +58,11 @@ pub(crate) struct _Support<A> {
 pub(crate) type _PhragmenAssignment<A> = (A, f64);
 pub(crate) type _SupportMap<A> = BTreeMap<A, _Support<A>>;
 
-pub(crate) type Balance = u128;
 pub(crate) type AccountId = u64;
 
 #[derive(Debug, Clone)]
 pub(crate) struct _PhragmenResult<A: Clone> {
-	pub winners: Vec<(A, Balance)>,
+	pub winners: Vec<(A, ExtendedBalance)>,
 	pub assignments: Vec<(A, Vec<_PhragmenAssignment<A>>)>
 }
 
@@ -87,9 +78,9 @@ pub(crate) fn elect_float<A, FS>(
 	stake_of: FS,
 ) -> Option<_PhragmenResult<A>> where
 	A: Default + Ord + Member + Copy,
-	for<'r> FS: Fn(&'r A) -> Balance,
+	for<'r> FS: Fn(&'r A) -> VoteWeight,
 {
-	let mut elected_candidates: Vec<(A, Balance)>;
+	let mut elected_candidates: Vec<(A, ExtendedBalance)>;
 	let mut assigned: Vec<(A, Vec<_PhragmenAssignment<A>>)>;
 	let mut c_idx_cache = BTreeMap::<A, usize>::new();
 	let num_voters = initial_candidates.len() + initial_voters.len();
@@ -161,7 +152,7 @@ pub(crate) fn elect_float<A, FS>(
 				}
 			}
 
-			elected_candidates.push((winner.who.clone(), winner.approval_stake as Balance));
+			elected_candidates.push((winner.who.clone(), winner.approval_stake as ExtendedBalance));
 		} else {
 			break
 		}
@@ -195,7 +186,7 @@ pub(crate) fn equalize_float<A, FS>(
 	iterations: usize,
 	stake_of: FS,
 ) where
-	for<'r> FS: Fn(&'r A) -> Balance,
+	for<'r> FS: Fn(&'r A) -> VoteWeight,
 	A: Ord + Clone + std::fmt::Debug,
 {
 	for _i in 0..iterations {
@@ -220,7 +211,7 @@ pub(crate) fn equalize_float<A, FS>(
 
 pub(crate) fn do_equalize_float<A>(
 	voter: &A,
-	budget_balance: Balance,
+	budget_balance: VoteWeight,
 	elected_edges: &mut Vec<_PhragmenAssignment<A>>,
 	support_map: &mut _SupportMap<A>,
 	tolerance: f64
@@ -310,12 +301,12 @@ pub(crate) fn do_equalize_float<A>(
 }
 
 
-pub(crate) fn create_stake_of(stakes: &[(AccountId, Balance)])
+pub(crate) fn create_stake_of(stakes: &[(AccountId, VoteWeight)])
-	-> Box<dyn Fn(&AccountId) -> Balance>
+	-> Box<dyn Fn(&AccountId) -> VoteWeight>
 {
-	let mut storage = BTreeMap::<AccountId, Balance>::new();
+	let mut storage = BTreeMap::<AccountId, VoteWeight>::new();
 	stakes.iter().for_each(|s| { storage.insert(s.0, s.1); });
-	let stake_of = move |who: &AccountId| -> Balance { storage.get(who).unwrap().to_owned() };
+	let stake_of = move |who: &AccountId| -> VoteWeight { storage.get(who).unwrap().to_owned() };
 	Box::new(stake_of)
 }
 
@@ -331,12 +322,12 @@ pub fn check_assignments_sum<T: PerThing>(assignments: Vec<Assignment<AccountId,
 pub(crate) fn run_and_compare<Output: PerThing>(
 	candidates: Vec<AccountId>,
 	voters: Vec<(AccountId, Vec<AccountId>)>,
-	stake_of: &Box<dyn Fn(&AccountId) -> Balance>,
+	stake_of: &Box<dyn Fn(&AccountId) -> VoteWeight>,
 	to_elect: usize,
 	min_to_elect: usize,
 ) {
 	// run fixed point code.
-	let PhragmenResult { winners, assignments } = elect::<_, _, TestCurrencyToVote, Output>(
+	let PhragmenResult { winners, assignments } = elect::<_, Output>(
 		to_elect,
 		min_to_elect,
 		candidates.clone(),
@@ -352,7 +343,7 @@ pub(crate) fn run_and_compare<Output: PerThing>(
 		&stake_of,
 	).unwrap();
 
-	assert_eq!(winners, truth_value.winners);
+	assert_eq!(winners.iter().map(|(x, _)| x).collect::<Vec<_>>(), truth_value.winners.iter().map(|(x, _)| x).collect::<Vec<_>>());
 
 	for Assignment { who, distribution } in assignments.clone() {
 		if let Some(float_assignments) = truth_value.assignments.iter().find(|x| x.0 == who) {
@@ -379,7 +370,7 @@ pub(crate) fn build_support_map_float<FS>(
 	result: &mut _PhragmenResult<AccountId>,
 	stake_of: FS,
 ) -> _SupportMap<AccountId>
-	where for<'r> FS: Fn(&'r AccountId) -> Balance
+	where for<'r> FS: Fn(&'r AccountId) -> VoteWeight
 {
 	let mut supports = <_SupportMap<AccountId>>::new();
 	result.winners

substrate/primitives/phragmen/src/tests.rs

@@ -20,11 +20,11 @@
 
 use crate::mock::*;
 use crate::{
-	elect, equalize, build_support_map, is_score_better,
+	elect, equalize, build_support_map, is_score_better, helpers::*,
 	Support, StakedAssignment, Assignment, PhragmenResult, ExtendedBalance,
 };
 use substrate_test_utils::assert_eq_uvec;
-use sp_runtime::{Perbill, Permill, Percent, PerU16, traits::Convert};
+use sp_runtime::{Perbill, Permill, Percent, PerU16};
 
 #[test]
 fn float_phragmen_poc_works() {
@@ -82,7 +82,7 @@ fn phragmen_poc_works() {
 	];
 
 	let stake_of = create_stake_of(&[(10, 10), (20, 20), (30, 30)]);
-	let PhragmenResult { winners, assignments } = elect::<_, _, TestCurrencyToVote, Perbill>(
+	let PhragmenResult { winners, assignments } = elect::<_, Perbill>(
 		2,
 		2,
 		candidates,
@@ -110,6 +110,77 @@ fn phragmen_poc_works() {
 			},
 		]
 	);
+
+	let mut staked = assignment_ratio_to_staked(assignments, &stake_of);
+	let winners = to_without_backing(winners);
+	let mut support_map = build_support_map::<AccountId>(&winners, &staked).0;
+
+	assert_eq_uvec!(
+		staked,
+		vec![
+			StakedAssignment {
+				who: 10u64,
+				distribution: vec![(2, 10)],
+			},
+			StakedAssignment {
+				who: 20,
+				distribution: vec![(3, 20)],
+			},
+			StakedAssignment {
+				who: 30,
+				distribution: vec![
+					(2, 15),
+					(3, 15),
+				],
+			},
+		]
+	);
+
+	assert_eq!(
+		*support_map.get(&2).unwrap(),
+		Support::<AccountId> { total: 25, voters: vec![(10, 10), (30, 15)] },
+	);
+	assert_eq!(
+		*support_map.get(&3).unwrap(),
+		Support::<AccountId> { total: 35, voters: vec![(20, 20), (30, 15)] },
+	);
+
+	equalize(
+		&mut staked,
+		&mut support_map,
+		0,
+		2,
+	);
+
+	assert_eq_uvec!(
+		staked,
+		vec![
+			StakedAssignment {
+				who: 10u64,
+				distribution: vec![(2, 10)],
+			},
+			StakedAssignment {
+				who: 20,
+				distribution: vec![(3, 20)],
+			},
+			StakedAssignment {
+				who: 30,
+				distribution: vec![
+					(2, 20),
+					(3, 10),
+				],
+			},
+		]
+	);
+
+	assert_eq!(
+		*support_map.get(&2).unwrap(),
+		Support::<AccountId> { total: 30, voters: vec![(10, 10), (30, 20)] },
+	);
+	assert_eq!(
+		*support_map.get(&3).unwrap(),
+		Support::<AccountId> { total: 30, voters: vec![(20, 20), (30, 10)] },
+	);
 }
 
 #[test]
@@ -168,7 +239,7 @@ fn phragmen_accuracy_on_large_scale_only_validators() {
 		(5, (u64::max_value() - 2).into()),
 	]);
 
-	let PhragmenResult { winners, assignments } = elect::<_, _, TestCurrencyToVote, Perbill>(
+	let PhragmenResult { winners, assignments } = elect::<_, Perbill>(
 		2,
 		2,
 		candidates.clone(),
@@ -198,7 +269,7 @@ fn phragmen_accuracy_on_large_scale_validators_and_nominators() {
 		(14, u64::max_value().into()),
 	]);
 
-	let PhragmenResult { winners, assignments } = elect::<_, _, TestCurrencyToVote, Perbill>(
+	let PhragmenResult { winners, assignments } = elect::<_, Perbill>(
 		2,
 		2,
 		candidates,
@@ -241,7 +312,7 @@ fn phragmen_accuracy_on_small_scale_self_vote() {
 		(30, 1),
 	]);
 
-	let PhragmenResult { winners, assignments: _ } = elect::<_, _, TestCurrencyToVote, Perbill>(
+	let PhragmenResult { winners, assignments: _ } = elect::<_, Perbill>(
 		3,
 		3,
 		candidates,
@@ -271,7 +342,7 @@ fn phragmen_accuracy_on_small_scale_no_self_vote() {
 		(3, 1),
 	]);
 
-	let PhragmenResult { winners, assignments: _ } = elect::<_, _, TestCurrencyToVote, Perbill>(
+	let PhragmenResult { winners, assignments: _ } = elect::<_, Perbill>(
 		3,
 		3,
 		candidates,
@@ -304,7 +375,7 @@ fn phragmen_large_scale_test() {
 		(50, 990000000000000000),
 	]);
 
-	let PhragmenResult { winners, assignments } = elect::<_, _, TestCurrencyToVote, Perbill>(
+	let PhragmenResult { winners, assignments } = elect::<_, Perbill>(
 		2,
 		2,
 		candidates,
@@ -330,7 +401,7 @@ fn phragmen_large_scale_test_2() {
 		(50, nom_budget.into()),
 	]);
 
-	let PhragmenResult { winners, assignments } = elect::<_, _, TestCurrencyToVote, Perbill>(
+	let PhragmenResult { winners, assignments } = elect::<_, Perbill>(
 		2,
 		2,
 		candidates,
@@ -406,7 +477,7 @@ fn elect_has_no_entry_barrier() {
 		(2, 10),
 	]);
 
-	let PhragmenResult { winners, assignments: _ } = elect::<_, _, TestCurrencyToVote, Perbill>(
+	let PhragmenResult { winners, assignments: _ } = elect::<_, Perbill>(
 		3,
 		3,
 		candidates,
@@ -433,7 +504,7 @@ fn minimum_to_elect_is_respected() {
 		(2, 10),
 	]);
 
-	let maybe_result = elect::<_, _, TestCurrencyToVote, Perbill>(
+	let maybe_result = elect::<_, Perbill>(
 		10,
 		10,
 		candidates,
@@ -459,7 +530,7 @@ fn self_votes_should_be_kept() {
 		(1, 8),
 	]);
 
-	let result = elect::<_, _, TestCurrencyToVote, Perbill>(
+	let result = elect::<_, Perbill>(
 		2,
 		2,
 		candidates,
@@ -480,16 +551,11 @@ fn self_votes_should_be_kept() {
 		],
 	);
 
-	let staked_assignments: Vec<StakedAssignment<AccountId>> = result.assignments
+	let mut staked_assignments = assignment_ratio_to_staked(result.assignments, &stake_of);
-		.into_iter()
+	let winners = to_without_backing(result.winners);
-		.map(|a| {
-			let stake = <TestCurrencyToVote as Convert<Balance, u64>>::convert(stake_of(&a.who)) as ExtendedBalance;
-			a.into_staked(stake, true)
-		}).collect();
 
-	let winners = result.winners.into_iter().map(|(who, _)| who).collect::<Vec<AccountId>>();
 	let (mut supports, _) = build_support_map::<AccountId>(
-		winners.as_slice(),
+		&winners,
 		&staked_assignments,
 	);
 
@@ -503,12 +569,11 @@ fn self_votes_should_be_kept() {
 		&Support { total: 24u128, voters: vec![(20u64, 20u128), (1u64, 4u128)] },
 	);
 
-	equalize::<Balance, AccountId, TestCurrencyToVote, _>(
+	equalize(
-		staked_assignments,
+		&mut staked_assignments,
 		&mut supports,
 		0,
 		2usize,
-		&stake_of,
 	);
 
 	assert_eq!(
@@ -526,7 +591,7 @@ fn assignment_convert_works() {
 	let staked = StakedAssignment {
 		who: 1 as AccountId,
 		distribution: vec![
-			(20, 100 as Balance),
+			(20, 100 as ExtendedBalance),
 			(30, 25),
 		],
 	};
@@ -578,10 +643,10 @@ fn score_comparison_is_lexicographical() {
 
 mod compact {
 	use codec::{Decode, Encode};
-	use crate::generate_compact_solution_type;
+	use crate::{generate_compact_solution_type, VoteWeight};
-	use super::{AccountId, Balance};
+	use super::{AccountId};
 	// these need to come from the same dev-dependency `sp-phragmen`, not from the crate.
-	use sp_phragmen::{Assignment, StakedAssignment, Error as PhragmenError};
+	use sp_phragmen::{Assignment, StakedAssignment, Error as PhragmenError, ExtendedBalance};
 	use sp_std::{convert::{TryInto, TryFrom}, fmt::Debug};
 	use sp_runtime::Percent;
 
@@ -736,7 +801,7 @@ mod compact {
 		let assignments = vec![
 			StakedAssignment {
 				who: 2 as AccountId,
-				distribution: vec![(20, 100 as Balance)]
+				distribution: vec![(20, 100 as ExtendedBalance)]
 			},
 			StakedAssignment {
 				who: 4,
@@ -773,7 +838,7 @@ mod compact {
 			targets.iter().position(|x| x == a).map(TryInto::try_into).unwrap().ok()
 		};
 
-		let compacted = <TestCompact<u16, u16, Balance>>::from_staked(
+		let compacted = <TestCompact<u16, u16, ExtendedBalance>>::from_staked(
 			assignments.clone(),
 			voter_index,
 			target_index,
@@ -794,7 +859,7 @@ mod compact {
 			}
 		);
 
-		let max_of_fn = |_: &AccountId| -> Balance { 100u128 };
+		let max_of_fn = |_: &AccountId| -> VoteWeight { 100 };
 		let voter_at = |a: u16| -> Option<AccountId> { voters.get(a as usize).cloned() };
 		let target_at = |a: u16| -> Option<AccountId> { targets.get(a as usize).cloned() };
 
@@ -812,14 +877,14 @@ mod compact {
 	fn compact_into_stake_must_report_overflow() {
 		// The last edge which is computed from the rest should ALWAYS be positive.
 		// in votes2
-		let compact = TestCompact::<u16, u16, Balance> {
+		let compact = TestCompact::<u16, u16, ExtendedBalance> {
 			votes1: Default::default(),
 			votes2: vec![(0, (1, 10), 2)],
 			..Default::default()
 		};
 
 		let entity_at = |a: u16| -> Option<AccountId> { Some(a as AccountId) };
-		let max_of = |_: &AccountId| -> Balance { 5 };
+		let max_of = |_: &AccountId| -> VoteWeight { 5 };
 
 		assert_eq!(
 			compact.into_staked(&max_of, &entity_at, &entity_at).unwrap_err(),
@@ -827,7 +892,7 @@ mod compact {
 		);
 
 		// in votes3 onwards
-		let compact = TestCompact::<u16, u16, Balance> {
+		let compact = TestCompact::<u16, u16, ExtendedBalance> {
 			votes1: Default::default(),
 			votes2: Default::default(),
 			votes3: vec![(0, [(1, 7), (2, 8)], 3)],
@@ -840,7 +905,7 @@ mod compact {
 		);
 
 		// Also if equal
-		let compact = TestCompact::<u16, u16, Balance> {
+		let compact = TestCompact::<u16, u16, ExtendedBalance> {
 			votes1: Default::default(),
 			votes2: Default::default(),
 			// 5 is total, we cannot leave none for 30 here.
@@ -889,13 +954,13 @@ mod compact {
 		let assignments = vec![
 			StakedAssignment {
 				who: 1 as AccountId,
-				distribution: (10..26).map(|i| (i as AccountId, i as Balance)).collect::<Vec<_>>(),
+				distribution: (10..26).map(|i| (i as AccountId, i as ExtendedBalance)).collect::<Vec<_>>(),
 			},
 		];
 
 		let entity_index = |a: &AccountId| -> Option<u16> { Some(*a as u16) };
 
-		let compacted = <TestCompact<u16, u16, Balance>>::from_staked(
+		let compacted = <TestCompact<u16, u16, ExtendedBalance>>::from_staked(
 			assignments.clone(),
 			entity_index,
 			entity_index,
@@ -906,11 +971,11 @@ mod compact {
 		let assignments = vec![
 			StakedAssignment {
 				who: 1 as AccountId,
-				distribution: (10..27).map(|i| (i as AccountId, i as Balance)).collect::<Vec<_>>(),
+				distribution: (10..27).map(|i| (i as AccountId, i as ExtendedBalance)).collect::<Vec<_>>(),
 			},
 		];
 
-		let compacted = <TestCompact<u16, u16, Balance>>::from_staked(
+		let compacted = <TestCompact<u16, u16, ExtendedBalance>>::from_staked(
 			assignments.clone(),
 			entity_index,
 			entity_index,
@@ -948,7 +1013,7 @@ mod compact {
 		let assignments = vec![
 			StakedAssignment {
 				who: 1 as AccountId,
-				distribution: vec![(10, 100 as Balance), (11, 100)]
+				distribution: vec![(10, 100 as ExtendedBalance), (11, 100)]
 			},
 			StakedAssignment {
 				who: 2,
@@ -963,7 +1028,7 @@ mod compact {
 			targets.iter().position(|x| x == a).map(TryInto::try_into).unwrap().ok()
 		};
 
-		let compacted = <TestCompact<u16, u16, Balance>>::from_staked(
+		let compacted = <TestCompact<u16, u16, ExtendedBalance>>::from_staked(
 			assignments.clone(),
 			voter_index,
 			target_index,