Enable staking equalise. (#2886)

* Enable equalise.

* Bump.

* Line-width

* Add benchmarks for equalise.

* Line-width

* Re-trigger CI.

* Fix test.

* Some nits.

* Rename.

* Bump.

substrate/node/runtime/src/lib.rs

@@ -59,7 +59,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	impl_name: create_runtime_str!("substrate-node"),
 	authoring_version: 10,
 	spec_version: 99,
-	impl_version: 104,
+	impl_version: 105,
 	apis: RUNTIME_API_VERSIONS,
 };
 

substrate/srml/staking/Cargo.toml

@@ -23,8 +23,9 @@ balances = { package = "srml-balances", path = "../balances" }
 rand = "0.6.5"
 
 [features]
+equalize = []
 bench = []
-default = ["std"]
+default = ["std", "equalize"]
 std = [
 	"serde",
 	"safe-mix/std",

substrate/srml/staking/src/benches.rs

@@ -35,7 +35,15 @@ const EDGES: u64 = 2;
 const TO_ELECT: usize = 100;
 const STAKE: u64 = 1000;
 
-fn do_phragmen(b: &mut Bencher, num_vals: u64, num_noms: u64, count: usize, votes_per: u64) {
+fn do_phragmen(
+	b: &mut Bencher,
+	num_vals: u64,
+	num_noms: u64,
+	count: usize,
+	votes_per: u64,
+	eq_iters: usize,
+	eq_tolerance: u128,
+) {
 	with_externalities(&mut ExtBuilder::default().nominate(false).build(), || {
 		assert!(num_vals > votes_per);
 		let rr = |a, b| rand::thread_rng().gen_range(a as usize, b as usize) as u64;
@@ -53,62 +61,135 @@ fn do_phragmen(b: &mut Bencher, num_vals: u64, num_noms: u64, count: usize, vote
 				let mut stashes_to_vote = (1 ..= 2*num_vals)
 					.step_by(2)
 					.map(|ctrl| ctrl + 1)
-					.collect::<Vec<AccountIdType>>();
+					.collect::<Vec<AccountId>>();
 				let votes = (0 .. votes_per)
 					.map(|_| {
 						stashes_to_vote.remove(rr(0, stashes_to_vote.len()) as usize)
 					})
-					.collect::<Vec<AccountIdType>>();
+					.collect::<Vec<AccountId>>();
 				bond_nominator(acc, STAKE + rr(10, 50), votes);
 			});
 
 		b.iter(|| {
-			let _ = phragmen::elect::<Test, _, _, _>(
+			let r = phragmen::elect::<Test, _, _, _>(
 				count,
 				1_usize,
 				<Validators<Test>>::enumerate(),
 				<Nominators<Test>>::enumerate(),
 				Staking::slashable_balance_of
-			);
+			).unwrap();
+
+			// Do the benchmarking with equalize.
+			if eq_iters > 0 {
+				let elected_stashes = r.0;
+				let assignments = r.1;
+
+				let to_balance = |b: ExtendedBalance|
+				<<mock::Test as Trait>::CurrencyToVote as Convert<ExtendedBalance, Balance>>::convert(b);
+				let to_votes = |b: Balance|
+					<<mock::Test as Trait>::CurrencyToVote as Convert<Balance, u64>>::convert(b) as ExtendedBalance;
+				let ratio_of = |b, p| (p as ExtendedBalance).saturating_mul(to_votes(b)) / ACCURACY;
+
+				let assignments_with_stakes = assignments.into_iter().map(|(n, a)|(
+					n,
+					Staking::slashable_balance_of(&n),
+					a.into_iter().map(|(acc, r)| (
+						acc.clone(),
+						r,
+						to_balance(ratio_of(Staking::slashable_balance_of(&n), r)),
+					))
+					.collect::<Vec<Assignment<Test>>>()
+				)).collect::<Vec<(AccountId, Balance, Vec<Assignment<Test>>)>>();
+
+				let mut exposures = <ExpoMap<Test>>::new();
+				elected_stashes
+					.into_iter()
+					.map(|e| (e, Staking::slashable_balance_of(&e)))
+					.for_each(|(e, s)| {
+						let item = Exposure { own: s, total: s, ..Default::default() };
+						exposures.insert(e, item);
+					});
+
+				for (n, _, assignment) in &assignments_with_stakes {
+					for (c, _, s) in assignment {
+						if let Some(expo) = exposures.get_mut(c) {
+							expo.total = expo.total.saturating_add(*s);
+							expo.others.push( IndividualExposure { who: n.clone(), value: *s } );
+						}
+					}
+				}
+
+				let mut assignments_with_votes = assignments_with_stakes.into_iter()
+					.map(|a| (
+						a.0, a.1,
+						a.2.into_iter()
+							.map(|e| (e.0, e.1, to_votes(e.2)))
+							.collect::<Vec<(AccountId, ExtendedBalance, ExtendedBalance)>>()
+					))
+					.collect::<Vec<(
+						AccountId,
+						Balance,
+						Vec<(AccountId, ExtendedBalance, ExtendedBalance)>
+					)>>();
+				equalize::<Test>(&mut assignments_with_votes, &mut exposures, eq_tolerance, eq_iters);
+			}
 		})
 	})
 }
 
 macro_rules! phragmen_benches {
 	($($name:ident: $tup:expr,)*) => {
-    $(
+	$(
         #[bench]
         fn $name(b: &mut Bencher) {
-			let (v, n, t, e) = $tup;
+			let (v, n, t, e, eq_iter, eq_tol) = $tup;
 			println!("");
 			println!(
-				"++ Benchmark: {} Validators // {} Nominators // {} Edges-per-nominator // {} total edges // electing {}",
-				v, n, e, e * n, t
+				r#"
+++ Benchmark: {} Validators // {} Nominators // {} Edges-per-nominator // {} total edges //
+electing {} // Equalize: {} iterations -- {} tolerance"#,
+				v, n, e, e * n, t, eq_iter, eq_tol,
 			);
-			do_phragmen(b, v, n, t, e);
+			do_phragmen(b, v, n, t, e, eq_iter, eq_tol);
         }
     )*
 	}
 }
 
 phragmen_benches! {
-	bench_1_1: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES),
-	bench_1_2: (VALIDATORS*2, NOMINATORS, TO_ELECT, EDGES),
-	bench_1_3: (VALIDATORS*4, NOMINATORS, TO_ELECT, EDGES),
-	bench_1_4: (VALIDATORS*8, NOMINATORS, TO_ELECT, EDGES),
+	bench_1_1: (VALIDATORS, 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_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_3_eq: (VALIDATORS*4, 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),
-	bench_0_2: (VALIDATORS, NOMINATORS, TO_ELECT * 4, EDGES),
-	bench_0_3: (VALIDATORS, NOMINATORS, TO_ELECT * 8, EDGES),
-	bench_0_4: (VALIDATORS, NOMINATORS, TO_ELECT * 16, EDGES),
+	bench_0_1: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES, 0, 0),
+	bench_0_2: (VALIDATORS, NOMINATORS, TO_ELECT * 4, EDGES, 0, 0),
+	bench_0_3: (VALIDATORS, NOMINATORS, TO_ELECT * 8, EDGES, 0, 0),
+	bench_0_4: (VALIDATORS, NOMINATORS, TO_ELECT * 16, EDGES , 0, 0),
+	bench_0_1_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES, 2, 0),
+	bench_0_2_eq: (VALIDATORS, NOMINATORS, TO_ELECT * 4, EDGES, 2, 0),
+	bench_0_3_eq: (VALIDATORS, NOMINATORS, TO_ELECT * 8, EDGES, 2, 0),
+	bench_0_4_eq: (VALIDATORS, NOMINATORS, TO_ELECT * 16, EDGES , 2, 0),
 
-	bench_2_1: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES),
-	bench_2_2: (VALIDATORS, NOMINATORS*2, TO_ELECT, EDGES),
-	bench_2_3: (VALIDATORS, NOMINATORS*4, TO_ELECT, EDGES),
-	bench_2_4: (VALIDATORS, NOMINATORS*8, TO_ELECT, EDGES),
+	bench_2_1: (VALIDATORS, NOMINATORS, 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_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_3_eq: (VALIDATORS, NOMINATORS*4, 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),
-	bench_3_2: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES*2),
-	bench_3_3: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES*4),
-	bench_3_4: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES*8),
+	bench_3_1: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES, 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_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_3_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES*4, 2, 0),
+	bench_3_4_eq: (VALIDATORS, NOMINATORS, TO_ELECT, EDGES*8, 2, 0),
 }

