Reformat Validator Election (#2406)

* Add index caching to election

* Initial draft of the new phragmen API.

* Port post-processing to the new API.

* Fix tests and accuracy.

* Final fixes.

* Unify convert namings.

* Remove todo.

* Some typos.

* Bump.

* Add extended balance type doc.

* A bit more sophisticated weight compensation.

* Fix review feedbacks.

* Bump.

* Final updates

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: 73,
-	impl_version: 73,
+	impl_version: 74,
 	apis: RUNTIME_API_VERSIONS,
 };
 

substrate/srml/staking/src/lib.rs

@@ -237,7 +237,7 @@
 
 #[cfg(feature = "std")]
 use runtime_io::with_storage;
-use rstd::{prelude::*, result};
+use rstd::{prelude::*, result, collections::btree_map::BTreeMap};
 use parity_codec::{HasCompact, Encode, Decode};
 use srml_support::{StorageValue, StorageMap, EnumerableStorageMap, dispatch::Result};
 use srml_support::{decl_module, decl_event, decl_storage, ensure};
@@ -247,7 +247,7 @@ use srml_support::traits::{
 };
 use session::OnSessionChange;
 use primitives::Perbill;
-use primitives::traits::{Convert, Zero, One, As, StaticLookup, CheckedSub, Saturating, Bounded};
+use primitives::traits::{Convert, Zero, One, As, StaticLookup, CheckedSub, CheckedShl, Saturating, Bounded};
 #[cfg(feature = "std")]
 use primitives::{Serialize, Deserialize};
 use system::ensure_signed;
@@ -256,7 +256,7 @@ mod mock;
 mod tests;
 mod phragmen;
 
-use phragmen::{elect, ElectionConfig};
+use phragmen::{elect, ACCURACY, ExtendedBalance};
 
 const RECENT_OFFLINE_COUNT: usize = 32;
 const DEFAULT_MINIMUM_VALIDATOR_COUNT: u32 = 4;
@@ -394,12 +394,19 @@ type BalanceOf<T> = <<T as Trait>::Currency as Currency<<T as system::Trait>::Ac
 type PositiveImbalanceOf<T> = <<T as Trait>::Currency as Currency<<T as system::Trait>::AccountId>>::PositiveImbalance;
 type NegativeImbalanceOf<T> = <<T as Trait>::Currency as Currency<<T as system::Trait>::AccountId>>::NegativeImbalance;
 
