Improve handling of the generic baggage fields (#5656)

* Parse generic baggage types more carefully to preserve inner structure

* Add example

* Way too many clones

* Allow multiple generic arguments for baggage fields

* Try to detect errors earlier

* Support more types for the baggage fields, get rid of the path constraint

polkadot/node/orchestra/examples/duo.rs

@@ -18,7 +18,7 @@
 //! A dummy to be used with cargo expand
 
 use orchestra::{self as orchestra, Spawner, *};
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 mod misc;
 
 pub use self::misc::*;
@@ -61,7 +61,7 @@ impl<Context> Fortified {
 }
 
 #[orchestra(signal=SigSigSig, event=EvX, error=Yikes, gen=AllMessages)]
-struct Duo<T> {
+struct Duo<T, U, V, W> {
 	#[subsystem(consumes: MsgStrukt, sends: [Plinko])]
 	sub0: Awesome,
 
@@ -69,19 +69,21 @@ struct Duo<T> {
 	plinkos: GoblinTower,
 
 	i_like_pi: f64,
-	i_like_generic: T,
-	i_like_hash: HashMap<f64, f64>,
+	i_like_tuple: (f64, f64),
+	i_like_generic: Arc<T>,
+	i_like_hash: HashMap<(U, V), Arc<W>>,
 }
 
 fn main() {
 	use futures::{executor, pin_mut};
 
 	executor::block_on(async move {
-		let (orchestra, _handle): (Duo<_, f64>, _) = Duo::builder()
+		let (orchestra, _handle): (Duo<_, f64, u32, f32, f64>, _) = Duo::builder()
 			.sub0(AwesomeSubSys::default())
 			.plinkos(Fortified::default())
 			.i_like_pi(::std::f64::consts::PI)
-			.i_like_generic(42.0)
+			.i_like_tuple((::std::f64::consts::PI, ::std::f64::consts::PI))
+			.i_like_generic(Arc::new(42.0))
 			.i_like_hash(HashMap::new())
 			.spawner(DummySpawner)
 			.build()

polkadot/node/orchestra/proc-macro/Cargo.toml

@@ -20,7 +20,7 @@ proc-macro2 = "1.0.37"
 proc-macro-crate = "1.1.3"
 expander = { version = "0.0.6", default-features = false }
 petgraph = "0.6.0"
-itertools = { version = "0.10.3", optional = true }
+itertools = { version = "0.10.3" }
 
 [dev-dependencies]
 assert_matches = "1.5"
@@ -35,4 +35,4 @@ default = [] # enable "graph" by default, blocked by <https://github.com/parityt
 expand = []
 # Create directional message consuming / outgoing graph.
 # Generates: `${OUT_DIR}/${orchestra|lowercase}-subsystem-messaging.dot`
-graph = ["itertools"]
+graph = []

polkadot/node/orchestra/proc-macro/src/impl_builder.rs

@@ -14,7 +14,7 @@
 // limitations under the License.
 
 use quote::{format_ident, quote};
-use syn::{parse_quote, Path, PathSegment};
+use syn::{parse_quote, Path, PathSegment, TypePath};
 
 use super::*;
 
@@ -89,7 +89,12 @@ pub(crate) fn impl_builder(info: &OrchestraInfo) -> proc_macro2::TokenStream {
 	let field_name = subsystem_name.iter().chain(baggage_name.iter()).collect::<Vec<_>>();
 	let field_type = subsystem_generics
 		.iter()
-		.map(|ident| Path::from(PathSegment::from(ident.clone())))
+		.map(|ident| {
+			syn::Type::Path(TypePath {
+				qself: None,
+				path: Path::from(PathSegment::from(ident.clone())),
+			})
+		})
 		.chain(info.baggage().iter().map(|bag| bag.field_ty.clone()))
 		.collect::<Vec<_>>();
 
@@ -249,14 +254,9 @@ pub(crate) fn impl_builder(info: &OrchestraInfo) -> proc_macro2::TokenStream {
 		post_setter_generics[idx] = parse_quote! { Init<#field_type> };
 
 		// Baggage can also be generic, so we need to include that to a signature
-		let preserved_baggage_generic = if bag_field.generic {
-			quote! {#field_type,}
-		} else {
-			TokenStream::new()
-		};
-
+		let preserved_baggage_generics = &bag_field.generic_types;
 		quote! {
-			impl <InitStateSpawner, #preserved_baggage_generic #( #subsystem_passthrough_state_generics, )* #( #impl_baggage_state_generics, )* >
+			impl <InitStateSpawner, #( #preserved_baggage_generics, )* #( #subsystem_passthrough_state_generics, )* #( #impl_baggage_state_generics, )* >
 			#builder <InitStateSpawner, #( #subsystem_passthrough_state_generics, )* #( #pre_setter_generics, )* >
 			{
 				/// Specify the baggage in the builder when it was not initialized before
@@ -278,7 +278,7 @@ pub(crate) fn impl_builder(info: &OrchestraInfo) -> proc_macro2::TokenStream {
 					}
 				}
 			}
-			impl <InitStateSpawner, #preserved_baggage_generic #( #subsystem_passthrough_state_generics, )* #( #impl_baggage_state_generics, )* >
+			impl <InitStateSpawner, #( #preserved_baggage_generics, )* #( #subsystem_passthrough_state_generics, )* #( #impl_baggage_state_generics, )* >
 			#builder <InitStateSpawner, #( #subsystem_passthrough_state_generics, )* #( #post_setter_generics, )* > {
 				/// Specify the baggage in the builder when it has been previously initialized
 				pub fn #fname (self, var: #field_type ) ->

polkadot/node/orchestra/proc-macro/src/parse/parse_orchestra_struct.rs

@@ -13,6 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use itertools::Itertools;
 use proc_macro2::{Span, TokenStream};
 use std::collections::{hash_map::RandomState, HashMap, HashSet};
 use syn::{
@@ -21,8 +22,8 @@ use syn::{
 	punctuated::Punctuated,
 	spanned::Spanned,
 	token::Bracket,
-	AttrStyle, Error, Field, FieldsNamed, GenericParam, Ident, ItemStruct, Path, Result, Token,
-	Type, Visibility,
+	AttrStyle, Error, Field, FieldsNamed, GenericParam, Ident, ItemStruct, Path, PathSegment,
+	Result, Token, Type, Visibility,
 };
 
 use quote::{quote, ToTokens};
@@ -106,13 +107,73 @@ pub(crate) struct SubSysField {
 	pub(crate) wip: bool,
 }
 
-fn try_type_to_path(ty: Type, span: Span) -> Result<Path> {
+// Converts a type enum to a path if this type is a TypePath
+fn try_type_to_path(ty: &Type, span: Span) -> Result<Path> {
 	match ty {
-		Type::Path(path) => Ok(path.path),
+		Type::Path(path) => Ok(path.path.clone()),
 		_ => Err(Error::new(span, "Type must be a path expression.")),
 	}
 }
 
+// Converts a Rust type to a list of idents recursively checking the possible values
+fn flatten_type(ty: &Type, span: Span) -> Result<Vec<Ident>> {
+	match ty {
+		syn::Type::Array(ar) => flatten_type(&ar.elem, span),
+		syn::Type::Paren(par) => flatten_type(&par.elem, span),
+		syn::Type::Path(type_path) => type_path
+			.path
+			.segments
+			.iter()
+			.map(|seg| flatten_path_segments(seg, span.clone()))
+			.flatten_ok()
+			.collect::<Result<Vec<_>>>(),
+		syn::Type::Tuple(tup) => tup
+			.elems
+			.iter()
+			.map(|element| flatten_type(element, span.clone()))
+			.flatten_ok()
+			.collect::<Result<Vec<_>>>(),
+		_ => Err(Error::new(span, format!("Unsupported type: {:?}", ty))),
+	}
+}
+
+// Flatten segments of some path to a list of idents used in these segments
+fn flatten_path_segments(path_segment: &PathSegment, span: Span) -> Result<Vec<Ident>> {
+	let mut result = vec![path_segment.ident.clone()];
+
+	match &path_segment.arguments {
+		syn::PathArguments::AngleBracketed(args) => {
+			let mut recursive_idents = args
+				.args
+				.iter()
+				.map(|generic_argument| match generic_argument {
+					syn::GenericArgument::Type(ty) => flatten_type(ty, span.clone()),
+					_ => Err(Error::new(
+						span,
+						format!(
+							"Field has a generic with an unsupported parameter {:?}",
+							generic_argument
+						),
+					)),
+				})
+				.flatten_ok()
+				.collect::<Result<Vec<_>>>()?;
+			result.append(&mut recursive_idents);
+		},
+		syn::PathArguments::None => {},
+		_ =>
+			return Err(Error::new(
+				span,
+				format!(
+					"Field has a generic with an unsupported path {:?}",
+					path_segment.arguments
+				),
+			)),
+	}
+
+	Ok(result)
+}
+
 macro_rules! extract_variant {
 	($unique:expr, $variant:ident ; default = $fallback:expr) => {
 		extract_variant!($unique, $variant).unwrap_or_else(|| $fallback)
@@ -254,8 +315,8 @@ impl Parse for SubSystemAttrItems {
 #[derive(Debug, Clone)]
 pub(crate) struct BaggageField {
 	pub(crate) field_name: Ident,
-	pub(crate) field_ty: Path,
-	pub(crate) generic: bool,
+	pub(crate) field_ty: Type,
+	pub(crate) generic_types: Vec<Ident>,
 	pub(crate) vis: Visibility,
 }
 
@@ -359,8 +420,7 @@ impl OrchestraInfo {
 	pub(crate) fn baggage_generic_types(&self) -> Vec<Ident> {
 		self.baggage
 			.iter()
-			.filter(|bag| bag.generic)
-			.filter_map(|bag| bag.field_ty.get_ident().cloned())
+			.flat_map(|bag| bag.generic_types.clone())
 			.collect::<Vec<_>>()
 	}
 
@@ -423,7 +483,7 @@ impl OrchestraGuts {
 			let ident = ident.ok_or_else(|| {
 				Error::new(
 					ty.span(),
-					"Missing identifier for field, only named fields are expceted.",
+					"Missing identifier for field, only named fields are expected.",
 				)
 			})?;
 
@@ -442,7 +502,7 @@ impl OrchestraGuts {
 				let attr_tokens = attr_tokens.clone();
 				let subsystem_attrs: SubSystemAttrItems = syn::parse2(attr_tokens.clone())?;
 
-				let field_ty = try_type_to_path(ty, span)?;
+				let field_ty = try_type_to_path(&ty, span)?;
 				let generic = field_ty
 					.get_ident()
 					.ok_or_else(|| {
@@ -484,12 +544,13 @@ impl OrchestraGuts {
 					blocking,
 				});
 			} else {
-				let field_ty = try_type_to_path(ty, ident.span())?;
-				let generic = field_ty
-					.get_ident()
-					.map(|ident| baggage_generics.contains(ident))
-					.unwrap_or(false);
-				baggage.push(BaggageField { field_name: ident, generic, field_ty, vis });
+				let flattened = flatten_type(&ty, ident.span())?;
+				let generic_types = flattened
+					.iter()
+					.filter(|flat_ident| baggage_generics.contains(flat_ident))
+					.cloned()
+					.collect::<Vec<_>>();
+				baggage.push(BaggageField { field_name: ident, generic_types, field_ty: ty, vis });
 			}
 		}
 		Ok(Self { name, subsystems, baggage })
@@ -503,7 +564,7 @@ impl Parse for OrchestraGuts {
 			syn::Fields::Named(named) => {
 				let name = ds.ident.clone();
 
-				// collect the indepedentent subsystem generics
+				// collect the independent subsystem generics
 				// which need to be carried along, there are the non-generated ones
 				let mut orig_generics = ds.generics;