substrate/srml/staking/src/lib.rs

@@ -290,7 +290,7 @@ use primitives::traits::{
 use primitives::{Serialize, Deserialize};
 use system::ensure_signed;
 
-use phragmen::{elect, ACCURACY, ExtendedBalance};
+use phragmen::{elect, ACCURACY, ExtendedBalance, equalize};
 
 const RECENT_OFFLINE_COUNT: usize = 32;
 const DEFAULT_MINIMUM_VALIDATOR_COUNT: u32 = 4;
@@ -1077,7 +1077,7 @@ impl<T: Trait> Module<T> {
 			let ratio_of = |b, p| (p as ExtendedBalance).saturating_mul(to_votes(b)) / ACCURACY;
 
 			// Compute the actual stake from nominator's ratio.
-			let mut assignments_with_stakes = assignments.iter().map(|(n, a)|(
+			let assignments_with_stakes = assignments.iter().map(|(n, a)|(
 				n.clone(),
 				Self::slashable_balance_of(n),
 				a.iter().map(|(acc, r)| (
@@ -1111,17 +1111,22 @@ impl<T: Trait> Module<T> {
 				}
 			}
 
-			// This optimization will most likely be only applied off-chain.
-			let do_equalize = false;
-			if do_equalize {
-				let tolerance = 10 as u128;
-				let iterations = 10 as usize;
-				phragmen::equalize::<T>(
-					&mut assignments_with_stakes,
-					&mut exposures,
-					tolerance,
-					iterations
-				);
+			if cfg!(feature = "equalize") {
+				let tolerance = 0_u128;
+				let iterations = 2_usize;
+				let mut assignments_with_votes = assignments_with_stakes.iter()
+					.map(|a| (
+						a.0.clone(), a.1,
+						a.2.iter()
+							.map(|e| (e.0.clone(), e.1, to_votes(e.2)))
+							.collect::<Vec<(T::AccountId, ExtendedBalance, ExtendedBalance)>>()
+					))
+					.collect::<Vec<(
+						T::AccountId,
+						BalanceOf<T>,
+						Vec<(T::AccountId, ExtendedBalance, ExtendedBalance)>
+					)>>();
+				equalize::<T>(&mut assignments_with_votes, &mut exposures, tolerance, iterations);
 			}
 
 			// Clear Stakers and reduce their slash_count.
@@ -1141,9 +1146,11 @@ impl<T: Trait> Module<T> {
 				}
 				<Stakers<T>>::insert(c.clone(), e.clone());
 			}
+
+			// Update slot stake.
 			<SlotStake<T>>::put(&slot_stake);
 
-			// Set the new validator set.
+			// Set the new validator set in sessions.
 			<CurrentElected<T>>::put(&elected_stashes);
 			let validators = elected_stashes.into_iter()
 				.map(|s| Self::bonded(s).unwrap_or_default())

substrate/srml/staking/src/mock.rs

@@ -22,12 +22,15 @@ use primitives::traits::{IdentityLookup, Convert, OpaqueKeys, OnInitialize};
 use primitives::testing::{Header, UintAuthorityId};
 use substrate_primitives::{H256, Blake2Hasher};
 use runtime_io;
-use srml_support::{impl_outer_origin, parameter_types, assert_ok, traits::Currency};
-use crate::{EraIndex, GenesisConfig, Module, Trait, StakerStatus, ValidatorPrefs, RewardDestination};
+use srml_support::{impl_outer_origin, parameter_types, assert_ok, traits::Currency, EnumerableStorageMap};
+use crate::{EraIndex, GenesisConfig, Module, Trait, StakerStatus,
+	ValidatorPrefs, RewardDestination, Nominators
+};
 
 /// The AccountId alias in this test module.
 pub type AccountId = u64;
 pub type BlockNumber = u64;
+pub type Balance = u64;
 
 /// Simple structure that exposes how u64 currency can be represented as... u64.
 pub struct CurrencyToVoteHandler;
@@ -261,20 +264,52 @@ pub type Session = session::Module<Test>;
 pub type Timestamp = timestamp::Module<Test>;
 pub type Staking = Module<Test>;
 
-pub fn check_exposure(acc: u64) {
-	let expo = Staking::stakers(&acc);
-	assert_eq!(expo.total as u128, expo.own as u128 + expo.others.iter().map(|e| e.value as u128).sum::<u128>());
-}
-
 pub fn check_exposure_all() {
 	Staking::current_elected().into_iter().for_each(|acc| check_exposure(acc));
 }
 
-pub fn assert_total_expo(acc: u64, val: u64) {
-	let expo = Staking::stakers(&acc);
+pub fn check_nominator_all() {
+	<Nominators<Test>>::enumerate().for_each(|(acc, _)| check_nominator_exposure(acc));
+}
+
+/// Check for each selected validator: expo.total = Sum(expo.other) + expo.own
+pub fn check_exposure(stash: u64) {
+	assert_is_stash(stash);
+	let expo = Staking::stakers(&stash);
+	assert_eq!(
+		expo.total as u128, expo.own as u128 + expo.others.iter().map(|e| e.value as u128).sum::<u128>(),
+		"wrong total exposure for {:?}: {:?}", stash, expo,
+	);
+}
+
+/// Check that for each nominator: slashable_balance > sum(used_balance)
+/// Note: we might not consume all of a nominator's balance, but we MUST NOT over spend it.
+pub fn check_nominator_exposure(stash: u64) {
+	assert_is_stash(stash);
+	let mut sum = 0;
+	Staking::current_elected()
+		.iter()
+		.map(|v| Staking::stakers(v))
+		.for_each(|e| e.others.iter()
+			.filter(|i| i.who == stash)
+			.for_each(|i| sum += i.value));
+	let nominator_stake = Staking::slashable_balance_of(&stash);
+	// a nominator cannot over-spend.
+	assert!(
+		nominator_stake >= sum,
+		"failed: Nominator({}) stake({}) >= sum divided({})", stash, nominator_stake, sum,
+	);
+}
+
+pub fn assert_total_expo(stash: u64, val: u64) {
+	let expo = Staking::stakers(&stash);
 	assert_eq!(expo.total, val);
 }
 
+pub fn assert_is_stash(acc: u64) {
+	assert!(Staking::bonded(&acc).is_some(), "Not a stash.");
+}
+
 pub fn bond_validator(acc: u64, val: u64) {
 	// a = controller
 	// a + 1 = stash

substrate/srml/staking/src/phragmen.rs

@@ -19,7 +19,7 @@
 use rstd::{prelude::*, collections::btree_map::BTreeMap};
 use primitives::{PerU128};
 use primitives::traits::{Zero, Convert, Saturating};
-use crate::{BalanceOf, Assignment, RawAssignment, ExpoMap, Trait, ValidatorPrefs};
+use crate::{BalanceOf, RawAssignment, ExpoMap, Trait, ValidatorPrefs, IndividualExposure};
 
 type Fraction = PerU128;
 /// Wrapper around the type used as the _safe_ wrapper around a `balance`.
@@ -275,7 +275,7 @@ pub fn elect<T: Trait + 'static, FV, FN, FS>(
 ///
 /// No value is returned from the function and the `expo_map` parameter is updated.
 pub fn equalize<T: Trait + 'static>(
-	assignments: &mut Vec<(T::AccountId, BalanceOf<T>, Vec<Assignment<T>>)>,
+	assignments: &mut Vec<(T::AccountId, BalanceOf<T>, Vec<(T::AccountId, ExtendedBalance, ExtendedBalance)>)>,
 	expo_map: &mut ExpoMap<T>,
 	tolerance: ExtendedBalance,
 	iterations: usize,
@@ -297,19 +297,19 @@ pub fn equalize<T: Trait + 'static>(
 fn do_equalize<T: Trait + 'static>(
 	nominator: &T::AccountId,
 	budget_balance: BalanceOf<T>,
-	elected_edges_balance: &mut Vec<Assignment<T>>,
+	elected_edges: &mut Vec<(T::AccountId, ExtendedBalance, ExtendedBalance)>,
 	expo_map: &mut ExpoMap<T>,
 	tolerance: ExtendedBalance
 ) -> ExtendedBalance {
-	let to_votes = |b: BalanceOf<T>| <T::CurrencyToVote as Convert<BalanceOf<T>, u64>>::convert(b) as ExtendedBalance;
-	let to_balance = |v: ExtendedBalance| <T::CurrencyToVote as Convert<ExtendedBalance, BalanceOf<T>>>::convert(v);
+	let to_votes = |b: BalanceOf<T>|
+		<T::CurrencyToVote as Convert<BalanceOf<T>, u64>>::convert(b) as ExtendedBalance;
+	let to_balance = |v: ExtendedBalance|
+		<T::CurrencyToVote as Convert<ExtendedBalance, BalanceOf<T>>>::convert(v);
 	let budget = to_votes(budget_balance);
 
-	// Convert all stakes to extended. Result is Vec<(Acc, Ratio, Balance)>
-	let mut elected_edges = elected_edges_balance
-		.into_iter()
-		.map(|e| (e.0.clone(), e.1, to_votes(e.2)))
-		.collect::<Vec<(T::AccountId, ExtendedBalance, ExtendedBalance)>>();
+	// Nothing to do. This nominator had nothing useful.
+	// Defensive only. Assignment list should always be populated.
+	if elected_edges.is_empty() { return 0; }
 
 	let stake_used = elected_edges
 		.iter()
@@ -350,18 +350,20 @@ fn do_equalize<T: Trait + 'static>(
 	elected_edges.iter_mut().for_each(|e| {
 		if let Some(expo) = expo_map.get_mut(&e.0) {
 			expo.total = expo.total.saturating_sub(to_balance(e.2));
+			expo.others.retain(|i_expo| i_expo.who != *nominator);
 		}
 		e.2 = 0;
 	});
 
-	elected_edges.sort_unstable_by_key(|e| e.2);
+	elected_edges.sort_unstable_by_key(|e|
+		if let Some(e) = expo_map.get(&e.0) { e.total } else { Zero::zero() }
+	);
 
 	let mut cumulative_stake: ExtendedBalance = 0;
 	let mut last_index = elected_edges.len() - 1;
 	elected_edges.iter_mut().enumerate().for_each(|(idx, e)| {
 		if let Some(expo) = expo_map.get_mut(&e.0) {
 			let stake: ExtendedBalance = to_votes(expo.total);
-
 			let stake_mul = stake.saturating_mul(idx as ExtendedBalance);
 			let stake_sub = stake_mul.saturating_sub(cumulative_stake);
 			if stake_sub > budget {
@@ -383,14 +385,9 @@ fn do_equalize<T: Trait + 'static>(
 				.saturating_add(last_stake)
 				.saturating_sub(to_votes(expo.total));
 			expo.total = expo.total.saturating_add(to_balance(e.2));
-			if let Some(i_expo) = expo.others.iter_mut().find(|i| i.who == nominator.clone()) {
-				i_expo.value = to_balance(e.2);
-			}
+			expo.others.push(IndividualExposure { who: nominator.clone(), value: to_balance(e.2)});
 		}
 	});
 
-	// Store back the individual edge weights.
-	elected_edges.iter().enumerate().for_each(|(idx, e)| elected_edges_balance[idx].2 = to_balance(e.2));
-
 	difference
 }

substrate/srml/staking/src/tests.rs

@@ -52,10 +52,30 @@ fn basic_setup_works() {
 		assert_eq!(Staking::ledger(100), Some(StakingLedger { stash: 101, total: 500, active: 500, unlocking: vec![] }));
 		assert_eq!(Staking::nominators(101), vec![11, 21]);
 
-		// Account 10 is exposed by 1000 * balance_factor from their own stash in account 11 + the default nominator vote
-		assert_eq!(Staking::stakers(11), Exposure { total: 1125, own: 1000, others: vec![ IndividualExposure { who: 101, value: 125 }] });
-		// Account 20 is exposed by 1000 * balance_factor from their own stash in account 21 + the default nominator vote
-		assert_eq!(Staking::stakers(21), Exposure { total: 1375, own: 1000, others: vec![ IndividualExposure { who: 101, value: 375 }] });
+		if cfg!(feature = "equalize") {
+			assert_eq!(
+				Staking::stakers(11),
+				Exposure { total: 1250, own: 1000, others: vec![ IndividualExposure { who: 101, value: 250 }] }
+			);
+			assert_eq!(
+				Staking::stakers(21),
+				Exposure { total: 1250, own: 1000, others: vec![ IndividualExposure { who: 101, value: 250 }] }
+			);
+			// initial slot_stake
+			assert_eq!(Staking::slot_stake(),  1250);
+		} else {
+			assert_eq!(
+				Staking::stakers(11),
+				Exposure { total: 1125, own: 1000, others: vec![ IndividualExposure { who: 101, value: 125 }] }
+			);
+			assert_eq!(
+				Staking::stakers(21),
+				Exposure { total: 1375, own: 1000, others: vec![ IndividualExposure { who: 101, value: 375 }] }
+			);
+			// initial slot_stake
+			assert_eq!(Staking::slot_stake(),  1125);
+		}
+
 
 		// The number of validators required.
 		assert_eq!(Staking::validator_count(), 2);
@@ -67,8 +87,6 @@ fn basic_setup_works() {
 		// initial rewards
 		assert_eq!(Staking::current_session_reward(), 10);
 
-		// initial slot_stake
-		assert_eq!(Staking::slot_stake(),  1125); // Naive
 
 		// initial slash_count of validators
 		assert_eq!(Staking::slash_count(&11), 0);
@@ -76,6 +94,7 @@ fn basic_setup_works() {
 
 		// All exposures must be correct.
 		check_exposure_all();
+		check_nominator_all();
 	});
 }
 
@@ -199,7 +218,10 @@ fn offline_grace_should_delay_slashing() {
 
 		// Check unstake_threshold is 3 (default)
 		let default_unstake_threshold = 3;
-		assert_eq!(Staking::validators(&11), ValidatorPrefs { unstake_threshold: default_unstake_threshold, validator_payment: 0 });
+		assert_eq!(
+			Staking::validators(&11),
+			ValidatorPrefs { unstake_threshold: default_unstake_threshold, validator_payment: 0 }
+		);
 
 		// Check slash count is zero
 		assert_eq!(Staking::slash_count(&11), 0);
@@ -499,7 +521,10 @@ fn staking_should_work() {
 		assert_eq_uvec!(Session::validators(), vec![20, 10]);
 
 		// Note: the stashed value of 4 is still lock
-		assert_eq!(Staking::ledger(&4), Some(StakingLedger { stash: 3, total: 1500, active: 1500, unlocking: vec![] }));
+		assert_eq!(
+			Staking::ledger(&4),
+			Some(StakingLedger { stash: 3, total: 1500, active: 1500, unlocking: vec![] })
+		);
 		// e.g. it cannot spend more than 500 that it has free from the total 2000
 		assert_noop!(Balances::reserve(&3, 501), "account liquidity restrictions prevent withdrawal");
 		assert_ok!(Balances::reserve(&3, 409));
@@ -528,6 +553,7 @@ fn less_than_needed_candidates_works() {
 		assert_eq!(Staking::stakers(20).others.len(), 0);
 		assert_eq!(Staking::stakers(30).others.len(), 0);
 		check_exposure_all();
+		check_nominator_all();
 	});
 }
 
@@ -597,7 +623,6 @@ fn nominating_and_rewards_should_work() {
 	// 10  with stake  400.0 20  with stake  600.0 30  with stake  0
 	// 4  has load  0.0005555555555555556 and supported
 	// 10  with stake  600.0 20  with stake  400.0 40  with stake  0.0
-
 	with_externalities(&mut ExtBuilder::default()
 		.nominate(false)
 		.validator_pool(true)
@@ -641,18 +666,57 @@ fn nominating_and_rewards_should_work() {
 
 		// ------ check the staked value of all parties.
 
-		// total expo of 10, with 1200 coming from nominators (externals), according to phragmen.
-		assert_eq!(Staking::stakers(11).own, 1000);
-		assert_eq!(Staking::stakers(11).total, 1000 + 800);
-		// 2 and 4 supported 10, each with stake 600, according to phragmen.
-		assert_eq!(Staking::stakers(11).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(), vec![400, 400]);
-		assert_eq!(Staking::stakers(11).others.iter().map(|e| e.who).collect::<Vec<BalanceOf<Test>>>(), vec![3, 1]);
-		// total expo of 20, with 500 coming from nominators (externals), according to phragmen.
-		assert_eq!(Staking::stakers(21).own, 1000);
-		assert_eq!(Staking::stakers(21).total, 1000 + 1198);
-		// 2 and 4 supported 20, each with stake 250, according to phragmen.
-		assert_eq!(Staking::stakers(21).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(), vec![599, 599]);
-		assert_eq!(Staking::stakers(21).others.iter().map(|e| e.who).collect::<Vec<BalanceOf<Test>>>(), vec![3, 1]);
+		if cfg!(feature = "equalize") {
+			// total expo of 10, with 1200 coming from nominators (externals), according to phragmen.
+			assert_eq!(Staking::stakers(11).own, 1000);
+			assert_eq!(Staking::stakers(11).total, 1000 + 999);
+			// 2 and 4 supported 10, each with stake 600, according to phragmen.
+			assert_eq!(
+				Staking::stakers(11).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(),
+				vec![599, 400]
+			);
+			assert_eq!(
+				Staking::stakers(11).others.iter().map(|e| e.who).collect::<Vec<u64>>(),
+				vec![3, 1]
+			);
+			// total expo of 20, with 500 coming from nominators (externals), according to phragmen.
+			assert_eq!(Staking::stakers(21).own, 1000);
+			assert_eq!(Staking::stakers(21).total, 1000 + 999);
+			// 2 and 4 supported 20, each with stake 250, according to phragmen.
+			assert_eq!(
+				Staking::stakers(21).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(),
+				vec![400, 599]
+			);
+			assert_eq!(
+				Staking::stakers(21).others.iter().map(|e| e.who).collect::<Vec<u64>>(),
+				vec![3, 1]
+			);
+		} else {
+			// total expo of 10, with 1200 coming from nominators (externals), according to phragmen.
+			assert_eq!(Staking::stakers(11).own, 1000);
+			assert_eq!(Staking::stakers(11).total, 1000 + 800);
+			// 2 and 4 supported 10, each with stake 600, according to phragmen.
+			assert_eq!(
+				Staking::stakers(11).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(),
+				vec![400, 400]
+			);
+			assert_eq!(
+				Staking::stakers(11).others.iter().map(|e| e.who).collect::<Vec<u64>>(),
+				vec![3, 1]
+			);
+			// total expo of 20, with 500 coming from nominators (externals), according to phragmen.
+			assert_eq!(Staking::stakers(21).own, 1000);
+			assert_eq!(Staking::stakers(21).total, 1000 + 1198);
+			// 2 and 4 supported 20, each with stake 250, according to phragmen.
+			assert_eq!(
+				Staking::stakers(21).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(),
+				vec![599, 599]
+			);
+			assert_eq!(
+				Staking::stakers(21).others.iter().map(|e| e.who).collect::<Vec<u64>>(),
+				vec![3, 1]
+			);
+		}
 
 		// They are not chosen anymore
 		assert_eq!(Staking::stakers(31).total, 0);
@@ -662,28 +726,35 @@ fn nominating_and_rewards_should_work() {
 		start_era(2);
 		// next session reward.
 		let new_session_reward = Staking::session_reward() * 3 * Staking::slot_stake();
-		// nothing else will happen, era ends and rewards are paid again,
-		// it is expected that nominators will also be paid. See below
 
-		// Approximation resulting from Perbill conversion
-		let approximation = 1;
-		// Nominator 2: has [400/1800 ~ 2/9 from 10] + [600/2200 ~ 3/11 from 20]'s reward. ==> 2/9 + 3/11
-		assert_eq!(
-			Balances::total_balance(&2),
-			initial_balance + (2*new_session_reward/9 + 3*new_session_reward/11) - 1 - approximation
-		);
-		// Nominator 4: has [400/1800 ~ 2/9 from 10] + [600/2200 ~ 3/11 from 20]'s reward. ==> 2/9 + 3/11
-		assert_eq!(
-			Balances::total_balance(&4),
-			initial_balance + (2*new_session_reward/9 + 3*new_session_reward/11) - 1 - approximation
-		);
+		let approximation = 3;
+		if cfg!(feature = "equalize") {
+			// Both have: has [400/2000 ~ 1/5 from 10] + [600/2000 ~ 3/10 from 20]'s reward. ==> 1/5 + 3/10 = 1/2
+			assert_eq!(Balances::total_balance(&2), initial_balance + new_session_reward/2 - approximation);
+			assert_eq!(Balances::total_balance(&4), initial_balance + new_session_reward/2 - approximation);
+			// Rest for validators.
+			assert_eq!(Balances::total_balance(&10), initial_balance + new_session_reward/2 + 1);
+			assert_eq!(Balances::total_balance(&20), initial_balance + new_session_reward/2 + 1);
+		} else {
+			// Nominator 2: has [400/1800 ~ 2/9 from 10] + [600/2200 ~ 3/11 from 20]'s reward. ==> 2/9 + 3/11
+			assert_eq!(
+				Balances::total_balance(&2),
+				initial_balance + (2*new_session_reward/9 + 3*new_session_reward/11) - approximation
+			);
+			// Nominator 4: has [400/1800 ~ 2/9 from 10] + [600/2200 ~ 3/11 from 20]'s reward. ==> 2/9 + 3/11
+			assert_eq!(
+				Balances::total_balance(&4),
+				initial_balance + (2*new_session_reward/9 + 3*new_session_reward/11) - approximation
+			);
 
-		// 10 got 800 / 1800 external stake => 8/18 =? 4/9 => Validator's share = 5/9
-		assert_eq!(Balances::total_balance(&10), initial_balance + 5*new_session_reward/9 - approximation);
-		// 10 got 1200 / 2200 external stake => 12/22 =? 6/11 => Validator's share = 5/11
-		assert_eq!(Balances::total_balance(&20), initial_balance + 5*new_session_reward/11+ 2);
+			// 10 got 800 / 1800 external stake => 8/18 =? 4/9 => Validator's share = 5/9
+			assert_eq!(Balances::total_balance(&10), initial_balance + 5*new_session_reward/9);
+			// 10 got 1200 / 2200 external stake => 12/22 =? 6/11 => Validator's share = 5/11
+			assert_eq!(Balances::total_balance(&20), initial_balance + 5*new_session_reward/11 + 2);
+		}
 
 		check_exposure_all();
+		check_nominator_all();
 	});
 }
 
@@ -730,6 +801,7 @@ fn nominators_also_get_slashed() {
 		assert_eq!(Balances::total_balance(&10), initial_balance + 30 - validator_slash);
 		assert_eq!(Balances::total_balance(&2), initial_balance - nominator_slash);
 		check_exposure_all();
+		check_nominator_all();
 		// Because slashing happened.
 		assert!(is_disabled(10));
 	});
@@ -746,9 +818,15 @@ fn double_staking_should_fail() {
 		|| {
 			let arbitrary_value = 5;
 			// 2 = controller, 1 stashed => ok
-			assert_ok!(Staking::bond(Origin::signed(1), 2, arbitrary_value, RewardDestination::default()));
+			assert_ok!(
+				Staking::bond(Origin::signed(1), 2, arbitrary_value,
+				RewardDestination::default())
+			);
 			// 4 = not used so far, 1 stashed => not allowed.
-			assert_noop!(Staking::bond(Origin::signed(1), 4, arbitrary_value, RewardDestination::default()), "stash already bonded");
+			assert_noop!(
+				Staking::bond(Origin::signed(1), 4, arbitrary_value,
+				RewardDestination::default()), "stash already bonded"
+			);
 			// 1 = stashed => attempting to nominate should fail.
 			assert_noop!(Staking::nominate(Origin::signed(1), vec![1]), "not a controller");
 			// 2 = controller  => nominating should work.
@@ -1033,6 +1111,7 @@ fn validator_payment_prefs_work() {
 		assert_eq!(Balances::total_balance(&2), 500 + shared_cut/2);
 
 		check_exposure_all();
+		check_nominator_all();
 	});
 
 }
@@ -1248,6 +1327,7 @@ fn slot_stake_is_least_staked_validator_and_exposure_defines_maximum_punishment(
 		assert_eq!(Balances::free_balance(&11), 1000 + 30 - 51 /*5% of 1030*/ * 8 /*2**3*/);
 
 		check_exposure_all();
+		check_nominator_all();
 	});
 }
 
@@ -1411,8 +1491,6 @@ fn phragmen_poc_works() {
 	// 10  with stake  166.66666666666674 20  with stake  333.33333333333326 30  with stake  0
 	// 4  has load  0.00075 and supported
 	// 10  with stake  333.3333333333333 20  with stake  166.66666666666666 40  with stake  0.0
-
-
 	with_externalities(&mut ExtBuilder::default()
 		.nominate(false)
 		.validator_pool(true)
@@ -1445,21 +1523,58 @@ fn phragmen_poc_works() {
 
 		assert_eq_uvec!(Session::validators(), vec![20, 10]);
 
-		// with stake 1666 and 1333 respectively
 		assert_eq!(Staking::stakers(11).own, 1000);
-		assert_eq!(Staking::stakers(11).total, 1000 + 332);
 		assert_eq!(Staking::stakers(21).own, 1000);
-		assert_eq!(Staking::stakers(21).total, 1000 + 666);
+
+		if cfg!(feature = "equalize") {
+			assert_eq!(Staking::stakers(11).total, 1000 + 499);
+			assert_eq!(Staking::stakers(21).total, 1000 + 499);
+		} else {
+			assert_eq!(Staking::stakers(11).total, 1000 + 332);
+			assert_eq!(Staking::stakers(21).total, 1000 + 666);
+		}
 
 		// Nominator's stake distribution.
-		assert_eq!(Staking::stakers(11).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(), vec![166, 166]);
-		assert_eq!(Staking::stakers(11).others.iter().map(|e| e.value).sum::<BalanceOf<Test>>(), 332);
 		assert_eq!(Staking::stakers(11).others.iter().map(|e| e.who).collect::<Vec<BalanceOf<Test>>>(), vec![3, 1]);
-
-		assert_eq!(Staking::stakers(21).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(), vec![333, 333]);
-		assert_eq!(Staking::stakers(21).others.iter().map(|e| e.value).sum::<BalanceOf<Test>>(), 666);
 		assert_eq!(Staking::stakers(21).others.iter().map(|e| e.who).collect::<Vec<BalanceOf<Test>>>(), vec![3, 1]);
+
+		if cfg!(feature = "equalize") {
+			assert_eq_uvec!(
+				Staking::stakers(11).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(),
+				vec![333, 166]
+			);
+			assert_eq!(
+				Staking::stakers(11).others.iter().map(|e| e.value).sum::<BalanceOf<Test>>(),
+				499
+			);
+			assert_eq_uvec!(
+				Staking::stakers(21).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(),
+				vec![333, 166]
+			);
+			assert_eq!(
+				Staking::stakers(21).others.iter().map(|e| e.value).sum::<BalanceOf<Test>>(),
+				499
+			);
+		} else {
+			assert_eq_uvec!(
+				Staking::stakers(11).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(),
+				vec![166, 166]
+			);
+			assert_eq!(
+				Staking::stakers(11).others.iter().map(|e| e.value).sum::<BalanceOf<Test>>(),
+				332
+			);
+			assert_eq_uvec!(
+				Staking::stakers(21).others.iter().map(|e| e.value).collect::<Vec<BalanceOf<Test>>>(),
+				vec![333, 333]
+			);
+			assert_eq!(
+				Staking::stakers(21).others.iter().map(|e| e.value).sum::<BalanceOf<Test>>(),
+				666
+			);
+		}
 		check_exposure_all();
+		check_nominator_all();
 	});
 }
 
@@ -1510,6 +1625,7 @@ fn phragmen_poc_2_works() {
 			(1, vec![(11, 4294967296)]),
 		]);
 		check_exposure_all();
+		check_nominator_all();
 	})
 }
 
@@ -1564,9 +1680,9 @@ fn switching_roles() {
 		assert_ok!(Staking::validate(Origin::signed(2), ValidatorPrefs::default()));
 		// new stakes:
 		// 10: 1000 self vote
-		// 20: 1000 self vote + 500 vote
+		// 20: 1000 self vote + 250 vote
 		// 6 : 1000 self vote
-		// 2 : 2000 self vote + 500 vote.
+		// 2 : 2000 self vote + 250 vote.
 		// Winners: 20 and 2
 
 		System::set_block_number(4);
@@ -1581,7 +1697,9 @@ fn switching_roles() {
 		System::set_block_number(6);
 		Session::on_initialize(System::block_number());
 		assert_eq_uvec!(Session::validators(), vec![2, 20]);
+
 		check_exposure_all();
+		check_nominator_all();
 	});
 }
 
@@ -1717,12 +1835,13 @@ fn bond_with_little_staked_value_bounded_by_slot_stake() {
 		// same for 10
 		assert_eq!(Balances::free_balance(&10), initial_balance_10 + 30 + reward.max(3));
 		check_exposure_all();
+		check_nominator_all();
 	});
 }
 
 
