macro modification

Author: gui1117

frame/support/procedural/src/construct_runtime/mod.rs

@@ -22,19 +22,79 @@ use frame_support_procedural_tools::{generate_crate_access, generate_hidden_incl
 use parse::{ModuleDeclaration, RuntimeDefinition, WhereSection};
 use proc_macro::TokenStream;
 use proc_macro2::TokenStream as TokenStream2;
-use quote::quote;
+use quote::{quote, quote_spanned};
 use syn::{Ident, Result, TypePath};
 
 /// The fixed name of the system module.
 const SYSTEM_MODULE_NAME: &str = "System";
 
 pub fn construct_runtime(input: TokenStream) -> TokenStream {
+	let input_clone = input.clone().into();
 	let definition = syn::parse_macro_input!(input as RuntimeDefinition);
+
+	if let Some(preprocess) = construct_runtime_preprocess(&definition, input_clone)
+		.unwrap_or_else(|e| Some(e.to_compile_error()))
+	{
+		return preprocess.into()
+	}
+
 	construct_runtime_parsed(definition)
 		.unwrap_or_else(|e| e.to_compile_error())
 		.into()
 }
 
+fn construct_runtime_preprocess(
+	def: &RuntimeDefinition,
+	input_clone: TokenStream2,
+) -> Result<Option<TokenStream2>> {
+	let mut auto_modules = vec![];
+	for module in def.modules.content.inner.iter() {
+		if module.module_parts.is_none() {
+			auto_modules.push(module.clone())
+		}
+	}
+
+	if !auto_modules.is_empty() {
+
+		// Make frame-support available to construct_runtime_args
+		let hidden_crate_name = "construct_runtime_preprocess";
+		let scrate_decl = generate_hidden_includes(&hidden_crate_name, "frame-support");
+		let scrate = generate_crate_access(&hidden_crate_name, "frame-support");
+
+		let mut expand = quote!( #scrate::construct_runtime! { #input_clone } );
+
+		while let Some(module) = auto_modules.pop()  {
+			let macro_call = if def.local_macro.as_ref().map_or(false, |m| *m == module.module) {
+				quote!( construct_runtime_args! )
+			} else {
+				let module = &module.module;
+				quote!( #module::construct_runtime_args! )
+			};
+
+			let module_name = &module.name;
+			let module_module = &module.module;
+			let module_instance = module.instance.as_ref()
+				.map(|instance| quote!( ::< #instance > ));
+
+			expand = quote_spanned!(module_name.span() =>
+				#macro_call {
+					{ #module_name : #module_module #module_instance}
+					#expand
+				}
+			);
+		}
+
+		expand = quote!(
+			#scrate_decl
+			#expand
+		);
+
+		Ok(Some(expand))
+	} else {
+		Ok(None)
+	}
+}
+
 fn construct_runtime_parsed(definition: RuntimeDefinition) -> Result<TokenStream2> {
 	let RuntimeDefinition {
 		name,
@@ -164,6 +224,12 @@ fn decl_outer_inherent<'a>(
 		})
 	});
 	quote!(
+		// Prevent UncheckedExtrinsic to print unused warning.
+		const _: () = {
+			#[allow(unused)]
+			type __hidden_use_unchecked_extrinsic = #unchecked_extrinsic;
+		};
+
 		#scrate::impl_outer_inherent!(
 			impl Inherents where Block = #block, UncheckedExtrinsic = #unchecked_extrinsic {
 				#(#modules_tokens)*
@@ -201,6 +267,7 @@ fn decl_outer_config<'a>(
 					#module #(#instance)* #(#generics)*,
 			)
 		});
+
 	quote!(
 		#scrate::sp_runtime::impl_outer_config! {
 			pub struct GenesisConfig for #runtime {
@@ -221,6 +288,7 @@ fn decl_runtime_metadata<'a>(
 			module_declaration.find_part("Module").map(|_| {
 				let filtered_names: Vec<_> = module_declaration
 					.module_parts()
+					.expect("Preprocessing has expanded module parts")
 					.into_iter()
 					.filter(|part| part.name() != "Module")
 					.map(|part| part.ident())

frame/support/procedural/src/construct_runtime/parse.rs

@@ -36,20 +36,44 @@ mod keyword {
 	syn::custom_keyword!(Origin);
 	syn::custom_keyword!(Inherent);
 	syn::custom_keyword!(ValidateUnsigned);
+	syn::custom_keyword!(local_macro);
 }
 
 #[derive(Debug)]
 pub struct RuntimeDefinition {
+	// Specified through optional inner attribute: `#[local_macro(my_pallet)]`
+	pub local_macro: Option<Ident>,
 	pub visibility_token: Token![pub],
 	pub enum_token: Token![enum],
 	pub name: Ident,
 	pub where_section: WhereSection,
 	pub modules: ext::Braces<ext::Punctuated<ModuleDeclaration, Token![,]>>,
 }
 
+pub struct LocalMacroDef(Ident);
+impl Parse for LocalMacroDef {
+	fn parse(input: ParseStream) -> Result<Self> {
+		input.parse::<Token![#]>()?;
+		let attr;
+		syn::bracketed!(attr in input);
+		attr.parse::<keyword::local_macro>()?;
+		let pallet;
+		syn::parenthesized!(pallet in attr);
+		let pallet = pallet.parse()?;
+		Ok(Self(pallet))
+	}
+}
+
 impl Parse for RuntimeDefinition {
 	fn parse(input: ParseStream) -> Result<Self> {
+		let local_macro = if input.peek(Token![#]) {
+			Some(input.parse::<LocalMacroDef>()?.0)
+		} else {
+			None
+		};
+
 		Ok(Self {
+			local_macro,
 			visibility_token: input.parse()?,
 			enum_token: input.parse()?,
 			name: input.parse()?,
@@ -149,12 +173,12 @@ impl Parse for WhereDefinition {
 	}
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct ModuleDeclaration {
 	pub name: Ident,
 	pub module: Ident,
 	pub instance: Option<Ident>,
-	pub module_parts: Vec<ModulePart>,
+	pub module_parts: Option<Vec<ModulePart>>,
 }
 
 impl Parse for ModuleDeclaration {
@@ -172,8 +196,12 @@ impl Parse for ModuleDeclaration {
 			None
 		};
 
-		let _: Token![::] = input.parse()?;
-		let module_parts = parse_module_parts(input)?;
+		let module_parts = if input.peek(Token![::]) {
+			let _: Token![::] = input.parse()?;
+			Some(parse_module_parts(input)?)
+		} else {
+			None
+		};
 
 		let parsed = Self {
 			name,
@@ -188,12 +216,12 @@ impl Parse for ModuleDeclaration {
 
 impl ModuleDeclaration {
 	/// Get resolved module parts
-	pub fn module_parts(&self) -> &[ModulePart] {
-		&self.module_parts
+	pub fn module_parts(&self) -> Option<&Vec<ModulePart>> {
+		self.module_parts.as_ref()
 	}
 
 	pub fn find_part(&self, name: &str) -> Option<&ModulePart> {
-		self.module_parts.iter().find(|part| part.name() == name)
+		self.module_parts.as_ref().and_then(|p| p.iter().find(|part| part.name() == name))
 	}
 
 	pub fn exists_part(&self, name: &str) -> bool {

frame/support/procedural/src/expand_after.rs

@@ -0,0 +1,112 @@
+use proc_macro2::{TokenStream, TokenTree, Group, Span};
+use syn::spanned::Spanned;
+use std::iter::once;
+
+pub fn expand_after(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+	let def = syn::parse_macro_input!(input as ExpandAfterDef);
+	let expand_in_span = def.expand_in.span();
+
+	match expand_in_stream(&def.expand_after, &mut Some(def.expand_with), def.expand_in) {
+		Ok(stream) => stream.into(),
+		Err(_) => {
+			let msg = format!("Cannot find pattern `{:?}` in given token stream", def.expand_after);
+			syn::Error::new(expand_in_span, msg).to_compile_error().into()
+		},
+	}
+}
+
+struct ExpandAfterDef {
+	// This is ensured to have no TokenTree::Group nor TokenTree::Literal and not being empty.
+	expand_after: Vec<TokenTree>,
+	expand_with: TokenStream,
+	expand_in: TokenStream,
+}
+
+impl syn::parse::Parse for ExpandAfterDef {
+	fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+		let expand_after;
+		let _replace_with_bracket: syn::token::Brace = syn::braced!(expand_after in input);
+		let expand_after = expand_after.parse::<TokenStream>()?
+			.into_iter().collect::<Vec<TokenTree>>();
+
+		if let Some(t) = expand_after.iter().find(|t| matches!(t, TokenTree::Group(_))) {
+			return Err(syn::Error::new(t.span(), "Unexpected group token tree"));
+		}
+		if let Some(t) = expand_after.iter().find(|t| matches!(t, TokenTree::Literal(_))) {
+			return Err(syn::Error::new(t.span(), "Unexpected literal token tree"));
+		}
+
+		if expand_after.is_empty() {
+			return Err(syn::Error::new(Span::call_site(), "empty match pattern is invalid"));
+		}
+
+		let expand_with;
+		let _replace_with_bracket: syn::token::Brace = syn::braced!(expand_with in input);
+		let expand_with: TokenStream = expand_with.parse()?;
+
+		Ok(Self {
+			expand_with,
+			expand_after,
+			expand_in: input.parse()?,
+		})
+	}
+}
+
+// Replace the first found `auto` ident by content of `with`.
+// `with` must be some (Option is used for internal simplification).
+// `after` musn't be empty and only contains Ident or Punct
+fn expand_in_stream(
+	after: &Vec<TokenTree>,
+	with: &mut Option<TokenStream>,
+	stream: TokenStream
+) -> Result<TokenStream, ()> {
+	assert!(with.is_some(), "`with` must be some, Option is used because `with` is used only once");
+	assert!(!after.is_empty(), "`after` mustn't be empty, otherwise it cannot be found");
+
+	let mut stream = stream.into_iter();
+	let mut extended = TokenStream::new();
+	let mut match_cursor = 0;
+
+	loop {
+		match stream.next() {
+			Some(TokenTree::Group(group)) => {
+				match_cursor = 0;
+				let stream = group.stream();
+				match expand_in_stream(after, with, stream) {
+					Ok(stream) => {
+						extended.extend(once(TokenTree::Group(Group::new(group.delimiter(), stream))));
+						break;
+					}
+					Err(_) => {
+						extended.extend(once(TokenTree::Group(group)));
+					}
+				}
+			}
+			Some(other) => {
+				use TokenTree::{Ident, Punct};
+
+				let other_match_pattern = match (&other, &after[match_cursor]) {
+					(Ident(i1), Ident(i2)) if i1 == i2 => true,
+					(Punct(p1), Punct(p2)) if p1.as_char() == p2.as_char() => true,
+					_ => false,
+				};
+
+				if other_match_pattern {
+					match_cursor += 1
+				}
+
+				extended.extend(once(other));
+
+				if match_cursor == after.len() {
+					extended.extend(once(with.take().expect("with is used to replace only once")));
+					break;
+				}
+			}
+			None => return Err(())
+		}
+	}
+
+	extended.extend(stream);
+
+	Ok(extended)
+}