+type RawAssignment<T> = (<T as system::Trait>::AccountId, ExtendedBalance);
+type Assignment<T> = (<T as system::Trait>::AccountId, ExtendedBalance, BalanceOf<T>);
+type ExpoMap<T> = BTreeMap::<<T as system::Trait>::AccountId, Exposure<<T as system::Trait>::AccountId, BalanceOf<T>>>;
+
 pub trait Trait: system::Trait + session::Trait {
 	/// The staking balance.
 	type Currency: LockableCurrency<Self::AccountId, Moment=Self::BlockNumber>;
 
 	/// Convert a balance into a number used for election calculation.
 	/// This must fit into a `u64` but is allowed to be sensibly lossy.
+	/// TODO: #1377
+	/// The backward convert should be removed as the new Phragmen API returns ratio.
+	/// The post-processing needs it but will be moved to off-chain.
 	type CurrencyToVote: Convert<BalanceOf<Self>, u64> + Convert<u128, BalanceOf<Self>>;
 
 	/// Some tokens minted.
@@ -927,20 +934,71 @@ impl<T: Trait> Module<T> {
 	///
 	/// Returns the new `SlotStake` value.
 	fn select_validators() -> BalanceOf<T> {
-		let maybe_elected_candidates = elect::<T, _, _, _>(
+		let maybe_elected_set = elect::<T, _, _, _>(
 			Self::validator_count() as usize,
 			Self::minimum_validator_count().max(1) as usize,
 			<Validators<T>>::enumerate(),
 			<Nominators<T>>::enumerate(),
 			Self::slashable_balance_of,
-			ElectionConfig::<BalanceOf<T>> {
-				equalize: false,
-				tolerance: <BalanceOf<T>>::sa(10 as u64),
-				iterations: 10,
-			}
 		);
 
-		if let Some(elected_candidates) = maybe_elected_candidates {
+		if let Some(elected_set) = maybe_elected_set {
+			let elected_stashes = elected_set.0;
+			let assignments = elected_set.1;
+
+			// helper closure.
+			let to_balance = |b: ExtendedBalance| <T::CurrencyToVote as Convert<ExtendedBalance, BalanceOf<T>>>::convert(b);
+			let to_votes = |b: BalanceOf<T>| <T::CurrencyToVote as Convert<BalanceOf<T>, u64>>::convert(b) as ExtendedBalance;
+
+			// The return value of this is safe to be converted to u64.
+			// The original balance, `b` is within the scope of u64. It is just extended to u128
+			// to be properly multiplied by a ratio, which will lead to another value
+			// less than u64 for sure. The result can then be safely passed to `to_balance`.
+			// For now the backward convert is used. A simple `TryFrom<u64>` is also safe.
+			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)|(
+				n.clone(),
+				Self::slashable_balance_of(n),
+				a.iter().map(|(acc, r)| (
+					acc.clone(),
+					*r,
+					to_balance(ratio_of(Self::slashable_balance_of(n), *r)),
+				))
+				.collect::<Vec<Assignment<T>>>()
+			)).collect::<Vec<(T::AccountId, BalanceOf<T>, Vec<Assignment<T>>)>>();
+
+			// update elected candidate exposures.
+			let mut exposures = <ExpoMap<T>>::new();
+			elected_stashes
+				.iter()
+				.map(|e| (e, Self::slashable_balance_of(e)))
+				.for_each(|(e, s)| {
+					exposures.insert(e.clone(), Exposure { own: s, total: s, ..Default::default() });
+				});
+
+			for (n, _, assignment) in &assignments_with_stakes {
+				for (c, _, s) in assignment {
+					if let Some(expo) = exposures.get_mut(c) {
+						// NOTE: simple example where this saturates:
+						// candidate with max_value stake. 1 nominator with max_value stake.
+						// Nuked. Sadly there is not much that we can do about this.
+						// See this test: phragmen_should_not_overflow_xxx()
+						expo.total = expo.total.saturating_add(*s);
+						expo.others.push( IndividualExposure { who: n.clone(), value: *s } );
+					}
+				}
+			}
+
+			// This optimization will most likely be only applied off-chain.
+			let do_equalise = false;
+			if do_equalise {
+				let tolerance = 10 as u128;
+				let iterations = 10 as usize;
+				phragmen::equalize::<T>(&mut assignments_with_stakes, &mut exposures, tolerance, iterations);
+			}
+
 			// Clear Stakers and reduce their slash_count.
 			for v in Self::current_elected().iter() {
 				<Stakers<T>>::remove(v);
@@ -951,17 +1009,16 @@ impl<T: Trait> Module<T> {
 			}
 
 			// Populate Stakers and figure out the minimum stake behind a slot.
-			let mut slot_stake = elected_candidates[0].exposure.total;
-			for c in &elected_candidates {
-				if c.exposure.total < slot_stake {
-					slot_stake = c.exposure.total;
+			let mut slot_stake = BalanceOf::<T>::max_value();
+			for (c, e) in exposures.iter() {
+				if e.total < slot_stake {
+					slot_stake = e.total;
 				}
-				<Stakers<T>>::insert(c.who.clone(), c.exposure.clone());
+				<Stakers<T>>::insert(c.clone(), e.clone());
 			}
 			<SlotStake<T>>::put(&slot_stake);
 
 			// Set the new validator set.
-			let elected_stashes = elected_candidates.into_iter().map(|i| i.who).collect::<Vec<_>>();
 			<CurrentElected<T>>::put(&elected_stashes);
 			<session::Module<T>>::set_validators(
 				&elected_stashes.into_iter().map(|s| Self::bonded(s).unwrap_or_default()).collect::<Vec<_>>()
@@ -974,6 +1031,7 @@ impl<T: Trait> Module<T> {
 			// We should probably disable all functionality except for block production
 			// and let the chain keep producing blocks until we can decide on a sufficiently
 			// substantial set.
+			// TODO: #2494
 			Self::slot_stake()
 		}
 	}
@@ -983,8 +1041,6 @@ impl<T: Trait> Module<T> {
 	///
 	/// NOTE: This is called with the controller (not the stash) account id.
 	pub fn on_offline_validator(controller: T::AccountId, count: usize) {
-		use primitives::traits::CheckedShl;
-
 		if let Some(l) = Self::ledger(&controller) {
 			let stash = l.stash;
 

substrate/srml/staking/src/phragmen.rs

@@ -16,43 +16,36 @@
 
 //! Rust implementation of the Phragmén election algorithm.
 
-use rstd::prelude::*;
-use primitives::PerU128;
-use primitives::traits::{Zero, Saturating, Convert};
-use parity_codec::{HasCompact, Encode, Decode};
-use crate::{Exposure, BalanceOf, Trait, ValidatorPrefs, IndividualExposure};
+use rstd::{prelude::*, collections::btree_map::BTreeMap};
+use primitives::{PerU128};
+use primitives::traits::{Zero, Convert, Saturating};
+use parity_codec::{Encode, Decode};
+use crate::{BalanceOf, Assignment, RawAssignment, ExpoMap, Trait, ValidatorPrefs};
 
 type Fraction = PerU128;
-type ExtendedBalance = u128;
-const SCALE_FACTOR_64: u128 = u64::max_value() as u128 + 1;
+/// Wrapper around the type used as the _safe_ wrapper around a `balance`.
+pub type ExtendedBalance = u128;
 
-/// Configure the behavior of the Phragmen election.
-/// Might be deprecated.
-pub struct ElectionConfig<Balance: HasCompact> {
-	/// Perform equalize?.
-	pub equalize: bool,
-	/// Number of equalize iterations.
-	pub iterations: usize,
-	/// Tolerance of max change per equalize iteration.
-	pub tolerance: Balance,
-}
+// this is only used while creating the candidate score. Due to reasons explained below
+// The more accurate this is, the less likely we choose a wrong candidate.
+const SCALE_FACTOR: ExtendedBalance = u32::max_value() as ExtendedBalance + 1;
+/// These are used to expose a fixed accuracy to the caller function. The bigger they are,
+/// the more accurate we get, but the more likely it is for us to overflow. The case of overflow
+/// is handled but accuracy will be lost. 32 or 16 are reasonable values.
+pub const ACCURACY: ExtendedBalance = u32::max_value() as ExtendedBalance + 1;
 
 /// Wrapper around validation candidates some metadata.
 #[derive(Clone, Encode, Decode, Default)]
 #[cfg_attr(feature = "std", derive(Debug))]
-pub struct Candidate<AccountId, Balance: HasCompact> {
+pub struct Candidate<AccountId> {
 	/// The validator's account
 	pub who: AccountId,
-	/// Exposure struct, holding info about the value that the validator has in stake.
-	pub exposure: Exposure<AccountId, Balance>,
 	/// Intermediary value used to sort candidates.
 	pub score: Fraction,
 	/// Accumulator of the stake of this candidate based on received votes.
 	approval_stake: ExtendedBalance,
 	/// Flag for being elected.
 	elected: bool,
-	/// This is most often equal to `Exposure.total` but not always. Needed for [`equalize`]
-	backing_stake: ExtendedBalance
 }
 
 /// Wrapper around the nomination info of a single nominator for a group of validators.
@@ -77,14 +70,10 @@ pub struct Edge<AccountId> {
 	who: AccountId,
 	/// Load of this vote.
 	load: Fraction,
-	/// Final backing stake of this vote.
-	backing_stake: ExtendedBalance,
-	/// Index of the candidate stored in the 'candidates' vector
+	/// Equal to `edge.load / nom.load`. Stored only to be used with post-processing.
+	ratio: ExtendedBalance,
+	/// Index of the candidate stored in the 'candidates' vector.
 	candidate_index: usize,
-	/// Index of the candidate stored in the 'elected_candidates' vector. Used only with equalize.
-	elected_idx: usize,
-	/// Indicates if this edge is a vote for an elected candidate. Used only with equalize.
-	elected: bool,
 }
 
 /// Perform election based on Phragmén algorithm.
@@ -93,52 +82,55 @@ pub struct Edge<AccountId> {
 ///
 /// Returns an Option of elected candidates, if election is performed.
 /// Returns None if not enough candidates exist.
+///
+/// The returned Option is a tuple consisting of:
+///   - The list of elected candidates.
+///   - The list of nominators and their associated vote weights.
 pub fn elect<T: Trait + 'static, FV, FN, FS>(
 	validator_count: usize,
 	minimum_validator_count: usize,
 	validator_iter: FV,
 	nominator_iter: FN,
 	stash_of: FS,
-	config: ElectionConfig<BalanceOf<T>>,
-) -> Option<Vec<Candidate<T::AccountId, BalanceOf<T>>>> where
+) -> Option<(Vec<T::AccountId>, Vec<(T::AccountId, Vec<RawAssignment<T>>)>)> where
 	FV: Iterator<Item=(T::AccountId, ValidatorPrefs<BalanceOf<T>>)>,
 	FN: Iterator<Item=(T::AccountId, Vec<T::AccountId>)>,
 	for <'r> FS: Fn(&'r T::AccountId) -> BalanceOf<T>,
 {
-	let into_currency = |b: BalanceOf<T>| <T::CurrencyToVote as Convert<BalanceOf<T>, u64>>::convert(b) as ExtendedBalance;
-	let into_votes = |b: ExtendedBalance| <T::CurrencyToVote as Convert<ExtendedBalance, BalanceOf<T>>>::convert(b);
-	let mut elected_candidates;
+	let to_votes = |b: BalanceOf<T>| <T::CurrencyToVote as Convert<BalanceOf<T>, u64>>::convert(b) as ExtendedBalance;
+
+	// return structures
+	let mut elected_candidates: Vec<T::AccountId>;
+	let mut assigned: Vec<(T::AccountId, Vec<RawAssignment<T>>)>;
+	let mut c_idx_cache = BTreeMap::<T::AccountId, usize>::new();
 
 	// 1- Pre-process candidates and place them in a container, optimisation and add phantom votes.
 	// Candidates who have 0 stake => have no votes or all null-votes. Kick them out not.
 	let mut nominators: Vec<Nominator<T::AccountId>> = Vec::with_capacity(validator_iter.size_hint().0 + nominator_iter.size_hint().0);
 	let mut candidates = validator_iter.map(|(who, _)| {
 			let stash_balance = stash_of(&who);
-			Candidate {
-				who,
-				exposure: Exposure { total: stash_balance, own: stash_balance, others: vec![] },
-				..Default::default()
-			}
+			(Candidate { who, ..Default::default() }, stash_balance)
 		})
-		.filter_map(|mut c| {
-			c.approval_stake += into_currency(c.exposure.total);
+		.filter_map(|(mut c, s)| {
+			c.approval_stake += to_votes(s);
 			if c.approval_stake.is_zero() {
 				None
 			} else {
-				Some(c)
+				Some((c, s))
 			}
 		})
 		.enumerate()
-		.map(|(idx, c)| {
+		.map(|(idx, (c, s))| {
 			nominators.push(Nominator {
 				who: c.who.clone(),
 				edges: vec![ Edge { who: c.who.clone(), candidate_index: idx, ..Default::default() }],
-				budget: into_currency(c.exposure.total),
+				budget: to_votes(s),
 				load: Fraction::zero(),
 			});
+			c_idx_cache.insert(c.who.clone(), idx);
 			c
 		})
-		.collect::<Vec<Candidate<T::AccountId, BalanceOf<T>>>>();
+		.collect::<Vec<Candidate<T::AccountId>>>();
 
 	// 2- Collect the nominators with the associated votes.
 	// Also collect approval stake along the way.
@@ -146,17 +138,17 @@ pub fn elect<T: Trait + 'static, FV, FN, FS>(
 		let nominator_stake = stash_of(&who);
 		let mut edges: Vec<Edge<T::AccountId>> = Vec::with_capacity(nominees.len());
 		for n in &nominees {
-			if let Some(idx) = candidates.iter_mut().position(|i| i.who == *n) {
-				candidates[idx].approval_stake = candidates[idx].approval_stake
-					.saturating_add(into_currency(nominator_stake));
-				edges.push(Edge { who: n.clone(), candidate_index: idx, ..Default::default() });
-			}
+			if let Some(idx) = c_idx_cache.get(n) {
+				// This candidate is valid + already cached.
+				candidates[*idx].approval_stake = candidates[*idx].approval_stake
+					.saturating_add(to_votes(nominator_stake));
+				edges.push(Edge { who: n.clone(), candidate_index: *idx, ..Default::default() });
+			} // else {} would be wrong votes. We don't really care about it.
 		}
-
 		Nominator {
 			who,
 			edges: edges,
-			budget: into_currency(nominator_stake),
+			budget: to_votes(nominator_stake),
 			load: Fraction::zero(),
 		}
 	}));
@@ -166,6 +158,7 @@ pub fn elect<T: Trait + 'static, FV, FN, FS>(
 		let validator_count = validator_count.min(candidates.len());
 
 		elected_candidates = Vec::with_capacity(validator_count);
+		assigned = Vec::with_capacity(validator_count);
 		// Main election loop
 		for _round in 0..validator_count {
 			// Loop 1: initialize score
@@ -179,12 +172,14 @@ pub fn elect<T: Trait + 'static, FV, FN, FS>(
 				for e in &n.edges {
 					let c = &mut candidates[e.candidate_index];
 					if !c.elected && !c.approval_stake.is_zero() {
+						// Basic fixed-point shifting by 32.
+						// `n.budget.saturating_mul(SCALE_FACTOR)` will never saturate
+						// since n.budget cannot exceed u64,despite being stored in u128. yet,
+						// `*n.load / SCALE_FACTOR` might collapse to zero. Hence, 32 or 16 bits are better scale factors.
+						// Note that left-associativity in operators precedence is crucially important here.
 						let temp =
-							// Basic fixed-point shifting by 64.
-							// This will never saturate since n.budget cannot exceed u64,despite being stored in u128.
-							// Note that left-associativity in operators precedence is crucially important here.
-							n.budget.saturating_mul(SCALE_FACTOR_64) / c.approval_stake
-							* (*n.load / SCALE_FACTOR_64);
+							n.budget.saturating_mul(SCALE_FACTOR) / c.approval_stake
+							* (*n.load / SCALE_FACTOR);
 						c.score = Fraction::from_max_value((*c.score).saturating_add(temp));
 					}
 				}
@@ -207,162 +202,195 @@ pub fn elect<T: Trait + 'static, FV, FN, FS>(
 					}
 				}
 
-				elected_candidates.push(winner.clone());
+				elected_candidates.push(winner.who.clone());
 			} else {
 				break
 			}
-		}
-		// end of all rounds
+		} // end of all rounds
 
 		// 4.1- Update backing stake of candidates and nominators
 		for n in &mut nominators {
+			let mut assignment = (n.who.clone(), vec![]);
 			for e in &mut n.edges {
-				// if the target of this vote is among the winners, otherwise let go.
-				if let Some(c) = elected_candidates.iter_mut().find(|c| c.who == e.who) {
-					e.elected = true;
-					e.backing_stake = n.budget.saturating_mul(*e.load) / n.load.max(1);
-					c.backing_stake = c.backing_stake.saturating_add(e.backing_stake);
-					if c.who != n.who {
-						// Only update the exposure if this vote is from some other account.
-						c.exposure.total = c.exposure.total.saturating_add(into_votes(e.backing_stake));
-						c.exposure.others.push(
-							IndividualExposure { who: n.who.clone(), value: into_votes(e.backing_stake) }
-						);
+				if let Some(c) = elected_candidates.iter().find(|c| **c == e.who) {
+					if *c != n.who {
+						let ratio = {
+							// Full support. No need to calculate.
+							if *n.load == *e.load { ACCURACY }
+							else {
+								// This should not saturate. Safest is to just check
+								if let Some(r) = ACCURACY.checked_mul(*e.load) {
+									r / n.load.max(1)
+								} else {
+									// Just a simple trick.
+									*e.load / (n.load.max(1) / ACCURACY)
+								}
+							}
+						};
+						e.ratio = ratio;
+						assignment.1.push((e.who.clone(), ratio));
 					}
 				}
 			}
-		}
 
-		// Optionally perform equalize post-processing.
-		if config.equalize {
-			let tolerance = config.tolerance;
-			let equalize_iterations = config.iterations;
+			if assignment.1.len() > 0 {
+				// To ensure an assertion indicating: no stake from the nominator going to waste,
+				// we add a minimal post-processing to equally assign all of the leftover stake ratios.
+				let vote_count = assignment.1.len() as ExtendedBalance;
+				let l = assignment.1.len();
+				let sum = assignment.1.iter().map(|a| a.1).sum::<ExtendedBalance>();
+				let diff = ACCURACY.checked_sub(sum).unwrap_or(0);
+				let diff_per_vote= diff / vote_count;
 
-			// Fix indexes
-			nominators.iter_mut().for_each(|n| {
-				n.edges.iter_mut().for_each(|e| {
-					if let Some(idx) = elected_candidates.iter().position(|c| c.who == e.who) {
-						e.elected_idx = idx;
+				if diff_per_vote > 0 {
+					for i in 0..l {
+						assignment.1[i%l].1 =
+							assignment.1[i%l].1
+							.saturating_add(diff_per_vote);
 					}
-				});
-			});
-
-			for _i in 0..equalize_iterations {
-				let mut max_diff = <BalanceOf<T>>::zero();
-				nominators.iter_mut().for_each(|mut n| {
-					let diff = equalize::<T>(&mut n, &mut elected_candidates, tolerance);
-					if diff > max_diff {
-						max_diff = diff;
-					}
-				});
-				if max_diff < tolerance {
-					break;
 				}
+
+				// `remainder` is set to be less than maximum votes of a nominator (currently 16).
+				// safe to cast it to usize.
+				let remainder = diff - diff_per_vote * vote_count;
+				for i in 0..remainder as usize {
+					assignment.1[i%l].1 =
+						assignment.1[i%l].1
+						.saturating_add(1);
+				}
+				assigned.push(assignment);
 			}
 		}
+
 	} else {
 		// if we have less than minimum, use the previous validator set.
 		return None
 	}
-	Some(elected_candidates)
+	Some((elected_candidates, assigned))
 }
 
 /// Performs equalize post-processing to the output of the election algorithm
 /// This function mutates the input parameters, most noticeably it updates the exposure of
 /// the elected candidates.
-/// The return value is to tolerance at which the function has stopped.
+///
+/// No value is returned from the function and the `expo_map` parameter is updated.
 pub fn equalize<T: Trait + 'static>(
-	nominator: &mut Nominator<T::AccountId>,
-	elected_candidates: &mut Vec<Candidate<T::AccountId, BalanceOf<T>>>,
-	_tolerance: BalanceOf<T>
-) -> BalanceOf<T> {
-	let into_currency = |b: BalanceOf<T>|    <T::CurrencyToVote as Convert<BalanceOf<T>, u64>>::convert(b) as ExtendedBalance;
-	let into_votes = |b: ExtendedBalance| <T::CurrencyToVote as Convert<ExtendedBalance, BalanceOf<T>>>::convert(b);
-	let tolerance = into_currency(_tolerance);
-
-	let mut elected_edges = nominator.edges
-		.iter_mut()
-		.filter(|e| e.elected)
-		.collect::<Vec<&mut Edge<T::AccountId>>>();
-
-	if elected_edges.len() == 0 {
-		return <BalanceOf<T>>::zero();
+	assignments: &mut Vec<(T::AccountId, BalanceOf<T>, Vec<Assignment<T>>)>,
+	expo_map: &mut ExpoMap<T>,
+	tolerance: ExtendedBalance,
+	iterations: usize,
+) {
+	for _i in 0..iterations {
+		let mut max_diff = 0;
+		assignments.iter_mut().for_each(|(n, budget, assignment)| {
+			let diff = do_equalize::<T>(&n, *budget, assignment, expo_map, tolerance);
+			if diff > max_diff {
+				max_diff = diff;
+			}
+		});
+		if max_diff < tolerance {
+			break;
+		}
 	}
+}
+
+fn do_equalize<T: Trait + 'static>(
+	nominator: &T::AccountId,
+	budget_balance: BalanceOf<T>,
+	elected_edges_balance: &mut Vec<Assignment<T>>,
+	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 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)>>();
 
 	let stake_used = elected_edges
 		.iter()
-		.fold(0, |s, e| s.saturating_add(e.backing_stake));
-	let backed_stakes = elected_edges
+		.fold(0 as ExtendedBalance, |s, e| s.saturating_add(e.2));
+
+	let backed_stakes_iter = elected_edges
 		.iter()
-		.map(|e| elected_candidates[e.elected_idx].backing_stake)
-		.collect::<Vec<ExtendedBalance>>();
+		.filter_map(|e| expo_map.get(&e.0))
+		.map(|e| to_votes(e.total));
+
 	let backing_backed_stake = elected_edges
 		.iter()
-		.filter(|e| e.backing_stake > 0)
-		.map(|e| elected_candidates[e.elected_idx].backing_stake)
+		.filter(|e| e.2 > 0)
+		.filter_map(|e| expo_map.get(&e.0))
+		.map(|e| to_votes(e.total))
 		.collect::<Vec<ExtendedBalance>>();
 
 	let mut difference;
 	if backing_backed_stake.len() > 0 {
-		let max_stake = *backing_backed_stake
+		let max_stake = backing_backed_stake
 			.iter()
 			.max()
 			.expect("vector with positive length will have a max; qed");
-		let min_stake = *backed_stakes
-			.iter()
+		let min_stake = backed_stakes_iter
 			.min()
-			.expect("vector with positive length will have a min; qed");
+			.expect("iterator with positive length will have a min; qed");
+
 		difference = max_stake.saturating_sub(min_stake);
-		difference = difference.saturating_add(nominator.budget.saturating_sub(stake_used));
+		difference = difference.saturating_add(budget.saturating_sub(stake_used));
 		if difference < tolerance {
-			return into_votes(difference);
+			return difference;
 		}
 	} else {
-		difference = nominator.budget;
+		difference = budget;
 	}
 
 	// Undo updates to exposure
 	elected_edges.iter_mut().for_each(|e| {
-		// NOTE: no assertions in the runtime, but this should nonetheless be indicative.
-		//assert_eq!(elected_candidates[e.elected_idx].who, e.who);
-		elected_candidates[e.elected_idx].backing_stake -= e.backing_stake;
-		elected_candidates[e.elected_idx].exposure.total -= into_votes(e.backing_stake);
-		e.backing_stake = 0;
+		if let Some(expo) = expo_map.get_mut(&e.0) {
+			expo.total = expo.total.saturating_sub(to_balance(e.2));
+		}
+		e.2 = 0;
 	});
 
-	elected_edges.sort_unstable_by_key(|e| elected_candidates[e.elected_idx].backing_stake);
+	elected_edges.sort_unstable_by_key(|e| e.2);
 
 	let mut cumulative_stake: ExtendedBalance = 0;
 	let mut last_index = elected_edges.len() - 1;
-	let budget = nominator.budget;
 	elected_edges.iter_mut().enumerate().for_each(|(idx, e)| {
-		let stake = elected_candidates[e.elected_idx].backing_stake;
+		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 {
-			last_index = idx.checked_sub(1).unwrap_or(0);
-			return
+			let stake_mul = stake.saturating_mul(idx as ExtendedBalance);
+			let stake_sub = stake_mul.saturating_sub(cumulative_stake);
+			if stake_sub > budget {
+				last_index = idx.checked_sub(1).unwrap_or(0);
+				return
+			}
+			cumulative_stake = cumulative_stake.saturating_add(stake);
 		}
-		cumulative_stake = cumulative_stake.saturating_add(stake);
 	});
 
-	let last_stake = elected_candidates[elected_edges[last_index].elected_idx].backing_stake;
+	let last_stake = elected_edges[last_index].2;
 	let split_ways = last_index + 1;
-	let excess = nominator.budget
+	let excess = budget
 		.saturating_add(cumulative_stake)
 		.saturating_sub(last_stake.saturating_mul(split_ways as ExtendedBalance));
-	let nominator_address = nominator.who.clone();
 	elected_edges.iter_mut().take(split_ways).for_each(|e| {
-		let c = &mut elected_candidates[e.elected_idx];
-		e.backing_stake = (excess / split_ways as ExtendedBalance)
-			.saturating_add(last_stake)
-			.saturating_sub(c.backing_stake);
-		c.exposure.total = c.exposure.total.saturating_add(into_votes(e.backing_stake));
-		c.backing_stake = c.backing_stake.saturating_add(e.backing_stake);
-		if let Some(i_expo) = c.exposure.others.iter_mut().find(|i| i.who == nominator_address) {
-			i_expo.value = into_votes(e.backing_stake);
+		if let Some(expo) = expo_map.get_mut(&e.0) {
+			e.2 = (excess / split_ways as ExtendedBalance)
+				.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);
+			}
 		}
 	});
-	into_votes(difference)
+
+	// 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

@@ -21,11 +21,45 @@
 use super::*;
 use runtime_io::with_externalities;
 use phragmen;
-use primitives::PerU128;
 use srml_support::{assert_ok, assert_noop, assert_eq_uvec, EnumerableStorageMap};
 use mock::{Balances, Session, Staking, System, Timestamp, Test, ExtBuilder, Origin};
 use srml_support::traits::{Currency, ReservableCurrency};
 
+#[inline]
+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>());
+}
+
+#[inline]
+fn check_exposure_all() {
+	Staking::current_elected().into_iter().for_each(|acc| check_exposure(acc));
+}
+
+#[inline]
+fn assert_total_expo(acc: u64, val: u64) {
+	let expo = Staking::stakers(&acc);
+	assert_eq!(expo.total, val);
+}
+
+#[inline]
+fn bond_validator(acc: u64, val: u64) {
+	// a = controller
+	// a + 1 = stash
+	let _ = Balances::make_free_balance_be(&(acc+1), val);
+	assert_ok!(Staking::bond(Origin::signed(acc+1), acc, val, RewardDestination::Controller));
+	assert_ok!(Staking::validate(Origin::signed(acc), ValidatorPrefs::default()));
+}
+
+#[inline]
+fn bond_nominator(acc: u64, val: u64, target: Vec<u64>) {
+	// a = controller
+	// a + 1 = stash
+	let _ = Balances::make_free_balance_be(&(acc+1), val);
+	assert_ok!(Staking::bond(Origin::signed(acc+1), acc, val, RewardDestination::Controller));
+	assert_ok!(Staking::nominate(Origin::signed(acc), target));
+}
+
 #[test]
 fn basic_setup_works() {
 	// Verifies initial conditions of mock
@@ -55,7 +89,7 @@ fn basic_setup_works() {
 		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: 1124, own: 1000, others: vec![ IndividualExposure { who: 101, value: 124 }] });
+		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 }] });
 
@@ -70,11 +104,14 @@ fn basic_setup_works() {
 		assert_eq!(Staking::current_session_reward(), 10);
 
 		// initial slot_stake
-		assert_eq!(Staking::slot_stake(),  1124); // Naive
+		assert_eq!(Staking::slot_stake(),  1125); // Naive
 
 		// initial slash_count of validators
 		assert_eq!(Staking::slash_count(&11), 0);
 		assert_eq!(Staking::slash_count(&21), 0);
+
+		// All exposures must be correct.
+		check_exposure_all();
 	});
 }
 
