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/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)
-			}
-		}
-	)
-}