+#[cfg(feature = "equalize")]
 #[test]
-#[ignore] // Enable this once post-processing is on.
 fn phragmen_linear_worse_case_equalize() {
 	with_externalities(&mut ExtBuilder::default()
 		.nominate(false)
@@ -1730,7 +1849,6 @@ fn phragmen_linear_worse_case_equalize() {
 		.fair(true)
 		.build(),
 	|| {
-		for i in &[10, 20, 30, 40] { assert_ok!(Staking::set_payee(Origin::signed(*i), RewardDestination::Controller)); }
 
 		bond_validator(50, 1000);
 		bond_validator(60, 1000);
@@ -1744,32 +1862,27 @@ fn phragmen_linear_worse_case_equalize() {
 		bond_nominator(120, 1000, vec![51, 61]);
 		bond_nominator(130, 1000, vec![61, 71]);
 
+		for i in &[10, 20, 30, 40, 50, 60, 70] {
+			assert_ok!(Staking::set_payee(Origin::signed(*i), RewardDestination::Controller));
+		}
+
 		assert_eq_uvec!(Session::validators(), vec![40, 30]);
 		assert_ok!(Staking::set_validator_count(7));
 
-		System::set_block_number(1);
-		Session::on_initialize(System::block_number());
+		start_era(1);
 
 		assert_eq_uvec!(Session::validators(), vec![10, 60, 40, 20, 50, 30, 70]);
 
-		// Sequential Phragmén with post processing gives
-		// 10  is elected with stake  3000.0 and score  0.00025
-		// 20  is elected with stake  2017.7421569824219 and score  0.0005277777777777777
-		// 30  is elected with stake  2008.8712884829595 and score  0.0003333333333333333
-		// 40  is elected with stake  2000.0001049958742 and score  0.0005555555555555556
-		// 50  is elected with stake  2000.0001049958742 and score  0.0003333333333333333
-		// 60  is elected with stake  1991.128921508789 and score  0.0004444444444444444
-		// 70  is elected with stake  1982.2574230340813 and score  0.0007222222222222222
+		assert_eq!(Staking::stakers(11).total, 3000);
+		assert_eq!(Staking::stakers(21).total, 2254);
+		assert_eq!(Staking::stakers(31).total, 2254);
+		assert_eq!(Staking::stakers(41).total, 1926);
+		assert_eq!(Staking::stakers(51).total, 1871);
+		assert_eq!(Staking::stakers(61).total, 1892);
+		assert_eq!(Staking::stakers(71).total, 1799);
 
 		check_exposure_all();
-
-		assert_eq!(Staking::stakers(11).total, 3000);
-		assert_eq!(Staking::stakers(21).total, 2209);
-		assert_eq!(Staking::stakers(31).total, 2027);
-		assert_eq!(Staking::stakers(41).total, 2010);
-		assert_eq!(Staking::stakers(51).total, 2010);
-		assert_eq!(Staking::stakers(61).total, 1998);
-		assert_eq!(Staking::stakers(71).total, 1983);
+		check_nominator_all();
 	})
 }
 
@@ -1790,6 +1903,7 @@ fn phragmen_chooses_correct_number_of_validators() {
 
 		assert_eq!(Session::validators().len(), 1);
 		check_exposure_all();
+		check_nominator_all();
 	})
 }
 