@@ -97,6 +134,31 @@ fn no_offline_should_work() {
 	});
 }
 
+#[test]
+fn change_controller_works() {
+	with_externalities(&mut ExtBuilder::default().build(),
+	|| {
+		assert_eq!(Staking::bonded(&11), Some(10));
+
+		assert!(<Validators<Test>>::enumerate().map(|(c, _)| c).collect::<Vec<u64>>().contains(&11));
+		// 10 can control 11 who is initially a validator.
+		assert_ok!(Staking::chill(Origin::signed(10)));
+		assert!(!<Validators<Test>>::enumerate().map(|(c, _)| c).collect::<Vec<u64>>().contains(&11));
+
+		assert_ok!(Staking::set_controller(Origin::signed(11), 5));
+
+		System::set_block_number(1);
+		Session::check_rotate_session(System::block_number());
+		assert_eq!(Staking::current_era(), 1);
+
+		assert_noop!(
+			Staking::validate(Origin::signed(10), ValidatorPrefs::default()),
+			"not a controller"
+		);
+		assert_ok!(Staking::validate(Origin::signed(5), ValidatorPrefs::default()));
+	})
+}
+
 #[test]
 fn invulnerability_should_work() {
 	// Test that users can be invulnerable from slashing and being kicked
@@ -282,7 +344,7 @@ fn rewards_should_work() {
 		.sessions_per_era(3)
 	.build(),
 	|| {
-		let delay = 0;
+		let delay = 1;
 		// this test is only in the scope of one era. Since this variable changes
 		// at the last block/new era, we'll save it.
 		let session_reward = 10;
@@ -305,7 +367,6 @@ fn rewards_should_work() {
 		assert_eq!(Balances::total_balance(&2), 500);
 
 		// add a dummy nominator.
-		// NOTE: this nominator is being added 'manually'. a Further test (nomination_and_reward..) will add it via '.nominate()'
 		<Stakers<Test>>::insert(&11, Exposure {
 			own: 500, // equal division indicates that the reward will be equally divided among validator and nominator.
 			total: 1000,
@@ -316,8 +377,7 @@ fn rewards_should_work() {
 		assert_eq!(Staking::payee(2), RewardDestination::Stash);
 		assert_eq!(Staking::payee(11), RewardDestination::Controller);
 
-		let mut block = 3;
-		// Block 3 => Session 1 => Era 0
+		let mut block = 3; // Block 3 => Session 1 => Era 0
 		System::set_block_number(block);
 		Timestamp::set_timestamp(block*5);	// on time.
 		Session::check_rotate_session(System::block_number());
@@ -354,16 +414,16 @@ fn rewards_should_work() {
 
 #[test]
 fn multi_era_reward_should_work() {
-	// should check that:
+	// Should check that:
 	// The value of current_session_reward is set at the end of each era, based on
-	// slot_stake and session_reward. Check and verify this.
+	// slot_stake and session_reward.
 	with_externalities(&mut ExtBuilder::default()
 		.session_length(3)
 		.sessions_per_era(3)
 		.nominate(false)
 		.build(),
 	|| {
-		let delay = 0;
+		let delay = 1;
 		let session_reward = 10;
 
 		// This is set by the test config builder.
@@ -378,12 +438,12 @@ fn multi_era_reward_should_work() {
 		let mut block = 3;
 		// Block 3 => Session 1 => Era 0
 		System::set_block_number(block);
-		Timestamp::set_timestamp(block*5);	// on time.
+		Timestamp::set_timestamp(block*5);
 		Session::check_rotate_session(System::block_number());
 		assert_eq!(Staking::current_era(), 0);
 		assert_eq!(Session::current_index(), 1);
 
-		// session triggered: the reward value stashed should be 10 -- defined in ExtBuilder genesis.
+		// session triggered: the reward value stashed should be 10
 		assert_eq!(Staking::current_session_reward(), session_reward);
 		assert_eq!(Staking::current_era_reward(), session_reward);
 
@@ -413,14 +473,14 @@ fn multi_era_reward_should_work() {
 		assert_eq!(Staking::current_session_reward(), new_session_reward);
 
 		// fast forward to next era:
-		block=12;System::set_block_number(block);Timestamp::set_timestamp(block*5);Session::check_rotate_session(System::block_number());
-		block=15;System::set_block_number(block);Timestamp::set_timestamp(block*5);Session::check_rotate_session(System::block_number());
+		block=12; System::set_block_number(block);Timestamp::set_timestamp(block*5);Session::check_rotate_session(System::block_number());
+		block=15; System::set_block_number(block);Timestamp::set_timestamp(block*5);Session::check_rotate_session(System::block_number());
 
 		// intermediate test.
 		assert_eq!(Staking::current_era_reward(), 2*new_session_reward);
 
 		// new era is triggered here.
-		block=18;System::set_block_number(block);Timestamp::set_timestamp(block*5);Session::check_rotate_session(System::block_number());
+		block=18; System::set_block_number(block);Timestamp::set_timestamp(block*5);Session::check_rotate_session(System::block_number());
 
 		// pay time
 		assert_eq!(Balances::total_balance(&10), 3*new_session_reward + recorded_balance);
@@ -515,31 +575,28 @@ fn staking_should_work() {
 
 #[test]
 fn less_than_needed_candidates_works() {
-	// Test the situation where the number of validators are less than `ValidatorCount` but more than <MinValidators>
-	// The expected behavior is to choose all the candidates that have some vote.
 	with_externalities(&mut ExtBuilder::default()
 		.minimum_validator_count(1)
 		.validator_count(4)
 		.nominate(false)
 		.build(),
 	|| {
-		assert_eq!(Staking::era_length(), 1);
 		assert_eq!(Staking::validator_count(), 4);
 		assert_eq!(Staking::minimum_validator_count(), 1);
+		assert_eq_uvec!(Session::validators(), vec![30, 20, 10]);
 
-		// 10 and 20 are now valid candidates.
-		// trigger era
 		System::set_block_number(1);
 		Session::check_rotate_session(System::block_number());
 		assert_eq!(Staking::current_era(), 1);
 
-		// both validators will be chosen again. NO election algorithm is even executed.
+		// Previous set is selected. NO election algorithm is even executed.
 		assert_eq_uvec!(Session::validators(), vec![30, 20, 10]);
 
 		// But the exposure is updated in a simple way. No external votes exists. This is purely self-vote.
-		assert_eq!(Staking::stakers(10).others.iter().map(|e| e.who).collect::<Vec<BalanceOf<Test>>>(), vec![]);
-		assert_eq!(Staking::stakers(20).others.iter().map(|e| e.who).collect::<Vec<BalanceOf<Test>>>(), vec![]);
-		assert_eq!(Staking::stakers(30).others.iter().map(|e| e.who).collect::<Vec<BalanceOf<Test>>>(), vec![]);
+		assert_eq!(Staking::stakers(10).others.len(), 0);
+		assert_eq!(Staking::stakers(20).others.len(), 0);
+		assert_eq!(Staking::stakers(30).others.len(), 0);
+		check_exposure_all();
 	});
 }
 
@@ -560,13 +617,15 @@ fn no_candidate_emergency_condition() {
 		// initial validators
 		assert_eq_uvec!(Session::validators(), vec![10, 20, 30, 40]);
 
+		let _ = Staking::chill(Origin::signed(10));
+
 		// trigger era
 		System::set_block_number(1);
 		Session::check_rotate_session(System::block_number());
-		assert_eq!(Staking::current_era(), 1);
 
-		// No one nominates => no one has a proper vote => no change
+		// Previous ones are elected. chill is invalidates. TODO: #2494
 		assert_eq_uvec!(Session::validators(), vec![10, 20, 30, 40]);
+		assert_eq!(Staking::current_elected().len(), 0);
 	});
 }
 
@@ -656,15 +715,15 @@ fn nominating_and_rewards_should_work() {
 
 		// 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 + 798);
+		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![399, 399]);
+		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 + 1200);
+		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![600, 600]);
+		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]);
 
 		// They are not chosen anymore
@@ -685,9 +744,11 @@ fn nominating_and_rewards_should_work() {
 		assert_eq!(Balances::total_balance(&4), initial_balance + (2*new_session_reward/9 + 3*new_session_reward/11));
 
 		// 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 + 2) ;
+		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);
+
+		check_exposure_all();
 	});
 }
 
@@ -701,7 +762,7 @@ fn nominators_also_get_slashed() {
 		assert_eq!(Staking::offline_slash_grace(), 0);
 		// Account 10 has not been reported offline
 		assert_eq!(Staking::slash_count(&10), 0);
-	<OfflineSlash<Test>>::put(Perbill::from_percent(12));
+		<OfflineSlash<Test>>::put(Perbill::from_percent(12));
 
 		// Set payee to controller
 		assert_ok!(Staking::set_payee(Origin::signed(10), RewardDestination::Controller));
@@ -735,6 +796,7 @@ fn nominators_also_get_slashed() {
 		// initial + first era reward + slash
 		assert_eq!(Balances::total_balance(&10), initial_balance + 10 - validator_slash);
 		assert_eq!(Balances::total_balance(&2), initial_balance - nominator_slash);
+		check_exposure_all();
 		// Because slashing happened.
 		assert!(Staking::forcing_new_era().is_some());
 	});
@@ -1050,6 +1112,8 @@ fn validator_payment_prefs_work() {
 		assert_eq!(Balances::total_balance(&10), 1);
 		// Rest of the reward will be shared and paid to the nominator in stake.
 		assert_eq!(Balances::total_balance(&2), 500 + shared_cut/2);
+
+		check_exposure_all();
 	});
 
 }
@@ -1220,6 +1284,8 @@ fn slot_stake_is_least_staked_validator_and_exposure_defines_maximum_punishment(
 		assert_eq!(Staking::slash_count(&11), 4);
 		// check the balance of 10 (slash will be deducted from free balance.)
 		assert_eq!(Balances::free_balance(&11), 1000 + 10 - 50 /*5% of 1000*/ * 8 /*2**3*/);
+
+		check_exposure_all();
 	});
 }
 
@@ -1432,11 +1498,12 @@ fn phragmen_poc_works() {
 		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]);
