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/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 staked = assignment_ratio_to_staked(ratio, stake_of);
+		let stake_of = |_: &u32| -> VoteWeight { 100 };
+		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.
-///
-/// This module's functions expect a `Convert` type to convert all balances to u64. Hence, u128 is
-/// a safe type for arithmetic operations over them.
-///
-/// Balance types converted to `ExtendedBalance` are referred to as `Votes`.
+/// 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.
+pub type VoteWeight = u64;
+
+/// A type in which performing operations on vote weights are safe.
 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<Balance, AccountId, C, FS>(
-	mut assignments: Vec<StakedAssignment<AccountId>>,
+pub fn equalize<AccountId>(
+	assignments: &mut Vec<StakedAssignment<AccountId>>,
 	supports: &mut SupportMap<AccountId>,
 	tolerance: ExtendedBalance,
 	iterations: usize,
-	stake_of: FS,
-) where
-	C: Convert<Balance, u64> + Convert<u128, Balance>,
-	for<'r> FS: Fn(&'r AccountId) -> Balance,
-	AccountId: Ord + Clone,
-{
-	// prepare the data for equalise
-	for _i in 0..iterations {
+) -> usize where AccountId: Ord + Clone {
+	if iterations == 0 { return 0; }
+
+	let mut i = 0 ;
+	loop {
 		let mut max_diff = 0;
-
-		for StakedAssignment { who, distribution } in assignments.iter_mut() {
-			let voter_budget = stake_of(&who);
-
-			let diff = do_equalize::<_, _, C>(
+		for assignment in assignments.iter_mut() {
+			let voter_budget = assignment.total();
+			let StakedAssignment { who, distribution } =  assignment;
+			let diff = do_equalize(
 				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 {
-			break;
+		i += 1;
+		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
-	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);
-
+) -> ExtendedBalance where AccountId: Ord + Clone {
 	// 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)])
-	-> Box<dyn Fn(&AccountId) -> Balance>
+pub(crate) fn create_stake_of(stakes: &[(AccountId, VoteWeight)])
+	-> 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
-		.into_iter()
-		.map(|a| {
-			let stake = <TestCurrencyToVote as Convert<Balance, u64>>::convert(stake_of(&a.who)) as ExtendedBalance;
-			a.into_staked(stake, true)
-		}).collect();
+	let mut staked_assignments = assignment_ratio_to_staked(result.assignments, &stake_of);
+	let winners = to_without_backing(result.winners);
 
-	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, _>(
-		staked_assignments,
+	equalize(
+		&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 super::{AccountId, Balance};
+	use crate::{generate_compact_solution_type, VoteWeight};
+	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,