Add conditional compilation support for impl_runtime_apis!

primitives/api/proc-macro/src/impl_runtime_apis.rs

@@ -31,10 +31,11 @@ use quote::quote;
 
 use syn::{
 	fold::{self, Fold},
+	parenthesized,
 	parse::{Error, Parse, ParseStream, Result},
 	parse_macro_input, parse_quote,
 	spanned::Spanned,
-	Attribute, Ident, ImplItem, ItemImpl, Path, Signature, Type, TypePath,
+	Attribute, Ident, ImplItem, ItemImpl, LitInt, LitStr, Path, Signature, Type, TypePath,
 };
 
 use std::collections::HashSet;
@@ -67,6 +68,7 @@ fn generate_impl_call(
 	runtime: &Type,
 	input: &Ident,
 	impl_trait: &Path,
+	api_version: &ApiVersion,
 ) -> Result<TokenStream> {
 	let params =
 		extract_parameter_names_types_and_borrows(signature, AllowSelfRefInParameters::No)?;
@@ -111,11 +113,40 @@ fn generate_impl_call(
 		)
 	};
 
+	let fn_calls = if let Some(feature_gated) = &api_version.feature_gated {
+		let pnames = pnames2;
+		let pnames2 = pnames.clone();
+		let pborrow2 = pborrow.clone();
+
+		let feature_name = &feature_gated.0;
+		let impl_trait_fg = extend_with_api_version(impl_trait.clone(), Some(feature_gated.1));
+		let impl_trait = extend_with_api_version(impl_trait.clone(), api_version.custom);
+
+		quote!(
+			#[cfg(feature = #feature_name)]
+			#[allow(deprecated)]
+			let r = <#runtime as #impl_trait_fg>::#fn_name(#( #pborrow #pnames ),*);
+
+			#[cfg(not(feature = #feature_name))]
+			#[allow(deprecated)]
+			let r = <#runtime as #impl_trait>::#fn_name(#( #pborrow2 #pnames2 ),*);
+
+			r
+		)
+	} else {
+		let pnames = pnames2;
+		let impl_trait = extend_with_api_version(impl_trait.clone(), api_version.custom);
+
+		quote!(
+			#[allow(deprecated)]
+			<#runtime as #impl_trait>::#fn_name(#( #pborrow #pnames ),*)
+		)
+	};
+
 	Ok(quote!(
 		#decode_params
 
-		#[allow(deprecated)]
-		<#runtime as #impl_trait>::#fn_name(#( #pborrow #pnames2 ),*)
+		#fn_calls
 	))
 }
 
