Skip to main content

iota_open_rpc_macros/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use derive_syn_parse::Parse;
6use itertools::Itertools;
7use proc_macro::TokenStream;
8use proc_macro2::{Ident, Span, TokenStream as TokenStream2, TokenTree};
9use quote::{ToTokens, TokenStreamExt, quote};
10use syn::{
11    Attribute, GenericArgument, LitStr, PatType, Path, PathArguments, Token, TraitItem, Type,
12    parse,
13    parse::{Parse, ParseStream},
14    parse_macro_input,
15    punctuated::Punctuated,
16    spanned::Spanned,
17    token::{Comma, Paren},
18};
19use unescape::unescape;
20
21const IOTA_RPC_ATTRS: [&str; 2] = ["deprecated", "version"];
22
23/// Add a [Service name]OpenRpc struct and implementation providing access to
24/// Open RPC doc builder. This proc macro must be use in conjunction with
25/// `jsonrpsee_proc_macro::rpc`
26///
27/// The generated method `open_rpc` is added to [Service name]OpenRpc,
28/// ideally we want to add this to the trait generated by jsonrpsee framework,
29/// creating a new struct to provide access to the method is a workaround.
30///
31/// TODO: consider contributing the open rpc doc macro to jsonrpsee to simplify
32/// the logics.
33#[proc_macro_attribute]
34pub fn open_rpc(attr: TokenStream, item: TokenStream) -> TokenStream {
35    let attr: OpenRpcAttributes = parse_macro_input!(attr);
36
37    let mut trait_data: syn::ItemTrait = syn::parse(item).unwrap();
38    let rpc_definition = parse_rpc_method(&mut trait_data).unwrap();
39
40    let namespace = attr
41        .find_attr("namespace")
42        .map(|str| str.value())
43        .unwrap_or_default();
44
45    let tag = attr.find_attr("tag").to_quote();
46
47    let methods = rpc_definition.methods.iter().flat_map(|method|{
48        if method.deprecated {
49            return None;
50        }
51        let name = &method.name;
52        let deprecated = method.deprecated;
53        let doc = &method.doc;
54        let mut inputs = Vec::new();
55        for (name, ty, description) in &method.params {
56            let (ty, required) = extract_type_from_option(ty.clone());
57            let description = if let Some(description) = description {
58                quote! {Some(#description.to_string())}
59            } else {
60                quote! {None}
61            };
62
63            inputs.push(quote! {
64                let des = builder.create_content_descriptor::<#ty>(#name, None, #description, #required);
65                inputs.push(des);
66            })
67        }
68        let returns_ty = if let Some(ty) = &method.returns {
69            let (ty, required) = extract_type_from_option(ty.clone());
70            let name = quote! {#ty}.to_string();
71            quote! {Some(builder.create_content_descriptor::<#ty>(#name, None, None, #required));}
72        } else {
73            quote! {None;}
74        };
75
76        if method.is_pubsub {
77            Some(quote! {
78                let mut inputs: Vec<iota_open_rpc::ContentDescriptor> = Vec::new();
79                #(#inputs)*
80                let result = #returns_ty
81                builder.add_subscription(#namespace, #name, inputs, result, #doc, #tag, #deprecated);
82            })
83        } else {
84            Some(quote! {
85                let mut inputs: Vec<iota_open_rpc::ContentDescriptor> = Vec::new();
86                #(#inputs)*
87                let result = #returns_ty
88                builder.add_method(#namespace, #name, inputs, result, #doc, #tag, #deprecated);
89            })
90        }
91    }).collect::<Vec<_>>();
92
93    let routes = rpc_definition
94        .version_routing
95        .into_iter()
96        .map(|route| {
97            let name = route.name;
98            let route_to = route.route_to;
99            let comparator = route.token.to_string();
100            let version = route.version;
101            quote! {
102                builder.add_method_routing(#namespace, #name, #route_to, #comparator, #version);
103            }
104        })
105        .collect::<Vec<_>>();
106
107    let open_rpc_name = quote::format_ident!("{}OpenRpc", &rpc_definition.name);
108
109    quote! {
110        #trait_data
111        pub struct #open_rpc_name;
112        impl #open_rpc_name {
113            pub fn module_doc() -> iota_open_rpc::Module{
114                let mut builder = iota_open_rpc::RpcModuleDocBuilder::default();
115                #(#methods)*
116                #(#routes)*
117                builder.build()
118            }
119        }
120    }
121    .into()
122}
123
124trait OptionalQuote {
125    fn to_quote(&self) -> TokenStream2;
126}
127
128impl OptionalQuote for Option<LitStr> {
129    fn to_quote(&self) -> TokenStream2 {
130        if let Some(value) = self {
131            quote!(Some(#value.to_string()))
132        } else {
133            quote!(None)
134        }
135    }
136}
137
138struct RpcDefinition {
139    name: Ident,
140    methods: Vec<Method>,
141    version_routing: Vec<Routing>,
142}
143struct Method {
144    name: String,
145    params: Vec<(String, Type, Option<String>)>,
146    returns: Option<Type>,
147    doc: String,
148    is_pubsub: bool,
149    deprecated: bool,
150}
151struct Routing {
152    name: String,
153    route_to: String,
154    token: TokenStream2,
155    version: String,
156}
157
158fn parse_rpc_method(trait_data: &mut syn::ItemTrait) -> Result<RpcDefinition, syn::Error> {
159    let mut methods = Vec::new();
160    let mut version_routing = Vec::new();
161    for trait_item in &mut trait_data.items {
162        if let TraitItem::Method(method) = trait_item {
163            let doc = extract_doc_comments(&method.attrs).to_string();
164            let params: Vec<_> = method
165                .sig
166                .inputs
167                .iter_mut()
168                .filter_map(|arg| {
169                    match arg {
170                        syn::FnArg::Receiver(_) => None,
171                        syn::FnArg::Typed(arg) => {
172                            let description = if let Some(description) = arg.attrs.iter().position(|a|a.path.is_ident("doc")){
173                                let doc = extract_doc_comments(&arg.attrs);
174                                arg.attrs.remove(description);
175                                Some(doc)
176                            }else{
177                                None
178                            };
179                            match *arg.pat.clone() {
180                                syn::Pat::Ident(name) => {
181                                    Some(get_type(arg).map(|ty| (name.ident.to_string(), ty, description)))
182                                }
183                                syn::Pat::Wild(wild) => Some(Err(syn::Error::new(
184                                    wild.underscore_token.span(),
185                                    "Method argument names must be valid Rust identifiers; got `_` instead",
186                                ))),
187                                _ => Some(Err(syn::Error::new(
188                                    arg.span(),
189                                    format!("Unexpected method signature input; got {:?} ", *arg.pat),
190                                ))),
191                            }
192                        },
193                    }
194                })
195                .collect::<Result<_, _>>()?;
196
197            let (method_name, mut returns, is_pubsub, deprecated) = if let Some(attr) =
198                find_attr(&mut method.attrs, "method")
199            {
200                let token: TokenStream = attr.tokens.clone().into();
201                let returns = match &method.sig.output {
202                    syn::ReturnType::Default => None,
203                    syn::ReturnType::Type(_, output) => extract_type_from(output, "RpcResult"),
204                };
205                let mut attributes = parse::<Attributes>(token)?;
206                let method_name = attributes.get_value("name");
207
208                let deprecated = attributes.find("deprecated").is_some();
209
210                if let Some(version_attr) = attributes.find("version") {
211                    if let (Some(token), Some(version)) = (&version_attr.token, &version_attr.value)
212                    {
213                        let route_to =
214                            format!("{method_name}_{}", version.value().replace('.', "_"));
215                        version_routing.push(Routing {
216                            name: method_name,
217                            route_to: route_to.clone(),
218                            token: token.to_token_stream(),
219                            version: version.value(),
220                        });
221                        if let Some(name) = attributes.find_mut("name") {
222                            name.value
223                                .replace(LitStr::new(&route_to, Span::call_site()));
224                        }
225                        attr.tokens = remove_iota_rpc_attributes(attributes);
226                        continue;
227                    }
228                }
229                attr.tokens = remove_iota_rpc_attributes(attributes);
230                (method_name, returns, false, deprecated)
231            } else if let Some(attr) = find_attr(&mut method.attrs, "subscription") {
232                let token: TokenStream = attr.tokens.clone().into();
233                let attributes = parse::<Attributes>(token)?;
234                let name = attributes.get_value("name");
235                let type_ = attributes
236                    .find("item")
237                    .expect("Subscription should have a [item] attribute")
238                    .type_
239                    .clone()
240                    .expect("[item] attribute should have a value");
241                let deprecated = attributes.find("deprecated").is_some();
242                attr.tokens = remove_iota_rpc_attributes(attributes);
243                (name, Some(type_), true, deprecated)
244            } else {
245                continue;
246            };
247
248            // Allow overriding the return type via #[schemars(with = "X")] on
249            // the method, matching the parameter-level override in `get_type`.
250            if let Some((pos, schemars_attr)) = method
251                .attrs
252                .iter()
253                .find_position(|a| a.path.is_ident("schemars"))
254            {
255                let attribute = parse::<NamedAttribute>(schemars_attr.tokens.clone().into())?;
256                let stream: TokenStream2 = syn::parse_str(&attribute.value.value())?;
257                let tokens = respan_token_stream(stream, attribute.value.span());
258                let ty: Type = syn::parse2(tokens)?;
259                method.attrs.remove(pos);
260                returns = Some(ty);
261            }
262
263            methods.push(Method {
264                name: method_name,
265                params,
266                returns,
267                doc,
268                is_pubsub,
269                deprecated,
270            });
271        }
272    }
273    Ok(RpcDefinition {
274        name: trait_data.ident.clone(),
275        methods,
276        version_routing,
277    })
278}
279// Remove IOTA rpc specific attributes.
280fn remove_iota_rpc_attributes(attributes: Attributes) -> TokenStream2 {
281    let attrs = attributes
282        .attrs
283        .into_iter()
284        .filter(|r| !IOTA_RPC_ATTRS.contains(&r.key.to_string().as_str()))
285        .collect::<Punctuated<Attr, Comma>>();
286    quote! {(#attrs)}
287}
288
289fn extract_type_from(ty: &Type, from_ty: &str) -> Option<Type> {
290    fn path_is(path: &Path, from_ty: &str) -> bool {
291        path.leading_colon.is_none()
292            && path.segments.len() == 1
293            && path.segments.iter().next().unwrap().ident == from_ty
294    }
295
296    if let Type::Path(p) = ty {
297        if p.qself.is_none() && path_is(&p.path, from_ty) {
298            if let PathArguments::AngleBracketed(a) = &p.path.segments[0].arguments {
299                if let Some(GenericArgument::Type(ty)) = a.args.first() {
300                    return Some(ty.clone());
301                }
302            }
303        }
304    }
305    None
306}
307
308fn extract_type_from_option(ty: Type) -> (Type, bool) {
309    if let Some(ty) = extract_type_from(&ty, "Option") {
310        (ty, false)
311    } else {
312        (ty, true)
313    }
314}
315
316fn get_type(pat_type: &mut PatType) -> Result<Type, syn::Error> {
317    Ok(
318        if let Some((pos, attr)) = pat_type
319            .attrs
320            .iter()
321            .find_position(|a| a.path.is_ident("schemars"))
322        {
323            let attribute = parse::<NamedAttribute>(attr.tokens.clone().into())?;
324
325            let stream = syn::parse_str(&attribute.value.value())?;
326            let tokens = respan_token_stream(stream, attribute.value.span());
327
328            let path = syn::parse2(tokens)?;
329            pat_type.attrs.remove(pos);
330            path
331        } else {
332            pat_type.ty.as_ref().clone()
333        },
334    )
335}
336
337fn find_attr<'a>(attrs: &'a mut [Attribute], ident: &str) -> Option<&'a mut Attribute> {
338    attrs.iter_mut().find(|a| a.path.is_ident(ident))
339}
340
341fn respan_token_stream(stream: TokenStream2, span: Span) -> TokenStream2 {
342    stream
343        .into_iter()
344        .map(|mut token| {
345            if let TokenTree::Group(g) = &mut token {
346                *g = proc_macro2::Group::new(g.delimiter(), respan_token_stream(g.stream(), span));
347            }
348            token.set_span(span);
349            token
350        })
351        .collect()
352}
353
354/// Find doc comments by looking for #[doc = "..."] attributes.
355///
356/// Consecutive attributes are combined together. If there is a leading space,
357/// it will be removed, and if there is trailing whitespace it will also be
358/// removed. Single newlines in doc comments are replaced by spaces (soft
359/// wrapping), but double newlines (an empty line) are preserved.
360fn extract_doc_comments(attrs: &[Attribute]) -> String {
361    let mut s = String::new();
362    let mut sep = "";
363
364    for attr in attrs {
365        if !attr.path.is_ident("doc") {
366            continue;
367        }
368
369        let Ok(syn::Meta::NameValue(meta)) = attr.parse_meta() else {
370            continue;
371        };
372
373        let syn::Lit::Str(lit) = &meta.lit else {
374            continue;
375        };
376
377        let token = lit.value();
378        let line = token.strip_prefix(" ").unwrap_or(&token).trim_end();
379
380        if line.is_empty() {
381            s.push_str("\n\n");
382            sep = "";
383        } else {
384            s.push_str(sep);
385            sep = " ";
386        }
387
388        s.push_str(line);
389    }
390
391    unescape(&s).unwrap_or_else(|| panic!("Cannot unescape doc comments : [{s}]"))
392}
393
394#[derive(Parse, Debug)]
395struct OpenRpcAttributes {
396    #[parse_terminated(OpenRpcAttribute::parse)]
397    fields: Punctuated<OpenRpcAttribute, Token![,]>,
398}
399
400impl OpenRpcAttributes {
401    fn find_attr(&self, name: &str) -> Option<LitStr> {
402        self.fields
403            .iter()
404            .find(|attr| attr.label == name)
405            .map(|attr| attr.value.clone())
406    }
407}
408
409#[derive(Parse, Debug)]
410struct OpenRpcAttribute {
411    label: Ident,
412    _eq_token: Token![=],
413    value: syn::LitStr,
414}
415
416#[derive(Parse, Debug)]
417struct NamedAttribute {
418    #[paren]
419    _paren_token: Paren,
420    #[inside(_paren_token)]
421    _ident: Ident,
422    #[inside(_paren_token)]
423    _eq_token: Token![=],
424    #[inside(_paren_token)]
425    value: syn::LitStr,
426}
427
428#[derive(Debug)]
429struct Attributes {
430    pub attrs: Punctuated<Attr, syn::token::Comma>,
431}
432
433impl Attributes {
434    pub fn find(&self, attr_name: &str) -> Option<&Attr> {
435        self.attrs.iter().find(|attr| attr.key == attr_name)
436    }
437    pub fn find_mut(&mut self, attr_name: &str) -> Option<&mut Attr> {
438        self.attrs.iter_mut().find(|attr| attr.key == attr_name)
439    }
440    pub fn get_value(&self, attr_name: &str) -> String {
441        self.attrs
442            .iter()
443            .find(|attr| attr.key == attr_name)
444            .unwrap_or_else(|| panic!("Method should have a [{attr_name}] attribute."))
445            .value
446            .as_ref()
447            .unwrap_or_else(|| panic!("[{attr_name}] attribute should have a value"))
448            .value()
449    }
450}
451
452impl Parse for Attributes {
453    fn parse(input: ParseStream) -> syn::Result<Self> {
454        let content;
455        let _paren = syn::parenthesized!(content in input);
456        let attrs = content.parse_terminated(Attr::parse)?;
457        Ok(Self { attrs })
458    }
459}
460
461#[derive(Debug)]
462struct Attr {
463    pub key: Ident,
464    pub token: Option<TokenStream2>,
465    pub value: Option<syn::LitStr>,
466    pub type_: Option<Type>,
467}
468
469impl ToTokens for Attr {
470    fn to_tokens(&self, tokens: &mut TokenStream2) {
471        tokens.append(self.key.clone());
472        if let Some(token) = &self.token {
473            tokens.extend(token.to_token_stream());
474        }
475        if let Some(value) = &self.value {
476            tokens.append(value.token());
477        }
478        if let Some(type_) = &self.type_ {
479            tokens.extend(type_.to_token_stream());
480        }
481    }
482}
483
484impl Parse for Attr {
485    fn parse(input: ParseStream) -> syn::Result<Self> {
486        let key = input.parse()?;
487        let token = if input.peek(Token!(=)) {
488            Some(input.parse::<Token!(=)>()?.to_token_stream())
489        } else if input.peek(Token!(<=)) {
490            Some(input.parse::<Token!(<=)>()?.to_token_stream())
491        } else {
492            None
493        };
494
495        let value = if token.is_some() && input.peek(syn::LitStr) {
496            Some(input.parse::<syn::LitStr>()?)
497        } else {
498            None
499        };
500
501        let type_ = if token.is_some() && input.peek(syn::Ident) {
502            Some(input.parse::<Type>()?)
503        } else {
504            None
505        };
506
507        Ok(Self {
508            key,
509            token,
510            value,
511            type_,
512        })
513    }
514}