[proc macros]: support generic type params

examples/proc_macro.rs

@@ -26,33 +26,40 @@
 
 use jsonrpsee::{
 	proc_macros::rpc,
-	types::{async_trait, error::Error},
+	types::{async_trait, error::Error, Subscription},
 	ws_client::WsClientBuilder,
-	ws_server::{RpcModule, SubscriptionSink, WsServerBuilder},
+	ws_server::{SubscriptionSink, WsServerBuilder},
 };
 use std::net::SocketAddr;
 
+type ExampleHash = [u8; 32];
+type ExampleStorageKey = Vec<u8>;
+
 #[rpc(server, client, namespace = "state")]
-pub trait Rpc {
+pub trait Rpc<Hash: std::fmt::Debug, StorageKey> {
 	/// Async method call example.
-	#[method(name = "getPairs")]
-	async fn storage_pairs(&self, prefix: usize, hash: Option<u128>) -> Result<Vec<usize>, Error>;
+	#[method(name = "getKeys")]
+	async fn storage_keys(&self, storage_key: StorageKey, hash: Option<Hash>) -> Result<Vec<StorageKey>, Error>;
 
-	/// Subscription that take `Option<Vec<u8>>` as input and produces output `Vec<usize>`.
-	#[subscription(name = "subscribeStorage", unsub = "unsubscribeStorage", item = Vec<usize>)]
-	fn subscribe_storage(&self, keys: Option<Vec<u8>>);
+	/// Subscription that takes a `StorageKey` as input and produces a `Vec<Hash>`.
+	#[subscription(name = "subscribeStorage", unsub = "unsubscribeStorage", item = Vec<Hash>)]
+	fn subscribe_storage(&self, keys: Option<Vec<StorageKey>>);
 }
 
 pub struct RpcServerImpl;
 
 #[async_trait]
-impl RpcServer for RpcServerImpl {
-	async fn storage_pairs(&self, _prefix: usize, _hash: Option<u128>) -> Result<Vec<usize>, Error> {
-		Ok(vec![1, 2, 3, 4])
+impl RpcServer<ExampleHash, ExampleStorageKey> for RpcServerImpl {
+	async fn storage_keys(
+		&self,
+		storage_key: ExampleStorageKey,
+		_hash: Option<ExampleHash>,
+	) -> Result<Vec<ExampleStorageKey>, Error> {
+		Ok(vec![storage_key])
 	}
 
-	fn subscribe_storage(&self, mut sink: SubscriptionSink, keys: Option<Vec<u8>>) {
-		sink.send(&keys.unwrap_or_default()).unwrap();
+	fn subscribe_storage(&self, mut sink: SubscriptionSink, _keys: Option<Vec<ExampleStorageKey>>) {
+		sink.send(&vec![[0; 32]]).unwrap();
 	}
 }
 
@@ -64,18 +71,17 @@ async fn main() -> anyhow::Result<()> {
 	let url = format!("ws://{}", server_addr);
 
 	let client = WsClientBuilder::default().build(&url).await?;
-	assert_eq!(client.storage_pairs(10, None).await.unwrap(), vec![1, 2, 3, 4]);
+	assert_eq!(client.storage_keys(vec![1, 2, 3, 4], None::<ExampleHash>).await.unwrap(), vec![vec![1, 2, 3, 4]]);
 
-	let mut sub = client.subscribe_storage(None).await.unwrap();
-	assert_eq!(Some(vec![]), sub.next().await.unwrap());
+	let mut sub: Subscription<Vec<ExampleHash>> =
+		RpcClient::<ExampleHash, ExampleStorageKey>::subscribe_storage(&client, None).await.unwrap();
+	assert_eq!(Some(vec![[0; 32]]), sub.next().await.unwrap());
 
 	Ok(())
 }
 
 async fn run_server() -> anyhow::Result<SocketAddr> {
 	let server = WsServerBuilder::default().build("127.0.0.1:0").await?;
-	let mut module = RpcModule::new(());
-	module.register_method("state_getPairs", |_, _| Ok(vec![1, 2, 3]))?;
 
 	let addr = server.local_addr()?;
 	tokio::spawn(async move { server.start(RpcServerImpl.into_rpc()).await });

proc-macros/Cargo.toml

@@ -16,7 +16,7 @@ proc-macro = true
 Inflector = "0.11.4"
 proc-macro2 = "1.0"
 quote = "1.0"
-syn = { version = "1.0", default-features = false, features = ["extra-traits", "full"] }
+syn = { version = "1.0", default-features = false, features = ["extra-traits", "full", "visit", "parsing"] }
 proc-macro-crate = "1"
 bae = "0.1.6"
 

proc-macros/src/helpers.rs

@@ -24,9 +24,12 @@
 // IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
+use crate::visitor::{FindSubTyParams, FindTyParams};
 use proc_macro2::Span;
 use proc_macro_crate::{crate_name, FoundCrate};
 use quote::quote;
+use std::collections::HashSet;
+use syn::{parse_quote, punctuated::Punctuated, visit::Visit, GenericParam, Generics, Token};
 
 /// Search for client-side `jsonrpsee` in `Cargo.toml`.
 pub(crate) fn find_jsonrpsee_client_crate() -> Result<proc_macro2::TokenStream, syn::Error> {
@@ -57,3 +60,84 @@ fn find_jsonrpsee_crate(http_name: &str, ws_name: &str) -> Result<proc_macro2::T
 		},
 	}
 }
+
+pub(crate) fn client_add_trait_bounds(item_trait: &syn::ItemTrait, sub_tys: &[syn::Type]) -> Generics {
+	let visitor = visit_trait(item_trait, sub_tys);
+	let mut generics = item_trait.generics.clone();
+
+	for param in &mut generics.params {
+		if let GenericParam::Type(ty) = param {
+			ty.bounds.push(parse_quote!(Send));
+			ty.bounds.push(parse_quote!(Sync));
+			ty.bounds.push(parse_quote!('static));
+
+			if visitor.input_params.contains(&ty.ident) {
+				ty.bounds.push(parse_quote!(jsonrpsee::types::Serialize))
+			}
+			if visitor.ret_params.contains(&ty.ident) {
+				ty.bounds.push(parse_quote!(jsonrpsee::types::DeserializeOwned))
+			}
+			if visitor.sub_params.contains(&ty.ident) {
+				ty.bounds.push(parse_quote!(jsonrpsee::types::DeserializeOwned))
+			}
+		}
+	}
+	generics
+}
+
+pub(crate) fn server_generate_where_clause(
+	item_trait: &syn::ItemTrait,
+	sub_tys: &[syn::Type],
+) -> Vec<syn::WherePredicate> {
+	let visitor = visit_trait(item_trait, sub_tys);
+	let additional_where_clause = item_trait.generics.where_clause.clone();
+
+	item_trait
+		.generics
+		.type_params()
+		.map(|ty| {
+			let ty_path = syn::TypePath { qself: None, path: ty.ident.clone().into() };
+			let mut bounds: Punctuated<syn::TypeParamBound, Token![+]> = parse_quote!(Send + Sync + 'static);
+
+			if visitor.input_params.contains(&ty.ident) {
+				bounds.push(parse_quote!(jsonrpsee::types::DeserializeOwned))
+			}
+
+			if visitor.ret_params.contains(&ty.ident) {
+				bounds.push(parse_quote!(jsonrpsee::types::Serialize))
+			}
+
+			if visitor.sub_params.contains(&ty.ident) {
+				bounds.push(parse_quote!(jsonrpsee::types::Serialize))
+			}
+
+			// Add the trait bounds specified in the trait.
+			if let Some(ref where_clause) = additional_where_clause {
+				for predicate in where_clause.predicates.iter() {
+					if let syn::WherePredicate::Type(where_ty) = predicate {
+						if let syn::Type::Path(ref predicate) = where_ty.bounded_ty {
+							if *predicate == ty_path {
+								bounds.extend(where_ty.bounds.clone().into_iter());
+							}
+						}
+					}
+				}
+			}
+
+			syn::WherePredicate::Type(syn::PredicateType {
+				lifetimes: None,
+				bounded_ty: syn::Type::Path(ty_path),
+				colon_token: <Token![:]>::default(),
+				bounds,
+			})
+		})
+		.collect()
+}
+
+fn visit_trait(item_trait: &syn::ItemTrait, sub_tys: &[syn::Type]) -> FindTyParams {
+	let type_params: HashSet<_> = item_trait.generics.type_params().map(|t| t.ident.clone()).collect();
+	let sub_tys = FindSubTyParams::new(type_params).visit(sub_tys);
+	let mut visitor = FindTyParams::new(sub_tys);
+	visitor.visit_item_trait(item_trait);
+	visitor
+}

proc-macros/src/lib.rs

@@ -37,6 +37,7 @@ mod render_client;
 mod render_server;
 mod respan;
 mod rpc_macro;
+pub(crate) mod visitor;
 
 /// Main RPC macro.
 ///

proc-macros/src/render_client.rs

@@ -24,17 +24,21 @@
 // IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
-use crate::rpc_macro::{RpcMethod, RpcSubscription};
-use crate::RpcDescription;
-
+use crate::helpers::client_add_trait_bounds;
+use crate::rpc_macro::{RpcDescription, RpcMethod, RpcSubscription};
 use proc_macro2::TokenStream as TokenStream2;
 use quote::quote;
+use syn::TypeParam;
 
 impl RpcDescription {
 	pub(super) fn render_client(&self) -> Result<TokenStream2, syn::Error> {
 		let jsonrpsee = self.jsonrpsee_client_path.as_ref().unwrap();
+		let sub_tys: Vec<syn::Type> = self.subscriptions.clone().into_iter().map(|s| s.item).collect();
 
 		let trait_name = quote::format_ident!("{}Client", &self.trait_def.ident);
+		let tweaked_generics = client_add_trait_bounds(&self.trait_def, &sub_tys);
+		let type_idents = tweaked_generics.type_params().collect::<Vec<&TypeParam>>();
+		let (impl_generics, type_generics, where_clause) = tweaked_generics.split_for_impl();
 
 		let super_trait = if self.subscriptions.is_empty() {
 			quote! { #jsonrpsee::types::traits::Client }
@@ -54,12 +58,13 @@ impl RpcDescription {
 		let trait_impl = quote! {
 			#[#async_trait]
 			#[doc = #doc_comment]
-			pub trait #trait_name: #super_trait {
+			pub trait #trait_name #impl_generics: #super_trait where #where_clause {
 				#(#method_impls)*
 				#(#sub_impls)*
 			}
 
-			impl<T> #trait_name for T where T: #super_trait {}
+			// TODO(niklasad1): support for where clause on trait def.
+			impl<T #(,#type_idents)*> #trait_name #type_generics for T where T: #super_trait {}
 		};
 
 		Ok(trait_impl)

proc-macros/src/render_server.rs

@@ -26,13 +26,15 @@
 
 use super::lifetimes::replace_lifetimes;
 use super::RpcDescription;
+use crate::helpers::server_generate_where_clause;
 use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::{quote, quote_spanned};
 use std::collections::HashSet;
 
 impl RpcDescription {
 	pub(super) fn render_server(&self) -> Result<TokenStream2, syn::Error> {
 		let trait_name = quote::format_ident!("{}Server", &self.trait_def.ident);
+		let generics = self.trait_def.generics.clone();
 
 		let method_impls = self.render_methods()?;
 		let into_rpc_impl = self.render_into_rpc()?;
@@ -45,7 +47,7 @@ impl RpcDescription {
 		let trait_impl = quote! {
 			#[#async_trait]
 			#[doc = #doc_comment]
-			pub trait #trait_name: Sized + Send + Sync + 'static {
+			pub trait #trait_name #generics: Sized + Send + Sync + 'static {
 				#method_impls
 				#into_rpc_impl
 			}
@@ -163,9 +165,13 @@ impl RpcDescription {
 		let doc_comment = "Collects all the methods and subscriptions defined in the trait \
 								and adds them into a single `RpcModule`.";
 
+		let sub_tys: Vec<syn::Type> = self.subscriptions.clone().into_iter().map(|s| s.item).collect();
+		let where_clause = server_generate_where_clause(&self.trait_def, &sub_tys);
+
+		// NOTE(niklasad1): empty where clause is valid rust syntax.
 		Ok(quote! {
 			#[doc = #doc_comment]
-			fn into_rpc(self) -> #rpc_module<Self> {
+			fn into_rpc(self) -> #rpc_module<Self> where #(#where_clause,)* {
 				let mut rpc = #rpc_module::new(self);
 
 				#(#errors)*
@@ -202,12 +208,11 @@ impl RpcDescription {
 			quote! {
 				let mut seq = params.sequence();
 				#(#decode_fields);*
-				(#params_fields)
 			}
 		};
 
 		// Code to decode sequence of parameters from a JSON object (aka map).
-		let decode_map = {
+		let _decode_map = {
 			let mut generics = None;
 
 			let serde = self.jrps_server_item(quote! { types::__reexports::serde });
@@ -246,11 +251,13 @@ impl RpcDescription {
 
 		// Parsing of `serde_json::Value`.
 		let parsing = quote! {
-			let (#params_fields) = if params.is_object() {
+			// TODO(niklasad1): add support for JSON object.
+			/*let (#params_fields) = if params.is_object() {
 				#decode_map
 			} else {
 				#decode_array
-			};
+			};*/
+			#decode_array;
 		};
 
 		(parsing, params_fields)

proc-macros/src/rpc_macro.rs

@@ -85,6 +85,7 @@ impl RpcSubscription {
 		let name = attributes.name;
 		let unsub_method = attributes.unsub;
 		let item = attributes.item;
+
 		let params: Vec<_> = sig
 			.inputs
 			.into_iter()
@@ -116,7 +117,7 @@ pub struct RpcDescription {
 	pub(crate) trait_def: syn::ItemTrait,
 	/// List of RPC methods defined in the trait.
 	pub(crate) methods: Vec<RpcMethod>,
-	/// List of RPC subscritpions defined in the trait.
+	/// List of RPC subscriptions defined in the trait.
 	pub(crate) subscriptions: Vec<RpcSubscription>,
 }