+		check_exposure_all();
 	});
 }
 
 #[test]
-fn phragmen_election_works_with_post_processing() {
+fn phragmen_poc_2_works() {
 	// tests the encapsulated phragmen::elect function.
 	// Votes  [
 	// 	('10', 1000, ['10']),
@@ -1448,32 +1515,6 @@ fn phragmen_election_works_with_post_processing() {
 	// Sequential Phragmén gives
 	// 10  is elected with stake  1705.7377049180327 and score  0.0004878048780487805
 	// 30  is elected with stake  1344.2622950819673 and score  0.0007439024390243903
-
-	// 10  has load  0.0004878048780487805 and supported
-	// 10  with stake  1000.0
-	// 20  has load  0 and supported
-	// 20  with stake  0
-	// 30  has load  0.0007439024390243903 and supported
-	// 30  with stake  1000.0
-	// 2  has load  0.0004878048780487805 and supported
-	// 10  with stake  50.0 20  with stake  0.0
-	// 4  has load  0.0007439024390243903 and supported
-	// 10  with stake  655.7377049180328 30  with stake  344.26229508196724
-
-	// Sequential Phragmén with post processing gives
-	// 10  is elected with stake  1525.0 and score  0.0004878048780487805
-	// 30  is elected with stake  1525.0 and score  0.0007439024390243903
-
-	// 10  has load  0.0004878048780487805 and supported
-	// 10  with stake  1000.0
-	// 20  has load  0 and supported
-	// 20  with stake  0
-	// 30  has load  0.0007439024390243903 and supported
-	// 30  with stake  1000.0
-	// 2  has load  0.0004878048780487805 and supported
-	// 10  with stake  50.0 20  with stake  0.0
-	// 4  has load  0.0007439024390243903 and supported
-	// 10  with stake  475.0 30  with stake  525.0
 	with_externalities(&mut ExtBuilder::default().nominate(false).build(), || {
 		// initial setup of 10 and 20, both validators
 		assert_eq_uvec!(Session::validators(), vec![20, 10]);
@@ -1497,30 +1538,17 @@ fn phragmen_election_works_with_post_processing() {
 			<Validators<Test>>::enumerate(),
 			<Nominators<Test>>::enumerate(),
 			Staking::slashable_balance_of,
-			ElectionConfig::<BalanceOf<Test>> {
-				equalize: true,
-				tolerance: <BalanceOf<Test>>::sa(10 as u64),
-				iterations: 10,
-			}
 		);
 
-		let winners = winners.unwrap();
+		let (winners, assignment) = winners.unwrap();
 
 		// 10 and 30 must be the winners
-		assert_eq!(winners.iter().map(|w| w.who).collect::<Vec<BalanceOf<Test>>>(), vec![11, 31]);
-
-		let winner_10 = winners.iter().filter(|w| w.who == 11).nth(0).unwrap();
-		let winner_30 = winners.iter().filter(|w| w.who == 31).nth(0).unwrap();
-
-		// Check exposures
-		assert_eq!(winner_10.exposure.total, 1000 + 525);
-		assert_eq!(winner_10.score, PerU128::from_max_value(165991398498018762665060784113057664));
-		assert_eq!(winner_10.exposure.others[0].value, 475);
-		assert_eq!(winner_10.exposure.others[1].value, 50);
-
-		assert_eq!(winner_30.exposure.total, 1000 + 525);
-		assert_eq!(winner_30.score, PerU128::from_max_value(253136882709478612992230401826229321));
-		assert_eq!(winner_30.exposure.others[0].value, 525);
+		assert_eq!(winners, vec![11, 31]);
+		assert_eq!(assignment, vec![
+			(3, vec![(11, 2816371998), (31, 1478595298)]),
+			(1, vec![(11, 4294967296)]),
+		]);
+		check_exposure_all();
 	})
 }
 
@@ -1593,6 +1621,7 @@ fn switching_roles() {
 		System::set_block_number(6);
 		Session::check_rotate_session(System::block_number());
 		assert_eq_uvec!(Session::validators(), vec![2, 20]);
+		check_exposure_all();
 	});
 }
 
