Custom Codec Implenetation for NPoS Election (#6720)

* Fancy compact encode/decode impl for compact solution

* Make it optional

* Remove extra file

* Update primitives/npos-elections/compact/src/lib.rs

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

* Final fixes.

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

substrate/frame/staking/src/lib.rs

@@ -326,7 +326,7 @@ use frame_system::{
 };
 use sp_npos_elections::{
 	ExtendedBalance, Assignment, ElectionScore, ElectionResult as PrimitiveElectionResult,
-	build_support_map, evaluate_support, seq_phragmen, generate_compact_solution_type,
+	build_support_map, evaluate_support, seq_phragmen, generate_solution_type,
 	is_score_better, VotingLimit, SupportMap, VoteWeight,
 };
 
@@ -368,7 +368,25 @@ pub type EraIndex = u32;
 pub type RewardPoint = u32;
 
 // Note: Maximum nomination limit is set here -- 16.
-generate_compact_solution_type!(pub GenericCompactAssignments, 16);
+generate_solution_type!(
+	#[compact]
+	pub struct CompactAssignments::<NominatorIndex, ValidatorIndex, OffchainAccuracy>(16)
+);
+
+/// Accuracy used for on-chain election.
+pub type ChainAccuracy = Perbill;
+
+/// Accuracy used for off-chain election. This better be small.
+pub type OffchainAccuracy = PerU16;
+
+/// The balance type of this module.
+pub type BalanceOf<T> =
+	<<T as Trait>::Currency as Currency<<T as frame_system::Trait>::AccountId>>::Balance;
+
+type PositiveImbalanceOf<T> =
+	<<T as Trait>::Currency as Currency<<T as frame_system::Trait>::AccountId>>::PositiveImbalance;
+type NegativeImbalanceOf<T> =
+	<<T as Trait>::Currency as Currency<<T as frame_system::Trait>::AccountId>>::NegativeImbalance;
 
 /// Information regarding the active era (era in used in session).
 #[derive(Encode, Decode, RuntimeDebug)]
@@ -382,25 +400,6 @@ pub struct ActiveEraInfo {
 	start: Option<u64>,
 }
 
-/// Accuracy used for on-chain election.
-pub type ChainAccuracy = Perbill;
-
-/// Accuracy used for off-chain election. This better be small.
-pub type OffchainAccuracy = PerU16;
-
-/// The balance type of this module.
-pub type BalanceOf<T> =
-	<<T as Trait>::Currency as Currency<<T as frame_system::Trait>::AccountId>>::Balance;
-
-/// The compact type for election solutions.
-pub type CompactAssignments =
-	GenericCompactAssignments<NominatorIndex, ValidatorIndex, OffchainAccuracy>;
-
-type PositiveImbalanceOf<T> =
-	<<T as Trait>::Currency as Currency<<T as frame_system::Trait>::AccountId>>::PositiveImbalance;
-type NegativeImbalanceOf<T> =
-	<<T as Trait>::Currency as Currency<<T as frame_system::Trait>::AccountId>>::NegativeImbalance;
-
 /// Reward points of an era. Used to split era total payout between validators.
 ///
 /// This points will be used to reward validators and their respective nominators.

substrate/primitives/arithmetic/fuzzer/src/per_thing_rational.rs

@@ -118,6 +118,5 @@ fn assert_per_thing_equal_error<P: PerThing>(a: P, b: P, err: u128) {
 	let a_abs = a.deconstruct().saturated_into::<u128>();
 	let b_abs = b.deconstruct().saturated_into::<u128>();
 	let diff = a_abs.max(b_abs) - a_abs.min(b_abs);
-	dbg!(&diff);
 	assert!(diff <= err, "{:?} !~ {:?}", a, b);
 }

substrate/primitives/npos-elections/compact/src/assignment.rs

@@ -15,11 +15,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! Code generation for the ratio assignment type.
+//! Code generation for the ratio assignment type' compact representation.
 
 use crate::field_name_for;
 use proc_macro2::TokenStream as TokenStream2;
-use syn::GenericArgument;
 use quote::quote;
 
 fn from_impl(count: usize) -> TokenStream2 {
@@ -27,8 +26,8 @@ fn from_impl(count: usize) -> TokenStream2 {
 		let name = field_name_for(1);
 		quote!(1 => compact.#name.push(
 			(
-				index_of_voter(&who).ok_or(_phragmen::Error::CompactInvalidIndex)?,
-				index_of_target(&distribution[0].0).ok_or(_phragmen::Error::CompactInvalidIndex)?,
+				index_of_voter(&who).or_invalid_index()?,
+				index_of_target(&distribution[0].0).or_invalid_index()?,
 			)
 		),)
 	};
@@ -37,29 +36,29 @@ fn from_impl(count: usize) -> TokenStream2 {
 		let name = field_name_for(2);
 		quote!(2 => compact.#name.push(
 			(
-				index_of_voter(&who).ok_or(_phragmen::Error::CompactInvalidIndex)?,
+				index_of_voter(&who).or_invalid_index()?,
 				(
-					index_of_target(&distribution[0].0).ok_or(_phragmen::Error::CompactInvalidIndex)?,
+					index_of_target(&distribution[0].0).or_invalid_index()?,
 					distribution[0].1,
 				),
-				index_of_target(&distribution[1].0).ok_or(_phragmen::Error::CompactInvalidIndex)?,
+				index_of_target(&distribution[1].0).or_invalid_index()?,
 			)
 		),)
 	};
 
 	let from_impl_rest = (3..=count).map(|c| {
 		let inner = (0..c-1).map(|i|
-			quote!((index_of_target(&distribution[#i].0).ok_or(_phragmen::Error::CompactInvalidIndex)?, distribution[#i].1),)
+			quote!((index_of_target(&distribution[#i].0).or_invalid_index()?, distribution[#i].1),)
 		).collect::<TokenStream2>();
 
 		let field_name = field_name_for(c);
 		let last_index = c - 1;
-		let last = quote!(index_of_target(&distribution[#last_index].0).ok_or(_phragmen::Error::CompactInvalidIndex)?);
+		let last = quote!(index_of_target(&distribution[#last_index].0).or_invalid_index()?);
 
 		quote!(
 			#c => compact.#field_name.push(
 				(
-					index_of_voter(&who).ok_or(_phragmen::Error::CompactInvalidIndex)?,
+					index_of_voter(&who).or_invalid_index()?,
 					[#inner],
 					#last,
 				)
@@ -74,15 +73,15 @@ fn from_impl(count: usize) -> TokenStream2 {
 	)
 }
 
-fn into_impl(count: usize) -> TokenStream2 {
+fn into_impl(count: usize, per_thing: syn::Type) -> TokenStream2 {
 	let into_impl_single = {
 		let name = field_name_for(1);
 		quote!(
 			for (voter_index, target_index) in self.#name {
 				assignments.push(_phragmen::Assignment {
-					who: voter_at(voter_index).ok_or(_phragmen::Error::CompactInvalidIndex)?,
+					who: voter_at(voter_index).or_invalid_index()?,
 					distribution: vec![
-						(target_at(target_index).ok_or(_phragmen::Error::CompactInvalidIndex)?, Accuracy::one())
+						(target_at(target_index).or_invalid_index()?, #per_thing::one())
 					],
 				})
 			}
@@ -93,21 +92,21 @@ fn into_impl(count: usize) -> TokenStream2 {
 		let name = field_name_for(2);
 		quote!(
 			for (voter_index, (t1_idx, p1), t2_idx) in self.#name {
-				if p1 >= Accuracy::one() {
+				if p1 >= #per_thing::one() {
 					return Err(_phragmen::Error::CompactStakeOverflow);
 				}
 
 				// defensive only. Since Percent doesn't have `Sub`.
 				let p2 = _phragmen::sp_arithmetic::traits::Saturating::saturating_sub(
-					Accuracy::one(),
+					#per_thing::one(),
 					p1,
 				);
 
 				assignments.push( _phragmen::Assignment {
-					who: voter_at(voter_index).ok_or(_phragmen::Error::CompactInvalidIndex)?,
+					who: voter_at(voter_index).or_invalid_index()?,
 					distribution: vec![
-						(target_at(t1_idx).ok_or(_phragmen::Error::CompactInvalidIndex)?, p1),
-						(target_at(t2_idx).ok_or(_phragmen::Error::CompactInvalidIndex)?, p2),
+						(target_at(t1_idx).or_invalid_index()?, p1),
+						(target_at(t2_idx).or_invalid_index()?, p2),
 					]
 				});
 			}
@@ -118,30 +117,30 @@ fn into_impl(count: usize) -> TokenStream2 {
 		let name = field_name_for(c);
 		quote!(
 			for (voter_index, inners, t_last_idx) in self.#name {
-				let mut sum = Accuracy::zero();
+				let mut sum = #per_thing::zero();
 				let mut inners_parsed = inners
 					.iter()
 					.map(|(ref t_idx, p)| {
 						sum = _phragmen::sp_arithmetic::traits::Saturating::saturating_add(sum, *p);
-						let target = target_at(*t_idx).ok_or(_phragmen::Error::CompactInvalidIndex)?;
+						let target = target_at(*t_idx).or_invalid_index()?;
 						Ok((target, *p))
 					})
-					.collect::<Result<Vec<(A, Accuracy)>, _phragmen::Error>>()?;
+					.collect::<Result<Vec<(A, #per_thing)>, _phragmen::Error>>()?;
 
-				if sum >= Accuracy::one() {
+				if sum >= #per_thing::one() {
 					return Err(_phragmen::Error::CompactStakeOverflow);
 				}
 
 				// defensive only. Since Percent doesn't have `Sub`.
 				let p_last = _phragmen::sp_arithmetic::traits::Saturating::saturating_sub(
-					Accuracy::one(),
+					#per_thing::one(),
 					sum,
 				);
 
-				inners_parsed.push((target_at(t_last_idx).ok_or(_phragmen::Error::CompactInvalidIndex)?, p_last));
+				inners_parsed.push((target_at(t_last_idx).or_invalid_index()?, p_last));
 
 				assignments.push(_phragmen::Assignment {
-					who: voter_at(voter_index).ok_or(_phragmen::Error::CompactInvalidIndex)?,
+					who: voter_at(voter_index).or_invalid_index()?,
 					distribution: inners_parsed,
 				});
 			}
@@ -157,39 +156,28 @@ fn into_impl(count: usize) -> TokenStream2 {
 
 pub(crate) fn assignment(
 	ident: syn::Ident,
-	voter_type: GenericArgument,
-	target_type: GenericArgument,
+	voter_type: syn::Type,
+	target_type: syn::Type,
+	weight_type: syn::Type,
 	count: usize,
 ) -> TokenStream2 {
-
 	let from_impl = from_impl(count);
-	let into_impl = into_impl(count);
+	let into_impl = into_impl(count, weight_type.clone());
 
 	quote!(
-		impl<
-			#voter_type: _phragmen::codec::Codec + Default + Copy,
-			#target_type: _phragmen::codec::Codec + Default + Copy,
-			Accuracy:
-				_phragmen::codec::Codec + Default + Clone + _phragmen::sp_arithmetic::PerThing +
-				PartialOrd,
-		>
-		#ident<#voter_type, #target_type, Accuracy>
-		{
+		use _phragmen::__OrInvalidIndex;
+		impl #ident {
 			pub fn from_assignment<FV, FT, A>(
-				assignments: Vec<_phragmen::Assignment<A, Accuracy>>,
+				assignments: Vec<_phragmen::Assignment<A, #weight_type>>,
 				index_of_voter: FV,
 				index_of_target: FT,
 			) -> Result<Self, _phragmen::Error>
 				where
+					A: _phragmen::IdentifierT,
 					for<'r> FV: Fn(&'r A) -> Option<#voter_type>,
 					for<'r> FT: Fn(&'r A) -> Option<#target_type>,
-					A: _phragmen::IdentifierT,
 			{
-				let mut compact: #ident<
-					#voter_type,
-					#target_type,
-					Accuracy,
-				> = Default::default();
+				let mut compact: #ident = Default::default();
 
 				for _phragmen::Assignment { who, distribution } in assignments {
 					match distribution.len() {
@@ -207,8 +195,8 @@ pub(crate) fn assignment(
 				self,
 				voter_at: impl Fn(#voter_type) -> Option<A>,
 				target_at: impl Fn(#target_type) -> Option<A>,
-			) -> Result<Vec<_phragmen::Assignment<A, Accuracy>>, _phragmen::Error> {
-				let mut assignments: Vec<_phragmen::Assignment<A, Accuracy>> = Default::default();
+			) -> Result<Vec<_phragmen::Assignment<A, #weight_type>>, _phragmen::Error> {
+				let mut assignments: Vec<_phragmen::Assignment<A, #weight_type>> = Default::default();
 				#into_impl
 				Ok(assignments)
 			}

substrate/primitives/npos-elections/compact/src/codec.rs

@@ -0,0 +1,203 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2020 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Code generation for the ratio assignment type' encode/decode impl.
+
+use crate::field_name_for;
+use proc_macro2::TokenStream as TokenStream2;
+use quote::quote;
+
+pub(crate) fn codec_impl(
+	ident: syn::Ident,
+	voter_type: syn::Type,
+	target_type: syn::Type,
+	weight_type: syn::Type,
+	count: usize,
+) -> TokenStream2 {
+	let encode = encode_impl(ident.clone(), count);
+	let decode = decode_impl(ident, voter_type, target_type, weight_type, count);
+
+	quote! {
+		#encode
+		#decode
+	}
+}
+
+fn decode_impl(
+	ident: syn::Ident,
+	voter_type: syn::Type,
+	target_type: syn::Type,
+	weight_type: syn::Type,
+	count: usize,
+) -> TokenStream2 {
+	let decode_impl_single = {
+		let name = field_name_for(1);
+		quote! {
+			let #name =
+			<
+				Vec<(_phragmen::codec::Compact<#voter_type>, _phragmen::codec::Compact<#target_type>)>
+				as
+				_phragmen::codec::Decode
+			>::decode(value)?;
+			let #name = #name
+				.into_iter()
+				.map(|(v, t)| (v.0, t.0))
+				.collect::<Vec<_>>();
+		}
+	};
+
+	let decode_impl_double = {
+		let name = field_name_for(2);
+		quote! {
+			let #name =
+			<
+				Vec<(
+					_phragmen::codec::Compact<#voter_type>,
+					(_phragmen::codec::Compact<#target_type>, _phragmen::codec::Compact<#weight_type>),
+					_phragmen::codec::Compact<#target_type>,
+				)>
+				as
+				_phragmen::codec::Decode
+			>::decode(value)?;
+			let #name = #name
+				.into_iter()
+				.map(|(v, (t1, w), t2)| (v.0, (t1.0, w.0), t2.0))
+				.collect::<Vec<_>>();
+		}
+	};
+
+	let decode_impl_rest = (3..=count).map(|c| {
+		let name = field_name_for(c);
+
+		let inner_impl = (0..c-1).map(|i|
+			quote! { ( (inner[#i].0).0, (inner[#i].1).0 ), }
+		).collect::<TokenStream2>();
+
+		quote! {
+			let #name =
+			<
+				Vec<(
+					_phragmen::codec::Compact<#voter_type>,
+					[(_phragmen::codec::Compact<#target_type>, _phragmen::codec::Compact<#weight_type>); #c-1],
+					_phragmen::codec::Compact<#target_type>,
+				)>
+				as _phragmen::codec::Decode
+			>::decode(value)?;
+			let #name = #name
+				.into_iter()
+				.map(|(v, inner, t_last)| (
+					v.0,
+					[ #inner_impl ],
+					t_last.0,
+				))
+				.collect::<Vec<_>>();
+		}
+	}).collect::<TokenStream2>();
+
+
+	let all_field_names = (1..=count).map(|c| {
+		let name = field_name_for(c);
+		quote! { #name, }
+	}).collect::<TokenStream2>();
+
+	quote!(
+		impl _phragmen::codec::Decode for #ident {
+			fn decode<I: _phragmen::codec::Input>(value: &mut I) -> Result<Self, _phragmen::codec::Error> {
+				#decode_impl_single
+				#decode_impl_double
+				#decode_impl_rest
+
+				// The above code generates variables with the decoded value with the same name as
+				// filed names of the struct, i.e. `let votes4 = decode_value_of_votes4`. All we
+				// have to do is collect them into the main struct now.
+				Ok(#ident { #all_field_names })
+			}
+		}
+	)
+}
+
+// General attitude is that we will convert inner values to `Compact` and then use the normal
+// `Encode` implementation.
+fn encode_impl(ident: syn::Ident, count: usize) -> TokenStream2 {
+	let encode_impl_single = {
+		let name = field_name_for(1);
+		quote! {
+			let #name = self.#name
+				.iter()
+				.map(|(v, t)| (
+					_phragmen::codec::Compact(v.clone()),
+					_phragmen::codec::Compact(t.clone()),
+				))
+				.collect::<Vec<_>>();
+			#name.encode_to(&mut r);
+		}
+	};
+
+	let encode_impl_double = {
+		let name = field_name_for(2);
+		quote! {
+			let #name = self.#name
+				.iter()
+				.map(|(v, (t1, w), t2)| (
+					_phragmen::codec::Compact(v.clone()),
+					(
+						_phragmen::codec::Compact(t1.clone()),
+						_phragmen::codec::Compact(w.clone())
+					),
+					_phragmen::codec::Compact(t2.clone()),
+				))
+				.collect::<Vec<_>>();
+			#name.encode_to(&mut r);
+		}
+	};
+
+	let encode_impl_rest = (3..=count).map(|c| {
+		let name = field_name_for(c);
+
+		// we use the knowledge of the length to avoid copy_from_slice.
+		let inners_compact_array = (0..c-1).map(|i|
+			quote!{(
+				_phragmen::codec::Compact(inner[#i].0.clone()),
+				_phragmen::codec::Compact(inner[#i].1.clone()),
+			),}
+		).collect::<TokenStream2>();
+
+		quote! {
+			let #name = self.#name
+				.iter()
+				.map(|(v, inner, t_last)| (
+					_phragmen::codec::Compact(v.clone()),
+					[ #inners_compact_array ],
+					_phragmen::codec::Compact(t_last.clone()),
+				))
+				.collect::<Vec<_>>();
+			#name.encode_to(&mut r);
+		}
+	}).collect::<TokenStream2>();
+
+	quote!(
+		impl _phragmen::codec::Encode for #ident {
+			fn encode(&self) -> Vec<u8> {
+				let mut r = vec![];
+				#encode_impl_single
+				#encode_impl_double
+				#encode_impl_rest
+				r
+			}
+		}
+	)
+}

substrate/primitives/npos-elections/compact/src/lib.rs

@@ -21,103 +21,92 @@ use proc_macro::TokenStream;
 use proc_macro2::{TokenStream as TokenStream2, Span, Ident};
 use proc_macro_crate::crate_name;
 use quote::quote;
-use syn::{GenericArgument, Type, parse::{Parse, ParseStream, Result}};
+use syn::{parse::{Parse, ParseStream, Result}};
 
 mod assignment;
-mod staked;
+mod codec;
 
 // prefix used for struct fields in compact.
 const PREFIX: &'static str = "votes";
 
-/// Generates a struct to store the election assignments in a compact way. The struct can only store
-/// distributions up to the given input count. The given count must be greater than 2.
+pub(crate) fn syn_err(message: &'static str) -> syn::Error {
+	syn::Error::new(Span::call_site(), message)
+}
+
+/// Generates a struct to store the election result in a small way. This can encode a structure
+/// which is the equivalent of a `sp_npos_elections::Assignment<_>`.
+///
+/// The following data types can be configured by the macro.
+///
+/// - The identifier of the voter. This can be any type that supports `parity-scale-codec`'s compact
+///   encoding.
+/// - The identifier of the target. This can be any type that supports `parity-scale-codec`'s
+///   compact encoding.
+/// - The accuracy of the ratios. This must be one of the `PerThing` types defined in
+///   `sp-arithmetic`.
+///
+/// Moreover, the maximum number of edges per voter (distribution per assignment) also need to be
+/// specified. Attempting to convert from/to an assignment with more distributions will fail.
+///
+///
+/// For example, the following generates a public struct with name `TestSolution` with `u16` voter
+/// type, `u8` target type and `Perbill` accuracy with maximum of 8 edges per voter.
 ///
 /// ```ignore
-/// // generate a struct with nominator and edge weight u128, with maximum supported
-/// // edge per voter of 16.
-/// generate_compact_solution_type(pub TestCompact, 16)
+/// generate_solution_type!(pub struct TestSolution<u16, u8, Perbill>::(8))
 /// ```
 ///
-/// This generates:
+/// The given struct provides function to convert from/to Assignment:
+///
+/// - [`from_assignment()`].
+/// - [`fn into_assignment()`].
+///
+/// The generated struct is by default deriving both `Encode` and `Decode`. This is okay but could
+/// lead to many 0s in the solution. If prefixed with `#[compact]`, then a custom compact encoding
+/// for numbers will be used, similar to how `parity-scale-codec`'s `Compact` works.
 ///
 /// ```ignore
-/// pub struct TestCompact<V, T, W> {
-/// 	votes1: Vec<(V, T)>,
-/// 	votes2: Vec<(V, (T, W), T)>,
-/// 	votes3: Vec<(V, [(T, W); 2usize], T)>,
-/// 	votes4: Vec<(V, [(T, W); 3usize], T)>,
-/// 	votes5: Vec<(V, [(T, W); 4usize], T)>,
-/// 	votes6: Vec<(V, [(T, W); 5usize], T)>,
-/// 	votes7: Vec<(V, [(T, W); 6usize], T)>,
-/// 	votes8: Vec<(V, [(T, W); 7usize], T)>,
-/// 	votes9: Vec<(V, [(T, W); 8usize], T)>,
-/// 	votes10: Vec<(V, [(T, W); 9usize], T)>,
-/// 	votes11: Vec<(V, [(T, W); 10usize], T)>,
-/// 	votes12: Vec<(V, [(T, W); 11usize], T)>,
-/// 	votes13: Vec<(V, [(T, W); 12usize], T)>,
-/// 	votes14: Vec<(V, [(T, W); 13usize], T)>,
-/// 	votes15: Vec<(V, [(T, W); 14usize], T)>,
-/// 	votes16: Vec<(V, [(T, W); 15usize], T)>,
-/// }
+/// generate_solution_type!(
+///     #[compact]
+///     pub struct TestSolutionCompact<u16, u8, Perbill>::(8)
+/// )
 /// ```
-///
-/// The generic arguments are:
-/// - `V`: identifier/index for voter (nominator) types.
-/// - `T` identifier/index for candidate (validator) types.
-/// - `W` weight type.
-///
-/// Some conversion implementations are provided by default if
-/// - `W` is u128, or
-/// - `W` is anything that implements `PerThing` (such as `Perbill`)
-///
-/// The ideas behind the structure are as follows:
-///
-/// - For single distribution, no weight is stored. The weight is known to be 100%.
-/// - For all the rest, the weight if the last distribution is omitted. This value can be computed
-///   from the rest.
-///
 #[proc_macro]
-pub fn generate_compact_solution_type(item: TokenStream) -> TokenStream {
-	let CompactSolutionDef {
+pub fn generate_solution_type(item: TokenStream) -> TokenStream {
+	let SolutionDef {
 		vis,
 		ident,
 		count,
-	} = syn::parse_macro_input!(item as CompactSolutionDef);
-
-	let voter_type = GenericArgument::Type(Type::Verbatim(quote!(V)));
-	let target_type = GenericArgument::Type(Type::Verbatim(quote!(T)));
-	let weight_type = GenericArgument::Type(Type::Verbatim(quote!(W)));
+		voter_type,
+		target_type,
+		weight_type,
+		compact_encoding,
+	} = syn::parse_macro_input!(item as SolutionDef);
 
 	let imports = imports().unwrap_or_else(|e| e.to_compile_error());
 
-	let compact_def = struct_def(
+	let solution_struct = struct_def(
 		vis,
 		ident.clone(),
 		count,
 		voter_type.clone(),
 		target_type.clone(),
-		weight_type,
+		weight_type.clone(),
+		compact_encoding,
 	).unwrap_or_else(|e| e.to_compile_error());
 
 	let assignment_impls = assignment::assignment(
 		ident.clone(),
 		voter_type.clone(),
 		target_type.clone(),
-		count,
-	);
-
-	let staked_impls = staked::staked(
-		ident,
-		voter_type,
-		target_type,
+		weight_type.clone(),
 		count,
 	);
 
 	quote!(
 		#imports
-		#compact_def
+		#solution_struct
 		#assignment_impls
-		#staked_impls
 	).into()
 }
 
@@ -125,25 +114,27 @@ fn struct_def(
 	vis: syn::Visibility,
 	ident: syn::Ident,
 	count: usize,
-	voter_type: GenericArgument,
-	target_type: GenericArgument,
-	weight_type: GenericArgument,
+	voter_type: syn::Type,
+	target_type: syn::Type,
+	weight_type: syn::Type,
+	compact_encoding: bool,
 ) -> Result<TokenStream2> {
 	if count <= 2 {
-		Err(syn::Error::new(
-			Span::call_site(),
-			"cannot build compact solution struct with capacity less than 2."
-		))?
+		Err(syn_err("cannot build compact solution struct with capacity less than 3."))?
 	}
 
 	let singles = {
 		let name = field_name_for(1);
-		quote!(#name: Vec<(#voter_type, #target_type)>,)
+		quote!(
+			#name: Vec<(#voter_type, #target_type)>,
+		)
 	};
 
 	let doubles = {
 		let name = field_name_for(2);
-		quote!(#name: Vec<(#voter_type, (#target_type, #weight_type), #target_type)>,)
+		quote!(
+			#name: Vec<(#voter_type, (#target_type, #weight_type), #target_type)>,
+		)
 	};
 
 	let rest = (3..=count).map(|c| {
@@ -175,31 +166,34 @@ fn struct_def(
 		)
 	}).collect::<TokenStream2>();
 
+	let derives_and_maybe_compact_encoding = if compact_encoding {
+		// custom compact encoding.
+		let compact_impl = codec::codec_impl(
+			ident.clone(),
+			voter_type.clone(),
+			target_type.clone(),
+			weight_type.clone(),
+			count,
+		);
+		quote!{
+			#compact_impl
+			#[derive(Default, PartialEq, Eq, Clone, Debug)]
+		}
+	} else {
+		// automatically derived.
+		quote!(#[derive(Default, PartialEq, Eq, Clone, Debug, _phragmen::codec::Encode, _phragmen::codec::Decode)])
+	};
+
 	Ok(quote! (
 		/// A struct to encode a election assignment in a compact way.
-		#[derive(
-			Default,
-			PartialEq,
-			Eq,
-			Clone,
-			Debug,
-			_phragmen::codec::Encode,
-			_phragmen::codec::Decode,
-		)]
-		#vis struct #ident<#voter_type, #target_type, #weight_type> {
-			// _marker: sp_std::marker::PhantomData<A>,
-			#singles
-			#doubles
-			#rest
-		}
+		#derives_and_maybe_compact_encoding
+		#vis struct #ident { #singles #doubles #rest }
 
-		impl<#voter_type, #target_type, #weight_type> _phragmen::VotingLimit
-		for #ident<#voter_type, #target_type, #weight_type>
-		{
+		impl _phragmen::VotingLimit for #ident {
 			const LIMIT: usize = #count;
 		}
 
-		impl<#voter_type, #target_type, #weight_type> #ident<#voter_type, #target_type, #weight_type> {
+		impl #ident {
 			/// Get the length of all the assignments that this type is encoding. This is basically
 			/// the same as the number of assignments, or the number of voters in total.
 			pub fn len(&self) -> usize {
@@ -239,20 +233,79 @@ fn imports() -> Result<TokenStream2> {
 	}
 }
 
-struct CompactSolutionDef {
+struct SolutionDef {
 	vis: syn::Visibility,
 	ident: syn::Ident,
+	voter_type: syn::Type,
+	target_type: syn::Type,
+	weight_type: syn::Type,
 	count: usize,
+	compact_encoding: bool,
 }
 
-impl Parse for CompactSolutionDef {
+fn check_compact_attr(input: ParseStream) -> Result<bool> {
+	let mut attrs = input.call(syn::Attribute::parse_outer).unwrap_or_default();
+	if attrs.len() == 1 {
+		let attr = attrs.pop().expect("Vec with len 1 can be popped.");
+		if attr.path.segments.len() == 1 {
+			let segment = attr.path.segments.first().expect("Vec with len 1 can be popped.");
+			if segment.ident == Ident::new("compact", Span::call_site()) {
+				Ok(true)
+			} else {
+				Err(syn_err("generate_solution_type macro can only accept #[compact] attribute."))
+			}
+		} else {
+			Err(syn_err("generate_solution_type macro can only accept #[compact] attribute."))
+		}
+	} else {
+		Ok(false)
+	}
+}
+
+/// #[compact] pub struct CompactName::<u32, u32, u32>()
+impl Parse for SolutionDef {
 	fn parse(input: ParseStream) -> syn::Result<Self> {
+		// optional #[compact]
+		let compact_encoding = check_compact_attr(input)?;
+
+		// <vis> struct <name>
 		let vis: syn::Visibility = input.parse()?;
+		let _ = <syn::Token![struct]>::parse(input)?;
 		let ident: syn::Ident = input.parse()?;
-		let _ = <syn::Token![,]>::parse(input)?;
-		let count_literal: syn::LitInt = input.parse()?;
-		let count = count_literal.base10_parse::<usize>()?;
-		Ok(Self { vis, ident, count } )
+
+		// ::<V, T, W>
+		let _ = <syn::Token![::]>::parse(input)?;
+		let generics: syn::AngleBracketedGenericArguments = input.parse()?;
+
+		if generics.args.len() != 3 {
+			return Err(syn_err("Must provide 3 generic args."))
+		}
+
+		let mut types: Vec<syn::Type> = generics.args.iter().map(|t|
+			match t {
+				syn::GenericArgument::Type(ty) => Ok(ty.clone()),
+				_ => Err(syn_err("Wrong type of generic provided. Must be a `type`.")),
+			}
+		).collect::<Result<_>>()?;
+
+		let weight_type = types.pop().expect("Vector of length 3 can be popped; qed");
+		let target_type = types.pop().expect("Vector of length 2 can be popped; qed");
+		let voter_type = types.pop().expect("Vector of length 1 can be popped; qed");
+
+		// (<count>)
+		let count_expr: syn::ExprParen = input.parse()?;
+		let expr = count_expr.expr;
+		let expr_lit = match *expr {
+			syn::Expr::Lit(count_lit) => count_lit.lit,
+			_ => return Err(syn_err("Count must be literal."))
+		};
+		let int_lit = match expr_lit {
+			syn::Lit::Int(int_lit) => int_lit,
+			_ => return Err(syn_err("Count must be int literal."))
+		};
+		let count = int_lit.base10_parse::<usize>()?;
+
+		Ok(Self { vis, ident, voter_type, target_type, weight_type, count, compact_encoding } )
 	}
 }
 

substrate/primitives/npos-elections/compact/src/staked.rs

@@ -1,212 +0,0 @@
-// This file is part of Substrate.
-
-// Copyright (C) 2020 Parity Technologies (UK) Ltd.
-// SPDX-License-Identifier: Apache-2.0
-
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// 	http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-//! Code generation for the staked assignment type.
-
-use crate::field_name_for;
-use proc_macro2::{TokenStream as TokenStream2};
-use syn::{GenericArgument};
-use quote::quote;
-
-fn from_impl(count: usize) -> TokenStream2 {
-	let from_impl_single = {
-		let name = field_name_for(1);
-		quote!(1 => compact.#name.push(
-			(
-				index_of_voter(&who).ok_or(_phragmen::Error::CompactInvalidIndex)?,
-				index_of_target(&distribution[0].0).ok_or(_phragmen::Error::CompactInvalidIndex)?,
-			)
-		),)
-	};
-
-	let from_impl_double = {
-		let name = field_name_for(2);
-		quote!(2 => compact.#name.push(
-			(
-				index_of_voter(&who).ok_or(_phragmen::Error::CompactInvalidIndex)?,
-				(
-					index_of_target(&distribution[0].0).ok_or(_phragmen::Error::CompactInvalidIndex)?,
-					distribution[0].1,
-				),
-				index_of_target(&distribution[1].0).ok_or(_phragmen::Error::CompactInvalidIndex)?,
-			)
-		),)
-	};
-
-	let from_impl_rest = (3..=count).map(|c| {
-		let inner = (0..c-1).map(|i|
-			quote!((index_of_target(&distribution[#i].0).ok_or(_phragmen::Error::CompactInvalidIndex)?, distribution[#i].1),)
-		).collect::<TokenStream2>();
-
-		let field_name = field_name_for(c);
-		let last_index = c - 1;
-		let last = quote!(index_of_target(&distribution[#last_index].0).ok_or(_phragmen::Error::CompactInvalidIndex)?);
-
-		quote!(
-			#c => compact.#field_name.push(
-				(index_of_voter(&who).ok_or(_phragmen::Error::CompactInvalidIndex)?, [#inner], #last)
-			),
-		)
-	}).collect::<TokenStream2>();
-
-	quote!(
-		#from_impl_single
-		#from_impl_double
-		#from_impl_rest
-	)
-}
-
-fn into_impl(count: usize) -> TokenStream2 {
-	let into_impl_single = {
-		let name = field_name_for(1);
-		quote!(
-			for (voter_index, target_index) in self.#name {
-				let who = voter_at(voter_index).ok_or(_phragmen::Error::CompactInvalidIndex)?;
-				let all_stake: u128 = max_of(&who).into();
-				assignments.push(_phragmen::StakedAssignment {
-					who,
-					distribution: vec![(target_at(target_index).ok_or(_phragmen::Error::CompactInvalidIndex)?, all_stake)],
-				})
-			}
-		)
-	};
-
-	let into_impl_double = {
-		let name = field_name_for(2);
-		quote!(
-			for (voter_index, (t1_idx, w1), t2_idx) in self.#name {
-				let who = voter_at(voter_index).ok_or(_phragmen::Error::CompactInvalidIndex)?;
-				let all_stake: u128 = max_of(&who).into();
-
-				if w1 >= all_stake {
-					return Err(_phragmen::Error::CompactStakeOverflow);
-				}
-
-				// w2 is ensured to be positive.
-				let w2 = all_stake - w1;
-				assignments.push( _phragmen::StakedAssignment {
-					who,
-					distribution: vec![
-						(target_at(t1_idx).ok_or(_phragmen::Error::CompactInvalidIndex)?, w1),
-						(target_at(t2_idx).ok_or(_phragmen::Error::CompactInvalidIndex)?, w2),
-					]
-				});
-			}
-		)
-	};
-
-	let into_impl_rest = (3..=count).map(|c| {
-		let name = field_name_for(c);
-		quote!(
-			for (voter_index, inners, t_last_idx) in self.#name {
-				let who = voter_at(voter_index).ok_or(_phragmen::Error::CompactInvalidIndex)?;
-				let mut sum = u128::min_value();
-				let all_stake: u128 = max_of(&who).into();
-
-				let mut inners_parsed = inners
-					.iter()
-					.map(|(ref t_idx, w)| {
-						sum = sum.saturating_add(*w);
-						let target = target_at(*t_idx).ok_or(_phragmen::Error::CompactInvalidIndex)?;
-						Ok((target, *w))
-					}).collect::<Result<Vec<(A, u128)>, _phragmen::Error>>()?;
-
-				if sum >= all_stake {
-					return Err(_phragmen::Error::CompactStakeOverflow);
-				}
-				// w_last is proved to be positive.
-				let w_last = all_stake - sum;
-
-				inners_parsed.push((target_at(t_last_idx).ok_or(_phragmen::Error::CompactInvalidIndex)?, w_last));
-
-				assignments.push(_phragmen::StakedAssignment {
-					who,
-					distribution: inners_parsed,
-				});
-			}
-		)
-	}).collect::<TokenStream2>();
-
-	quote!(
-		#into_impl_single
-		#into_impl_double
-		#into_impl_rest
-	)
-}
-
-pub(crate) fn staked(
-	ident: syn::Ident,
-	voter_type: GenericArgument,
-	target_type: GenericArgument,
-	count: usize,
-) -> TokenStream2 {
-
-	let from_impl = from_impl(count);
-	let into_impl = into_impl(count);
-
-	quote!(
-		impl<
-			#voter_type: _phragmen::codec::Codec + Default + Copy,
-			#target_type: _phragmen::codec::Codec + Default + Copy,
-		>
-		#ident<#voter_type, #target_type, u128>
-		{
-			/// Generate self from a vector of `StakedAssignment`.
-			pub fn from_staked<FV, FT, A>(
-				assignments: Vec<_phragmen::StakedAssignment<A>>,
-				index_of_voter: FV,
-				index_of_target: FT,
-			) -> Result<Self, _phragmen::Error>
-				where
-					for<'r> FV: Fn(&'r A) -> Option<#voter_type>,
-					for<'r> FT: Fn(&'r A) -> Option<#target_type>,
-					A: _phragmen::IdentifierT
-			{
-				let mut compact: #ident<#voter_type, #target_type, u128> = Default::default();
-				for _phragmen::StakedAssignment { who, distribution }  in assignments {
-					match distribution.len() {
-						0 => continue,
-						#from_impl
-						_ => {
-							return Err(_phragmen::Error::CompactTargetOverflow);
-						}
-					}
-				};
-				Ok(compact)
-			}
-
-			/// Convert self into `StakedAssignment`. The given function should return the total
-			/// weight of a voter. It is used to subtract the sum of all the encoded weights to
-			/// infer the last one.
-			pub fn into_staked<FM, A>(
-				self,
-				max_of: FM,
-				voter_at: impl Fn(#voter_type) -> Option<A>,
-				target_at: impl Fn(#target_type) -> Option<A>,
-			)
-				-> Result<Vec<_phragmen::StakedAssignment<A>>, _phragmen::Error>
-			where
-				for<'r> FM: Fn(&'r A) -> u64,
-				A: _phragmen::IdentifierT,
-			{
-				let mut assignments: Vec<_phragmen::StakedAssignment<A>> = Default::default();
-				#into_impl
-				Ok(assignments)
-			}
-		}
-	)
-}

substrate/primitives/npos-elections/src/lib.rs

@@ -60,8 +60,22 @@ pub use codec;
 #[doc(hidden)]
 pub use sp_arithmetic;
 
+/// Simple Extension trait to easily convert `None` from index closures to `Err`.
+///
+/// This is only generated and re-exported for the compact solution code to use.
+#[doc(hidden)]
+pub trait __OrInvalidIndex<T> {
+	fn or_invalid_index(self) -> Result<T, Error>;
+}
+
+impl<T> __OrInvalidIndex<T> for Option<T> {
+	fn or_invalid_index(self) -> Result<T, Error> {
+		self.ok_or(Error::CompactInvalidIndex)
+	}
+}
+
 // re-export the compact solution type.
-pub use sp_npos_elections_compact::generate_compact_solution_type;
+pub use sp_npos_elections_compact::generate_solution_type;
 
 /// A trait to limit the number of votes per voter. The generated compact type will implement this.
 pub trait VotingLimit {

substrate/primitives/npos-elections/src/tests.rs

@@ -920,28 +920,76 @@ mod score {
 	}
 }
 
-mod compact {
+mod solution_type {
 	use codec::{Decode, Encode};
 	use super::AccountId;
 	// these need to come from the same dev-dependency `sp-npos-elections`, not from the crate.
 	use crate::{
-		generate_compact_solution_type, VoteWeight, Assignment, StakedAssignment,
-		Error as PhragmenError, ExtendedBalance,
+		generate_solution_type, Assignment,
+		Error as PhragmenError,
 	};
-	use sp_std::{convert::{TryInto, TryFrom}, fmt::Debug};
+	use sp_std::{convert::TryInto, fmt::Debug};
 	use sp_arithmetic::Percent;
 
-	type Accuracy = Percent;
+	type TestAccuracy = Percent;
 
-	generate_compact_solution_type!(TestCompact, 16);
+	generate_solution_type!(pub struct TestSolutionCompact::<u32, u8, TestAccuracy>(16));
+
+	#[allow(dead_code)]
+	mod __private {
+		// This is just to make sure that that the compact can be generated in a scope without any
+		// imports.
+		use crate::generate_solution_type;
+		use sp_arithmetic::Percent;
+		generate_solution_type!(
+			#[compact]
+			struct InnerTestSolutionCompact::<u32, u8, Percent>(12)
+		);
+
+	}
 
 	#[test]
-	fn compact_struct_is_codec() {
-		let compact = TestCompact::<_, _, _> {
-			votes1: vec![(2u64, 20), (4, 40)],
+	fn solution_struct_works_with_and_without_compact() {
+		// we use u32 size to make sure compact is smaller.
+		let without_compact = {
+			generate_solution_type!(pub struct InnerTestSolution::<u32, u32, Percent>(16));
+			let compact = InnerTestSolution {
+				votes1: vec![(2, 20), (4, 40)],
+				votes2: vec![
+					(1, (10, TestAccuracy::from_percent(80)), 11),
+					(5, (50, TestAccuracy::from_percent(85)), 51),
+				],
+				..Default::default()
+			};
+
+			compact.encode().len()
+		};
+
+		let with_compact = {
+			generate_solution_type!(#[compact] pub struct InnerTestSolutionCompact::<u32, u32, Percent>(16));
+			let compact = InnerTestSolutionCompact {
+				votes1: vec![(2, 20), (4, 40)],
+				votes2: vec![
+					(1, (10, TestAccuracy::from_percent(80)), 11),
+					(5, (50, TestAccuracy::from_percent(85)), 51),
+				],
+				..Default::default()
+			};
+
+			compact.encode().len()
+		};
+
+		dbg!(with_compact, without_compact);
+		assert!(with_compact < without_compact);
+	}
+
+	#[test]
+	fn solution_struct_is_codec() {
+		let compact = TestSolutionCompact {
+			votes1: vec![(2, 20), (4, 40)],
 			votes2: vec![
-				(1, (10, Accuracy::from_percent(80)), 11),
-				(5, (50, Accuracy::from_percent(85)), 51),
+				(1, (10, TestAccuracy::from_percent(80)), 11),
+				(5, (50, TestAccuracy::from_percent(85)), 51),
 			],
 			..Default::default()
 		};
@@ -956,14 +1004,8 @@ mod compact {
 		assert_eq!(compact.edge_count(), 2 + 4);
 	}
 
-	fn basic_ratio_test_with<V, T>() where
-		V: codec::Codec + Copy + Default + PartialEq + Eq + TryInto<usize> + TryFrom<usize> + From<u8> + Debug,
-		T: codec::Codec + Copy + Default + PartialEq + Eq + TryInto<usize> + TryFrom<usize> + From<u8> + Debug,
-		<V as TryFrom<usize>>::Error: std::fmt::Debug,
-		<T as TryFrom<usize>>::Error: std::fmt::Debug,
-		<V as TryInto<usize>>::Error: std::fmt::Debug,
-		<T as TryInto<usize>>::Error: std::fmt::Debug,
-	{
+	#[test]
+	fn basic_from_and_into_compact_works_assignments() {
 		let voters = vec![
 			2 as AccountId,
 			4,
@@ -986,44 +1028,44 @@ mod compact {
 		let assignments = vec![
 			Assignment {
 				who: 2 as AccountId,
-				distribution: vec![(20u64, Accuracy::from_percent(100))]
+				distribution: vec![(20u64, TestAccuracy::from_percent(100))]
 			},
 			Assignment {
 				who: 4,
-				distribution: vec![(40, Accuracy::from_percent(100))],
+				distribution: vec![(40, TestAccuracy::from_percent(100))],
 			},
 			Assignment {
 				who: 1,
 				distribution: vec![
-					(10, Accuracy::from_percent(80)),
-					(11, Accuracy::from_percent(20))
+					(10, TestAccuracy::from_percent(80)),
+					(11, TestAccuracy::from_percent(20))
 				],
 			},
 			Assignment {
 				who: 5,
 				distribution: vec![
-					(50, Accuracy::from_percent(85)),
-					(51, Accuracy::from_percent(15)),
+					(50, TestAccuracy::from_percent(85)),
+					(51, TestAccuracy::from_percent(15)),
 				]
 			},
 			Assignment {
 				who: 3,
 				distribution: vec![
-					(30, Accuracy::from_percent(50)),
-					(31, Accuracy::from_percent(25)),
-					(32, Accuracy::from_percent(25)),
+					(30, TestAccuracy::from_percent(50)),
+					(31, TestAccuracy::from_percent(25)),
+					(32, TestAccuracy::from_percent(25)),
 				],
 			},
 		];
 
-		let voter_index = |a: &AccountId| -> Option<V> {
+		let voter_index = |a: &AccountId| -> Option<u32> {
 			voters.iter().position(|x| x == a).map(TryInto::try_into).unwrap().ok()
 		};
-		let target_index = |a: &AccountId| -> Option<T> {
+		let target_index = |a: &AccountId| -> Option<u8> {
 			targets.iter().position(|x| x == a).map(TryInto::try_into).unwrap().ok()
 		};
 
-		let compacted = <TestCompact<V, T, Percent>>::from_assignment(
+		let compacted = TestSolutionCompact::from_assignment(
 			assignments.clone(),
 			voter_index,
 			target_index,
@@ -1038,25 +1080,29 @@ mod compact {
 
 		assert_eq!(
 			compacted,
-			TestCompact {
-				votes1: vec![(V::from(0u8), T::from(2u8)), (V::from(1u8), T::from(6u8))],
+			TestSolutionCompact {
+				votes1: vec![(0, 2), (1, 6)],
 				votes2: vec![
-					(V::from(2u8), (T::from(0u8), Accuracy::from_percent(80)), T::from(1u8)),
-					(V::from(3u8), (T::from(7u8), Accuracy::from_percent(85)), T::from(8u8)),
+					(2, (0, TestAccuracy::from_percent(80)), 1),
+					(3, (7, TestAccuracy::from_percent(85)), 8),
 				],
 				votes3: vec![
 					(
-						V::from(4),
-						[(T::from(3u8), Accuracy::from_percent(50)), (T::from(4u8), Accuracy::from_percent(25))],
-						T::from(5u8),
+						4,
+						[(3, TestAccuracy::from_percent(50)), (4, TestAccuracy::from_percent(25))],
+						5,
 					),
 				],
 				..Default::default()
 			}
 		);
 
-		let voter_at = |a: V| -> Option<AccountId> { voters.get(<V as TryInto<usize>>::try_into(a).unwrap()).cloned() };
-		let target_at = |a: T| -> Option<AccountId> { targets.get(<T as TryInto<usize>>::try_into(a).unwrap()).cloned() };
+		let voter_at = |a: u32| -> Option<AccountId> {
+			voters.get(<u32 as TryInto<usize>>::try_into(a).unwrap()).cloned()
+		};
+		let target_at = |a: u8| -> Option<AccountId> {
+			targets.get(<u8 as TryInto<usize>>::try_into(a).unwrap()).cloned()
+		};
 
 		assert_eq!(
 			compacted.into_assignment(voter_at, target_at).unwrap(),
@@ -1064,240 +1110,59 @@ mod compact {
 		);
 	}
 
-	#[test]
-	fn basic_from_and_into_compact_works_assignments() {
-		basic_ratio_test_with::<u16, u16>();
-		basic_ratio_test_with::<u16, u32>();
-		basic_ratio_test_with::<u8, u32>();
-	}
-
-	#[test]
-	fn basic_from_and_into_compact_works_staked_assignments() {
-		let voters = vec![
-			2 as AccountId,
-			4,
-			1,
-			5,
-			3,
-		];
-		let targets = vec![
-			10 as AccountId, 11,
-			20,
-			30, 31, 32,
-			40,
-			50, 51,
-		];
-
-		let assignments = vec![
-			StakedAssignment {
-				who: 2 as AccountId,
-				distribution: vec![(20, 100 as ExtendedBalance)]
-			},
-			StakedAssignment {
-				who: 4,
-				distribution: vec![(40, 100)],
-			},
-			StakedAssignment {
-				who: 1,
-				distribution: vec![
-					(10, 80),
-					(11, 20)
-				],
-			},
-			StakedAssignment {
-				who: 5, distribution:
-				vec![
-					(50, 85),
-					(51, 15),
-				]
-			},
-			StakedAssignment {
-				who: 3,
-				distribution: vec![
-					(30, 50),
-					(31, 25),
-					(32, 25),
-				],
-			},
-		];
-
-		let voter_index = |a: &AccountId| -> Option<u16> {
-			voters.iter().position(|x| x == a).map(TryInto::try_into).unwrap().ok()
-		};
-		let target_index = |a: &AccountId| -> Option<u16> {
-			targets.iter().position(|x| x == a).map(TryInto::try_into).unwrap().ok()
-		};
-
-		let compacted = <TestCompact<u16, u16, ExtendedBalance>>::from_staked(
-			assignments.clone(),
-			voter_index,
-			target_index,
-		).unwrap();
-		assert_eq!(compacted.len(), assignments.len());
-		assert_eq!(
-			compacted.edge_count(),
-			assignments.iter().fold(0, |a, b| a + b.distribution.len()),
-		);
-
-		assert_eq!(
-			compacted,
-			TestCompact {
-				votes1: vec![(0, 2), (1, 6)],
-				votes2: vec![
-					(2, (0, 80), 1),
-					(3, (7, 85), 8),
-				],
-				votes3: vec![
-					(4, [(3, 50), (4, 25)], 5),
-				],
-				..Default::default()
-			}
-		);
-
-		let max_of_fn = |_: &AccountId| -> VoteWeight { 100 };
-		let voter_at = |a: u16| -> Option<AccountId> { voters.get(a as usize).cloned() };
-		let target_at = |a: u16| -> Option<AccountId> { targets.get(a as usize).cloned() };
-
-		assert_eq!(
-			compacted.into_staked(
-				max_of_fn,
-				voter_at,
-				target_at,
-			).unwrap(),
-			assignments,
-		);
-	}
-
-	#[test]
-	fn compact_into_stake_must_report_overflow() {
-		// The last edge which is computed from the rest should ALWAYS be positive.
-		// in votes2
-		let compact = TestCompact::<u16, u16, ExtendedBalance> {
-			votes1: Default::default(),
-			votes2: vec![(0, (1, 10), 2)],
-			..Default::default()
-		};
-
-		let entity_at = |a: u16| -> Option<AccountId> { Some(a as AccountId) };
-		let max_of = |_: &AccountId| -> VoteWeight { 5 };
-
-		assert_eq!(
-			compact.into_staked(&max_of, &entity_at, &entity_at).unwrap_err(),
-			PhragmenError::CompactStakeOverflow,
-		);
-
-		// in votes3 onwards
-		let compact = TestCompact::<u16, u16, ExtendedBalance> {
-			votes1: Default::default(),
-			votes2: Default::default(),
-			votes3: vec![(0, [(1, 7), (2, 8)], 3)],
-			..Default::default()
-		};
-
-		assert_eq!(
-			compact.into_staked(&max_of, &entity_at, &entity_at).unwrap_err(),
-			PhragmenError::CompactStakeOverflow,
-		);
-
-		// Also if equal
-		let compact = TestCompact::<u16, u16, ExtendedBalance> {
-			votes1: Default::default(),
-			votes2: Default::default(),
-			// 5 is total, we cannot leave none for 30 here.
-			votes3: vec![(0, [(1, 3), (2, 2)], 3)],
-			..Default::default()
-		};
-
-		assert_eq!(
-			compact.into_staked(&max_of, &entity_at, &entity_at).unwrap_err(),
-			PhragmenError::CompactStakeOverflow,
-		);
-	}
-
 	#[test]
 	fn compact_into_assignment_must_report_overflow() {
 		// in votes2
-		let compact = TestCompact::<u16, u16, Accuracy> {
+		let compact = TestSolutionCompact {
 			votes1: Default::default(),
-			votes2: vec![(0, (1, Accuracy::from_percent(100)), 2)],
+			votes2: vec![(0, (1, TestAccuracy::from_percent(100)), 2)],
 			..Default::default()
 		};
 
-		let entity_at = |a: u16| -> Option<AccountId> { Some(a as AccountId) };
+		let voter_at = |a: u32| -> Option<AccountId> { Some(a as AccountId) };
+		let target_at = |a: u8| -> Option<AccountId> { Some(a as AccountId) };
+
 
 		assert_eq!(
-			compact.into_assignment(&entity_at, &entity_at).unwrap_err(),
+			compact.into_assignment(&voter_at, &target_at).unwrap_err(),
 			PhragmenError::CompactStakeOverflow,
 		);
 
 		// in votes3 onwards
-		let compact = TestCompact::<u16, u16, Accuracy> {
+		let compact = TestSolutionCompact {
 			votes1: Default::default(),
 			votes2: Default::default(),
-			votes3: vec![(0, [(1, Accuracy::from_percent(70)), (2, Accuracy::from_percent(80))], 3)],
+			votes3: vec![(0, [(1, TestAccuracy::from_percent(70)), (2, TestAccuracy::from_percent(80))], 3)],
 			..Default::default()
 		};
 
 		assert_eq!(
-			compact.into_assignment(&entity_at, &entity_at).unwrap_err(),
+			compact.into_assignment(&voter_at, &target_at).unwrap_err(),
 			PhragmenError::CompactStakeOverflow,
 		);
 	}
 
 	#[test]
 	fn target_count_overflow_is_detected() {
-		let assignments = vec![
-			StakedAssignment {
-				who: 1 as AccountId,
-				distribution: (10..26).map(|i| (i as AccountId, i as ExtendedBalance)).collect::<Vec<_>>(),
-			},
-		];
-
-		let entity_index = |a: &AccountId| -> Option<u16> { Some(*a as u16) };
-
-		let compacted = <TestCompact<u16, u16, ExtendedBalance>>::from_staked(
-			assignments.clone(),
-			entity_index,
-			entity_index,
-		);
-
-		assert!(compacted.is_ok());
-
-		let assignments = vec![
-			StakedAssignment {
-				who: 1 as AccountId,
-				distribution: (10..27).map(|i| (i as AccountId, i as ExtendedBalance)).collect::<Vec<_>>(),
-			},
-		];
-
-		let compacted = <TestCompact<u16, u16, ExtendedBalance>>::from_staked(
-			assignments.clone(),
-			entity_index,
-			entity_index,
-		);
-
-		assert_eq!(
-			compacted.unwrap_err(),
-			PhragmenError::CompactTargetOverflow,
-		);
+		let voter_index = |a: &AccountId| -> Option<u32> { Some(*a as u32) };
+		let target_index = |a: &AccountId| -> Option<u8> { Some(*a as u8) };
 
 		let assignments = vec![
 			Assignment {
 				who: 1 as AccountId,
-				distribution: (10..27).map(|i| (i as AccountId, Percent::from_parts(i as u8))).collect::<Vec<_>>(),
+				distribution:
+					(10..27)
+					.map(|i| (i as AccountId, Percent::from_parts(i as u8)))
+					.collect::<Vec<_>>(),
 			},
 		];
 
-		let compacted = <TestCompact<u16, u16, Percent>>::from_assignment(
+		let compacted = TestSolutionCompact::from_assignment(
 			assignments.clone(),
-			entity_index,
-			entity_index,
-		);
-
-		assert_eq!(
-			compacted.unwrap_err(),
-			PhragmenError::CompactTargetOverflow,
+			voter_index,
+			target_index,
 		);
+		assert_eq!(compacted.unwrap_err(), PhragmenError::CompactTargetOverflow);
 	}
 
 		#[test]
@@ -1306,24 +1171,24 @@ mod compact {
 		let targets = vec![10 as AccountId, 11];
 
 		let assignments = vec![
-			StakedAssignment {
+			Assignment {
 				who: 1 as AccountId,
-				distribution: vec![(10, 100 as ExtendedBalance), (11, 100)]
+				distribution: vec![(10, Percent::from_percent(50)), (11, Percent::from_percent(50))],
 			},
-			StakedAssignment {
+			Assignment {
 				who: 2,
 				distribution: vec![],
 			},
 		];
 
-		let voter_index = |a: &AccountId| -> Option<u16> {
+		let voter_index = |a: &AccountId| -> Option<u32> {
 			voters.iter().position(|x| x == a).map(TryInto::try_into).unwrap().ok()
 		};
-		let target_index = |a: &AccountId| -> Option<u16> {
+		let target_index = |a: &AccountId| -> Option<u8> {
 			targets.iter().position(|x| x == a).map(TryInto::try_into).unwrap().ok()
 		};
 
-		let compacted = <TestCompact<u16, u16, ExtendedBalance>>::from_staked(
+		let compacted = TestSolutionCompact::from_assignment(
 			assignments.clone(),
 			voter_index,
 			target_index,
@@ -1331,9 +1196,9 @@ mod compact {
 
 		assert_eq!(
 			compacted,
-			TestCompact {
+			TestSolutionCompact {
 				votes1: Default::default(),
-				votes2: vec![(0, (0, 100), 1)],
+				votes2: vec![(0, (0, Percent::from_percent(50)), 1)],
 				..Default::default()
 			}
 		);