@@ -130,7 +161,6 @@ fn generate_impl_calls(
 		let trait_api_ver = extract_api_version(&impl_.attrs, impl_.span())?;
 		let impl_trait_path = extract_impl_trait(impl_, RequireQualifiedTraitPath::Yes)?;
 		let impl_trait = extend_with_runtime_decl_path(impl_trait_path.clone());
-		let impl_trait = extend_with_api_version(impl_trait, trait_api_ver);
 		let impl_trait_ident = &impl_trait_path
 			.segments
 			.last()
@@ -139,14 +169,31 @@ fn generate_impl_calls(
 
 		for item in &impl_.items {
 			if let ImplItem::Fn(method) = item {
-				let impl_call =
-					generate_impl_call(&method.sig, &impl_.self_ty, input, &impl_trait)?;
+				let impl_call = generate_impl_call(
+					&method.sig,
+					&impl_.self_ty,
+					input,
+					&impl_trait,
+					&trait_api_ver,
+				)?;
+				let mut attrs = filter_cfg_attrs(&impl_.attrs);
+
+				// If the method has got a `#[cfg(feature = X)]` attribute match it against the impl
+				// block `cfg_attr` add a feature flag for this impl call. Note that the feature
+				// guard is added to `Vec<Attribute>` which is returned as a result and not to the
+				// method itself.
+				if let Some((feature_name, _)) = &trait_api_ver.feature_gated {
+					if method_feature_flag_matches_trait_feature_gate(&method.attrs, &feature_name)?
+					{
+						add_feature_guard(&mut attrs, &feature_name, true)
+					}
+				}
 
 				impl_calls.push((
 					impl_trait_ident.clone(),
 					method.sig.ident.clone(),
 					impl_call,
-					filter_cfg_attrs(&impl_.attrs),
+					attrs,
 				));
 			}
 		}
@@ -441,6 +488,18 @@ fn extend_with_api_version(mut trait_: Path, version: Option<u64>) -> Path {
 	trait_
 }
 
+/// Adds a feature guard to `attributes`.
+///
+/// Depending on `enable`, the feature guard either enables ('feature = "something"`) or disables
+/// (`not(feature = "something")`).
+fn add_feature_guard(attrs: &mut Vec<Attribute>, feature_name: &str, enable: bool) {
+	let attr = match enable {
+		true => parse_quote!(#[cfg(feature = #feature_name)]),
+		false => parse_quote!(#[cfg(not(feature = #feature_name))]),
+	};
+	attrs.push(attr);
+}
+
 /// Generates the implementations of the apis for the runtime.
 fn generate_api_impl_for_runtime(impls: &[ItemImpl]) -> Result<TokenStream> {
 	let mut impls_prepared = Vec::new();
@@ -451,12 +510,29 @@ fn generate_api_impl_for_runtime(impls: &[ItemImpl]) -> Result<TokenStream> {
 		let trait_api_ver = extract_api_version(&impl_.attrs, impl_.span())?;
 
 		let mut impl_ = impl_.clone();
+		impl_.attrs = filter_cfg_attrs(&impl_.attrs);
+
 		let trait_ = extract_impl_trait(&impl_, RequireQualifiedTraitPath::Yes)?.clone();
 		let trait_ = extend_with_runtime_decl_path(trait_);
-		let trait_ = extend_with_api_version(trait_, trait_api_ver);
+		// If the trait api contains feature gated version - there are staging methods in it. Handle
+		// them explicitly here by adding staging implementation with `#cfg(feature = ...)` and
+		// stable implementation with `#[cfg(not(feature = ...))]`.
+		if let Some(feature_gated) = trait_api_ver.feature_gated {
+			let mut feature_gated_impl = impl_.clone();
+			add_feature_guard(&mut feature_gated_impl.attrs, &feature_gated.0, true);
+			feature_gated_impl.trait_.as_mut().unwrap().1 =
+				extend_with_api_version(trait_.clone(), Some(feature_gated.1));
+
+			impls_prepared.push(feature_gated_impl);
+
+			// Finally add `#[cfg(not(feature = ...))]` for the stable implementation (which is
+			// generated outside this if).
+			add_feature_guard(&mut impl_.attrs, &feature_gated.0, false);
+		}
 
+		// Generate stable trait implementation.
+		let trait_ = extend_with_api_version(trait_, trait_api_ver.custom);
 		impl_.trait_.as_mut().unwrap().1 = trait_;
-		impl_.attrs = filter_cfg_attrs(&impl_.attrs);
 		impls_prepared.push(impl_);
 	}
 
@@ -663,7 +739,8 @@ fn generate_runtime_api_versions(impls: &[ItemImpl]) -> Result<TokenStream> {
 	let c = generate_crate_access();
 
 	for impl_ in impls {
-		let api_ver = extract_api_version(&impl_.attrs, impl_.span())?.map(|a| a as u32);
+		let versions = extract_api_version(&impl_.attrs, impl_.span())?;
+		let api_ver = versions.custom.map(|a| a as u32);
 
 		let mut path = extend_with_runtime_decl_path(
 			extract_impl_trait(impl_, RequireQualifiedTraitPath::Yes)?.clone(),
@@ -687,11 +764,34 @@ fn generate_runtime_api_versions(impls: &[ItemImpl]) -> Result<TokenStream> {
 		}
 
 		let id: Path = parse_quote!( #path ID );
-		let version = quote!( #path VERSION );
-		let attrs = filter_cfg_attrs(&impl_.attrs);
+		let mut attrs = filter_cfg_attrs(&impl_.attrs);
+
+		// Handle API versioning
+		// If feature gated version is set - handle it first
+		if let Some(feature_gated) = versions.feature_gated {
+			let feature_gated_version = feature_gated.1 as u32;
+			// the attributes for the feature gated staging api
+			let mut feature_gated_attrs = attrs.clone();
+			add_feature_guard(&mut feature_gated_attrs, &feature_gated.0, true);
+			populate_runtime_api_versions(
+				&mut result,
+				&mut sections,
+				feature_gated_attrs,
+				id.clone(),
+				quote!( #feature_gated_version ),
+				&c,
+			);
+
+			// Add `#[cfg(not(feature ...))]` to the initial attributes. If the staging feature flag
+			// is not set we want to set the stable api version
+			add_feature_guard(&mut attrs, &feature_gated.0, false);
+		}
 
-		let api_ver = api_ver.map(|a| quote!( #a )).unwrap_or_else(|| version);
-		populate_runtime_api_versions(&mut result, &mut sections, attrs, id, api_ver, &c)
+		// Now add the stable api version to the versions list. If the api has got staging functions
+		// there might be a `#[cfg(not(feature ...))]` attribute attached to the stable version.
+		let base_api_version = quote!( #path VERSION );
+		let api_ver = api_ver.map(|a| quote!( #a )).unwrap_or_else(|| base_api_version);
+		populate_runtime_api_versions(&mut result, &mut sections, attrs, id, api_ver, &c);
 	}
 
 	Ok(quote!(
@@ -758,12 +858,62 @@ fn filter_cfg_attrs(attrs: &[Attribute]) -> Vec<Attribute> {
 	attrs.iter().filter(|a| a.path().is_ident("cfg")).cloned().collect()
 }
 
+/// Parse feature flagged api_version.
+/// E.g. `#[cfg_attr(feature = "enable-staging-api", api_version(99))]`
+fn extract_cfg_api_version(attrs: &Vec<Attribute>, span: Span) -> Result<Option<(String, u64)>> {
+	let cfg_attrs = attrs.iter().filter(|a| a.path().is_ident("cfg_attr")).collect::<Vec<_>>();
+
+	let mut cfg_api_version_attr = Vec::new();
+	for cfg_attr in cfg_attrs {
+		let mut feature_name = None;
+		let mut api_version = None;
+		cfg_attr.parse_nested_meta(|m| {
+			if m.path.is_ident("feature") {
+				let a = m.value()?;
+				let b: LitStr = a.parse()?;
+				feature_name = Some(b.value());
+			} else if m.path.is_ident(API_VERSION_ATTRIBUTE) {
+				let content;
+				parenthesized!(content in m.input);
+				let ver: LitInt = content.parse()?;
+				api_version = Some(ver.base10_parse::<u64>()?);
+			}
+			Ok(())
+		})?;
+
+		// If there is a cfg attribute containing api_version - save if for processing
+		if let (Some(feature_name), Some(api_version)) = (feature_name, api_version) {
+			cfg_api_version_attr.push((feature_name, api_version, cfg_attr.span()));
+		}
+	}
+
+	if cfg_api_version_attr.len() > 1 {
+		let mut err = Error::new(span, format!("Found multiple feature gated api versions (cfg attribute with nested `{}` attribute). This is not supported.", API_VERSION_ATTRIBUTE));
+		for (_, _, attr_span) in cfg_api_version_attr {
+			err.combine(Error::new(attr_span, format!("`{}` found here", API_VERSION_ATTRIBUTE)));
+		}
+
+		return Err(err)
+	}
+
+	let result = cfg_api_version_attr.drain(..).next().map(|(feature, name, _)| (feature, name));
+	Ok(result)
+}
+
+/// Represents an API version.
+struct ApiVersion {
+	/// Corresponds to `#[api_version(X)]` attribute.
+	pub custom: Option<u64>,
+	/// Corresponds to `#[cfg_attr(feature = "enable-staging-api", api_version(99))]`
+	/// attribute. `String` is the feature name, `u64` the staging api version.
+	pub feature_gated: Option<(String, u64)>,
+}
+
 // Extracts the value of `API_VERSION_ATTRIBUTE` and handles errors.
 // Returns:
 // - Err if the version is malformed
-// - Some(u64) if the version is set
-// - None if the version is not set (this is valid).
-fn extract_api_version(attrs: &Vec<Attribute>, span: Span) -> Result<Option<u64>> {
+// - `ApiVersion` on success. If a version is set or not is determined by the fields of `ApiVersion`
+fn extract_api_version(attrs: &Vec<Attribute>, span: Span) -> Result<ApiVersion> {
 	// First fetch all `API_VERSION_ATTRIBUTE` values (should be only one)
 	let api_ver = attrs
 		.iter()
@@ -782,7 +932,41 @@ fn extract_api_version(attrs: &Vec<Attribute>, span: Span) -> Result<Option<u64>
 	}
 
 	// Parse the runtime version if there exists one.
-	api_ver.first().map(|v| parse_runtime_api_version(v)).transpose()
+	Ok(ApiVersion {
+		custom: api_ver.first().map(|v| parse_runtime_api_version(v)).transpose()?,
+		feature_gated: extract_cfg_api_version(attrs, span)?,
+	})
+}
+
+/// This function looks for a matching trait `cfg_attr` and method feature flag.
+/// It looks for a `#[feature = X]` attribute in `attrs`. Returns `true` if X matches
+/// `trait_feature_gate`.
+fn method_feature_flag_matches_trait_feature_gate(
+	attrs: &Vec<Attribute>,
+	trait_feature_gate: &String,
+) -> Result<bool> {
+	let cfg_attrs = attrs.iter().filter(|a| a.path().is_ident("cfg")).collect::<Vec<_>>();
+
+	for cfg_attr in cfg_attrs {
+		let mut feature_name = None;
+		cfg_attr.parse_nested_meta(|m| {
+			if m.path.is_ident("feature") {
+				let a = m.value()?;
+				let b: LitStr = a.parse()?;
+				feature_name = Some(b.value());
+			}
+			Ok(())
+		})?;
+
+		// If there is a cfg attribute containing api_version - save if for processing
+		if let Some(feature_name) = feature_name {
+			if feature_name == *trait_feature_gate {
+				return Ok(true)
+			}
+		}
+	}
+
+	Ok(false)
 }
 
 #[cfg(test)]

primitives/api/src/lib.rs

@@ -344,6 +344,54 @@ pub use sp_api_proc_macro::decl_runtime_apis;
 /// ```
 /// In this case `Balance` api version 3 is being implemented for `Runtime`. The `impl` block
 /// must contain all methods declared in version 3 and below.
+///
+/// # Conditional version implementation
+///
+/// `impl_runtime_apis!` supports `cfg_attr` attribute for conditional compilation. For example
+/// let's say you want to implement a staging version of the runtime api and put it behind a
+/// feature flag. You can do it this way:
+/// ```ignore
+/// pub struct Runtime {}
+/// sp_api::decl_runtime_apis! {
+///     pub trait ApiWithStagingMethod {
+///         fn stable_one(data: u64);
+///
+///         #[api_version(99)]
+///         fn staging_one();
+///     }
+/// }
+///
+/// sp_api::impl_runtime_apis! {
+///     #[cfg_attr(feature = "enable-staging-api", api_version(99))]
+///     impl self::ApiWithStagingMethod<Block> for Runtime {
+///         fn stable_one(_: u64) {}
+///
+///         #[cfg(feature = "enable-staging-api")]
+///         fn staging_one() {}
+///     }
+/// }
+/// ```
+///
+/// [`decl_runtime_apis!`] declares two version of the api - 1 (the default one, which is
+/// considered stable in our example) and 99 (which is considered staging). In
+/// `impl_runtime_apis!` a `cfg_attr` attribute is attached to the `ApiWithStagingMethod`
+/// implementation. If the code is compiled with  `enable-staging-api` feature a version 99 of
+/// the runtime api will be built which will include `staging_one`. Note that `staging_one`
+/// implementation is feature gated by `#[cfg(feature = ... )]` attribute.
+///
+/// If the code is compiled without `enable-staging-api` version 1 (the default one) will be
+/// built which doesn't include `staging_one`.
+///
+/// `cfg_attr` can also be used together with `api_version`. For the next snippet will build
+/// version 99 if `enable-staging-api` is enabled and version 2 otherwise because both
+/// `cfg_attr` and `api_version` are attached to the impl block:
+/// ```ignore
+/// #[cfg_attr(feature = "enable-staging-api", api_version(99))]
+/// #[api_version(2)]
+/// impl self::ApiWithStagingAndVersionedMethods<Block> for Runtime {
+///  // impl skipped
+/// }
+/// ```
 pub use sp_api_proc_macro::impl_runtime_apis;
 
 /// Mocks given trait implementations as runtime apis.

primitives/api/test/Cargo.toml

@@ -35,3 +35,6 @@ static_assertions = "1.1.0"
 [[bench]]
 name = "bench"
 harness = false
+
+[features]
+"enable-staging-api" = []