@@ -1728,6 +1757,7 @@ fn bond_with_little_staked_value_bounded_by_slot_stake() {
 		assert_eq!(Balances::free_balance(&2), initial_balance_2 + reward.max(1));
 		// same for 10
 		assert_eq!(Balances::free_balance(&10), initial_balance_10 + 10 + reward.max(1));
+		check_exposure_all();
 	});
 }
 
@@ -1741,17 +1771,6 @@ fn phragmen_linear_worse_case_equalize() {
 		.fair(true)
 		.build(),
 	|| {
-		let bond_validator = |a, b| {
-			let _ = Balances::deposit_creating(&(a-1), b);
-			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
-			assert_ok!(Staking::validate(Origin::signed(a), ValidatorPrefs::default()));
-		};
-		let bond_nominator = |a, b, v| {
-			let _ = Balances::deposit_creating(&(a-1), b);
-			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
-			assert_ok!(Staking::nominate(Origin::signed(a), v));
-		};
-
 		for i in &[10, 20, 30, 40] { assert_ok!(Staking::set_payee(Origin::signed(*i), RewardDestination::Controller)); }
 
 		bond_validator(50, 1000);
@@ -1763,8 +1782,8 @@ fn phragmen_linear_worse_case_equalize() {
 		bond_nominator(6, 1000, vec![21, 31]);
 		bond_nominator(8, 1000, vec![31, 41]);
 		bond_nominator(110, 1000, vec![41, 51]);
-		bond_nominator(112, 1000, vec![51, 61]);
-		bond_nominator(114, 1000, vec![61, 71]);
+		bond_nominator(120, 1000, vec![51, 61]);
+		bond_nominator(130, 1000, vec![61, 71]);
 
 		assert_eq_uvec!(Session::validators(), vec![40, 30]);
 		assert_ok!(Staking::set_validator_count(7));
@@ -1776,20 +1795,22 @@ fn phragmen_linear_worse_case_equalize() {
 
 		// 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
-		// 20  is elected with stake  2017.7421569824219 and score  0.0005277777777777777
-		// 40  is elected with stake  2000.0001049958742 and score  0.0005555555555555556
 		// 70  is elected with stake  1982.2574230340813 and score  0.0007222222222222222
 
+		check_exposure_all();
+
 		assert_eq!(Staking::stakers(11).total, 3000);
-		assert_eq!(Staking::stakers(31).total, 2035);
-		assert_eq!(Staking::stakers(51).total, 2000);
-		assert_eq!(Staking::stakers(61).total, 1968);
-		assert_eq!(Staking::stakers(21).total, 2035);
-		assert_eq!(Staking::stakers(41).total, 2024);
-		assert_eq!(Staking::stakers(71).total, 1936);
+		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);
 	})
 }
 