@@ -1809,6 +1923,7 @@ fn phragmen_score_should_be_accurate_on_large_stakes() {
 
 		assert_eq!(Session::validators(), vec![4, 2]);
 		check_exposure_all();
+		check_nominator_all();
 	})
 }
 
@@ -1929,32 +2044,8 @@ fn phragmen_large_scale_test() {
 
 		start_era(1);
 
-		// For manual inspection
-		println!("Validators are {:?}", Session::validators());
-		println!("Validators are {:#?}",
-			Session::validators()
-				.iter()
-				.map(|v| (v.clone(), Staking::stakers(v+1)))
-				.collect::<Vec<(u64, Exposure<u64, u64>)>>()
-		);
-
-		// Each exposure => total == own + sum(others)
 		check_exposure_all();
-
-		// aside from some error, stake must be divided correctly
-		let individual_expo_sum: u128 = Session::validators()
-			.iter()
-			.map(|v| Staking::stakers(v+1))
-			.fold(0u128, |s, v| if v.others.len() > 0 { s + v.others[0].value as u128 } else { s });
-		assert!(
-			990000000000000000 - individual_expo_sum < 100,
-			format!(
-				"Nominator stake = {} / SUM(individual expo) = {} / diff = {}",
-				990000000000000000u64,
-				individual_expo_sum,
-				990000000000000000 - individual_expo_sum
-			)
-		);
+		check_nominator_all();
 	})
 }
 
@@ -1980,6 +2071,7 @@ fn phragmen_large_scale_test_2() {
 
 		// Each exposure => total == own + sum(others)
 		check_exposure_all();
+		check_nominator_all();
 
 		assert_total_expo(3, nom_budget / 2 + c_budget);
 		assert_total_expo(5, nom_budget / 2 + c_budget);