@@ -1809,6 +1830,7 @@ fn phragmen_chooses_correct_number_of_validators() {
 		Session::check_rotate_session(System::block_number());
 
 		assert_eq!(Session::validators().len(), 1);
+		check_exposure_all();
 	})
 }
 
@@ -1816,18 +1838,9 @@ fn phragmen_chooses_correct_number_of_validators() {
 #[test]
 fn phragmen_score_should_be_accurate_on_large_stakes() {
 	with_externalities(&mut ExtBuilder::default()
-	.nominate(false)
-	.build()
-	, || {
-		let bond_validator = |a, b| {
-			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
-			assert_ok!(Staking::validate(Origin::signed(a), ValidatorPrefs::default()));
-		};
-
-		for i in 1..=8 {
-			let _ = Balances::make_free_balance_be(&i, u64::max_value());
-		}
-
+		.nominate(false)
+		.build(),
+	|| {
 		bond_validator(2, u64::max_value());
 		bond_validator(4, u64::max_value());
 		bond_validator(6, u64::max_value()-1);
@@ -1837,184 +1850,114 @@ fn phragmen_score_should_be_accurate_on_large_stakes() {
 		Session::check_rotate_session(System::block_number());
 
 		assert_eq!(Session::validators(), vec![4, 2]);
+		check_exposure_all();
 	})
 }
 
 #[test]
 fn phragmen_should_not_overflow_validators() {
 	with_externalities(&mut ExtBuilder::default()
-	.nominate(false)
-	.build()
-	, || {
-		let bond_validator = |a, b| {
-			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
-			assert_ok!(Staking::validate(Origin::signed(a), ValidatorPrefs::default()));
-		};
-		let bond_nominator = |a, b, v| {
-			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
-			assert_ok!(Staking::nominate(Origin::signed(a), v));
-		};
-		let check_exposure = |a| {
-			let expo = Staking::stakers(&a);
-			assert_eq!(expo.total, expo.own + expo.others.iter().map(|e| e.value).sum::<u64>());
-		};
-
-		for i in 1..=8 {
-			let _ = Balances::make_free_balance_be(&i, u64::max_value());
-		}
-
+		.nominate(false)
+		.build(),
+	|| {
 		let _ = Staking::chill(Origin::signed(10));
 		let _ = Staking::chill(Origin::signed(20));
 
 		bond_validator(2, u64::max_value());
 		bond_validator(4, u64::max_value());
 
-		bond_nominator(6, u64::max_value()/2, vec![1, 3]);
-		bond_nominator(8, u64::max_value()/2, vec![1, 3]);
+		bond_nominator(6, u64::max_value()/2, vec![3, 5]);
+		bond_nominator(8, u64::max_value()/2, vec![3, 5]);
 
 		System::set_block_number(2);
 		Session::check_rotate_session(System::block_number());
 
 		assert_eq_uvec!(Session::validators(), vec![4, 2]);
-		check_exposure(4);
-		check_exposure(2);
+
+		// This test will fail this. Will saturate.
+		// check_exposure_all();
+		assert_eq!(Staking::stakers(3).total, u64::max_value());
+		assert_eq!(Staking::stakers(5).total, u64::max_value());
 	})
 }
 
 #[test]
 fn phragmen_should_not_overflow_nominators() {
 	with_externalities(&mut ExtBuilder::default()
-	.nominate(false)
-	.build()
-	, || {
-		let bond_validator = |a, b| {
-			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
-			assert_ok!(Staking::validate(Origin::signed(a), ValidatorPrefs::default()));
-		};
-		let bond_nominator = |a, b, v| {
-			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
-			assert_ok!(Staking::nominate(Origin::signed(a), v));
-		};
-		let check_exposure = |a| {
-			let expo = Staking::stakers(&a);
-			assert_eq!(expo.total, expo.own + expo.others.iter().map(|e| e.value).sum::<u64>());
-		};
-
+		.nominate(false)
+		.build(),
+	|| {
 		let _ = Staking::chill(Origin::signed(10));
 		let _ = Staking::chill(Origin::signed(20));
 
-		for i in 1..=8 {
-			let _ = Balances::make_free_balance_be(&i, u64::max_value());
-		}
-
 		bond_validator(2, u64::max_value()/2);
 		bond_validator(4, u64::max_value()/2);
 
-		bond_nominator(6, u64::max_value(), vec![1, 3]);
-		bond_nominator(8, u64::max_value(), vec![1, 3]);
+		bond_nominator(6, u64::max_value(), vec![3, 5]);
+		bond_nominator(8, u64::max_value(), vec![3, 5]);
 
 		System::set_block_number(2);
 		Session::check_rotate_session(System::block_number());
 
 		assert_eq_uvec!(Session::validators(), vec![4, 2]);
-		check_exposure(4);
-		check_exposure(2);
+
+		// Saturate.
+		assert_eq!(Staking::stakers(3).total, u64::max_value());
+		assert_eq!(Staking::stakers(5).total, u64::max_value());
 	})
 }
 
 #[test]
 fn phragmen_should_not_overflow_ultimate() {
 	with_externalities(&mut ExtBuilder::default()
-	.nominate(false)
-	.build()
-	, || {
-		let bond_validator = |a, b| {
-			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
-			assert_ok!(Staking::validate(Origin::signed(a), ValidatorPrefs::default()));
-		};
-		let bond_nominator = |a, b, v| {
-			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
-			assert_ok!(Staking::nominate(Origin::signed(a), v));
-		};
-		let check_exposure = |a| {
-			let expo = Staking::stakers(&a);
-			assert_eq!(expo.total, expo.own + expo.others.iter().map(|e| e.value).sum::<u64>());
-		};
-
-		for i in 1..=8 {
-			let _ = Balances::make_free_balance_be(&i, u64::max_value());
-		}
-
+		.nominate(false)
+		.build(),
+	|| {
 		bond_validator(2, u64::max_value());
 		bond_validator(4, u64::max_value());
 
-		bond_nominator(6, u64::max_value(), vec![1, 3]);
-		bond_nominator(8, u64::max_value(), vec![1, 3]);
+		bond_nominator(6, u64::max_value(), vec![3, 5]);
+		bond_nominator(8, u64::max_value(), vec![3, 5]);
 
 		System::set_block_number(2);
 		Session::check_rotate_session(System::block_number());
 
 		assert_eq_uvec!(Session::validators(), vec![4, 2]);
-		check_exposure(4);
-		check_exposure(2);
+
+		// Saturate.
+		assert_eq!(Staking::stakers(3).total, u64::max_value());
+		assert_eq!(Staking::stakers(5).total, u64::max_value());
 	})
 }
 
 
 #[test]
-fn large_scale_test() {
+fn phragmen_large_scale_test() {
 	with_externalities(&mut ExtBuilder::default()
-	.nominate(false)
-	.minimum_validator_count(1)
-	.validator_count(20)
-	.build()
-	, || {
+		.nominate(false)
+		.minimum_validator_count(1)
+		.validator_count(20)
+		.build(),
+	|| {
 		let _ = Staking::chill(Origin::signed(10));
 		let _ = Staking::chill(Origin::signed(20));
-
-		let bond_validator = |a, b| {
-			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
-			assert_ok!(Staking::validate(Origin::signed(a), ValidatorPrefs::default()));
-		};
-		let bond_nominator = |a, b, v| {
-			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
-			assert_ok!(Staking::nominate(Origin::signed(a), v));
-		};
-		let check_expo = |e: Exposure<u64, u64>| {
-			assert_eq!(e.total, e.own + e.others.iter().fold(0, |s, v| s + v.value));
-		};
-
+		let _ = Staking::chill(Origin::signed(30));
 		let prefix = 200;
-		for i in prefix..=(prefix+50) {
-			let _ = Balances::make_free_balance_be(&i, u64::max_value());
-		}
 
 		bond_validator(prefix + 2,  1);
-		bond_validator(prefix + 4,  105);
-		bond_validator(prefix + 6,  3879248198);
+		bond_validator(prefix + 4,  100);
+		bond_validator(prefix + 6,  1000000);
 		bond_validator(prefix + 8,  100000000001000);
 		bond_validator(prefix + 10, 100000000002000);
-		bond_validator(prefix + 12, 150000000000000);
+		bond_validator(prefix + 12, 100000000003000);
 		bond_validator(prefix + 14, 400000000000000);
-		bond_validator(prefix + 16, 524611669156413);
-		bond_validator(prefix + 18, 700000000000000);
-		bond_validator(prefix + 20, 797663650978304);
-		bond_validator(prefix + 22, 900003879248198);
-		bond_validator(prefix + 24, 997530000000000);
-		bond_validator(prefix + 26, 1000000000010000);
-		bond_validator(prefix + 28, 1000000000020000);
-		bond_validator(prefix + 30, 1000003879248198);
-		bond_validator(prefix + 32, 1200000000000000);
-		bond_validator(prefix + 34, 7997659802817256);
-		bond_validator(prefix + 36, 18000000000000000);
-		bond_validator(prefix + 38, 20000033025753738);
-		bond_validator(prefix + 40, 500000000000100000);
-		bond_validator(prefix + 42, 500000000000200000);
+		bond_validator(prefix + 16, 400000000001000);
+		bond_validator(prefix + 18, 18000000000000000);
+		bond_validator(prefix + 20, 20000000000000000);
+		bond_validator(prefix + 22, 500000000000100000);
+		bond_validator(prefix + 24, 500000000000200000);
 
-		Balances::make_free_balance_be(&50, u64::max_value());
-		Balances::make_free_balance_be(&49, u64::max_value());
-		bond_nominator(50, 990000068998617227, vec![
-			prefix + 1,
+		bond_nominator(50, 990000000000000000, vec![
 			prefix + 3,
 			prefix + 5,
 			prefix + 7,
@@ -2026,38 +1969,66 @@ fn large_scale_test() {
 			prefix + 19,
 			prefix + 21,
 			prefix + 23,
-			prefix + 25,
-			prefix + 27,
-			prefix + 29,
-			prefix + 31,
-			prefix + 33,
-			prefix + 35,
-			prefix + 37,
-			prefix + 39,
-			prefix + 41,
-			prefix + 43,
-			prefix + 45]
+			prefix + 25]
 		);
 
 		System::set_block_number(1);
 		Session::check_rotate_session(System::block_number());
 
-		// Each exposure => total == own + sum(others)
-		Session::validators().iter().for_each(|acc| check_expo(Staking::stakers(acc-1)));
-
-		// aside from some error, stake must be divided correctly
-		assert!(
-			990000068998617227
-			- Session::validators()
-				.iter()
-				.map(|v| Staking::stakers(v-1))
-				.fold(0, |s, v| if v.others.len() > 0 { s + v.others[0].value } else { s })
-			< 100
-		);
-
 		// 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>)>>());
+			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
+			)
+		);
 	})
-}
+}
+
+#[test]
+fn phragmen_large_scale_test_2() {
+	with_externalities(&mut ExtBuilder::default()
+		.nominate(false)
+		.minimum_validator_count(1)
+		.validator_count(2)
+		.build(),
+	|| {
+		let _ = Staking::chill(Origin::signed(10));
+		let _ = Staking::chill(Origin::signed(20));
+		let nom_budget: u64 = 1_000_000_000_000_000_000;
+		let c_budget: u64 = 4_000_000;
+
+		bond_validator(2, c_budget as u64);
+		bond_validator(4, c_budget as u64);
+
+		bond_nominator(50, nom_budget, vec![3, 5]);
+
+		System::set_block_number(1);
+		Session::check_rotate_session(System::block_number());
+
+		// Each exposure => total == own + sum(others)
+		check_exposure_all();
+
+		assert_total_expo(3, nom_budget / 2 + c_budget);
+		assert_total_expo(5, nom_budget / 2 + c_budget);
+	})
+}