Skip to main content

jsonrpsee_proc_macros/
visitor.rs

1// Copyright 2021 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any
4// person obtaining a copy of this software and associated
5// documentation files (the "Software"), to deal in the
6// Software without restriction, including without
7// limitation the rights to use, copy, modify, merge,
8// publish, distribute, sublicense, and/or sell copies of
9// the Software, and to permit persons to whom the Software
10// is furnished to do so, subject to the following
11// conditions:
12//
13// The above copyright notice and this permission notice
14// shall be included in all copies or substantial portions
15// of the Software.
16//
17// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
18// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
19// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
20// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
21// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
22// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
23// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
24// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
25// DEALINGS IN THE SOFTWARE.
26
27use std::collections::HashSet;
28
29use syn::{
30    visit::{self, Visit},
31    Ident,
32};
33
34/// Visitor that parses generic type parameters from `syn::Type` by traversing the AST.
35/// A `syn::Type` can any type such as `Vec<T>, T, Foo<A<B<V>>>, usize or similar`.
36/// The implementation is based on <https://github.com/serde-rs/serde/blob/master/serde_derive/src/bound.rs>.
37#[derive(Default, Debug)]
38pub(crate) struct FindSubscriptionParams {
39    pub(crate) generic_sub_params: HashSet<Ident>,
40    pub(crate) all_type_params: HashSet<Ident>,
41}
42
43/// Visitor for the entire `RPC trait`.
44pub struct FindAllParams {
45    pub(crate) trait_generics: HashSet<syn::Ident>,
46    pub(crate) input_params: HashSet<syn::Ident>,
47    pub(crate) ret_params: HashSet<syn::Ident>,
48    pub(crate) sub_params: HashSet<syn::Ident>,
49    pub(crate) visiting_return_type: bool,
50    pub(crate) visiting_fn_arg: bool,
51}
52
53impl FindAllParams {
54    /// Create a visitor to traverse the entire RPC trait.
55    /// It takes the already visited subscription parameters as input.
56    pub fn new(sub_params: HashSet<syn::Ident>) -> Self {
57        Self {
58            trait_generics: HashSet::new(),
59            input_params: HashSet::new(),
60            ret_params: HashSet::new(),
61            sub_params,
62            visiting_return_type: false,
63            visiting_fn_arg: false,
64        }
65    }
66}
67
68impl<'ast> Visit<'ast> for FindAllParams {
69    /// Visit generic type param.
70    fn visit_type_param(&mut self, ty_param: &'ast syn::TypeParam) {
71        self.trait_generics.insert(ty_param.ident.clone());
72    }
73
74    /// Visit return type and mark it as `visiting_return_type`.
75    /// To know whether a given Ident is a function argument or return type when traversing.
76    fn visit_return_type(&mut self, return_type: &'ast syn::ReturnType) {
77        self.visiting_return_type = true;
78        visit::visit_return_type(self, return_type);
79        self.visiting_return_type = false
80    }
81
82    /// Visit ident.
83    fn visit_ident(&mut self, ident: &'ast syn::Ident) {
84        if self.trait_generics.contains(ident) {
85            if self.visiting_return_type {
86                self.ret_params.insert(ident.clone());
87            }
88            if self.visiting_fn_arg {
89                self.input_params.insert(ident.clone());
90            }
91        }
92    }
93
94    /// Visit function argument and mark it as `visiting_fn_arg`.
95    /// To know whether a given Ident is a function argument or return type when traversing.
96    fn visit_fn_arg(&mut self, arg: &'ast syn::FnArg) {
97        self.visiting_fn_arg = true;
98        visit::visit_fn_arg(self, arg);
99        self.visiting_fn_arg = false;
100    }
101}
102
103impl FindSubscriptionParams {
104    /// Visit all types and returns all generic [`struct@syn::Ident`]'s that are subscriptions.
105    pub fn visit(mut self, tys: &[syn::Type]) -> HashSet<Ident> {
106        for ty in tys {
107            self.visit_type(ty);
108        }
109        self.generic_sub_params
110    }
111
112    /// Create a new subscription parameters visitor that takes all
113    /// generic parameters on the RPC trait as input in order to determine
114    /// whether a given ident is a generic type param or not when traversing
115    /// one or more types in `FindSubscriptionParams::visit`.
116    pub fn new(all_type_params: HashSet<Ident>) -> Self {
117        Self { generic_sub_params: HashSet::new(), all_type_params }
118    }
119
120    /// Visit path, if it's a leaf path and generic type param then add it as a subscription param.
121    fn visit_path(&mut self, path: &syn::Path) {
122        if path.leading_colon.is_none() && path.segments.len() == 1 {
123            let id = &path.segments[0].ident;
124            if self.all_type_params.contains(id) {
125                self.generic_sub_params.insert(id.clone());
126            }
127        }
128        for segment in &path.segments {
129            self.visit_path_segment(segment);
130        }
131    }
132
133    /// Traverse syntax tree.
134    fn visit_type(&mut self, ty: &syn::Type) {
135        match ty {
136            syn::Type::Array(ty) => self.visit_type(&ty.elem),
137            syn::Type::BareFn(ty) => {
138                for arg in &ty.inputs {
139                    self.visit_type(&arg.ty);
140                }
141                self.visit_return_type(&ty.output);
142            }
143            syn::Type::Group(ty) => self.visit_type(&ty.elem),
144            syn::Type::ImplTrait(ty) => {
145                for bound in &ty.bounds {
146                    self.visit_type_param_bound(bound);
147                }
148            }
149            syn::Type::Macro(ty) => self.visit_macro(&ty.mac),
150            syn::Type::Paren(ty) => self.visit_type(&ty.elem),
151            syn::Type::Path(ty) => {
152                if let Some(qself) = &ty.qself {
153                    self.visit_type(&qself.ty);
154                }
155                self.visit_path(&ty.path);
156            }
157            syn::Type::Ptr(ty) => self.visit_type(&ty.elem),
158            syn::Type::Reference(ty) => self.visit_type(&ty.elem),
159            syn::Type::Slice(ty) => self.visit_type(&ty.elem),
160            syn::Type::TraitObject(ty) => {
161                for bound in &ty.bounds {
162                    self.visit_type_param_bound(bound);
163                }
164            }
165            syn::Type::Tuple(ty) => {
166                for elem in &ty.elems {
167                    self.visit_type(elem);
168                }
169            }
170            syn::Type::Infer(_) | syn::Type::Never(_) | syn::Type::Verbatim(_) => {}
171            _ => {}
172        }
173    }
174
175    /// Traverse syntax tree.
176    fn visit_path_segment(&mut self, segment: &syn::PathSegment) {
177        self.visit_path_arguments(&segment.arguments);
178    }
179
180    /// Traverse syntax tree.
181    fn visit_path_arguments(&mut self, arguments: &syn::PathArguments) {
182        match arguments {
183            syn::PathArguments::None => {}
184            syn::PathArguments::AngleBracketed(arguments) => {
185                for arg in &arguments.args {
186                    match arg {
187                        syn::GenericArgument::Type(arg) => self.visit_type(arg),
188                        syn::GenericArgument::AssocType(arg) => self.visit_type(&arg.ty),
189                        _ => {}
190                    }
191                }
192            }
193            syn::PathArguments::Parenthesized(arguments) => {
194                for argument in &arguments.inputs {
195                    self.visit_type(argument);
196                }
197                self.visit_return_type(&arguments.output);
198            }
199        }
200    }
201
202    /// Traverse syntax tree.
203    fn visit_return_type(&mut self, return_type: &syn::ReturnType) {
204        match return_type {
205            syn::ReturnType::Default => {}
206            syn::ReturnType::Type(_, output) => self.visit_type(output),
207        }
208    }
209
210    /// Traverse syntax tree.
211    fn visit_type_param_bound(&mut self, bound: &syn::TypeParamBound) {
212        if let syn::TypeParamBound::Trait(bound) = bound {
213            self.visit_path(&bound.path);
214        }
215    }
216
217    // Type parameter should not be considered used by a macro path.
218    //
219    //     struct TypeMacro<T> {
220    //         mac: T!(),
221    //         marker: PhantomData<T>,
222    //     }
223    fn visit_macro(&mut self, _mac: &syn::Macro) {}
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use syn::{parse_quote, Type};
230
231    #[test]
232    fn it_works() {
233        let t: Type = parse_quote!(Vec<T>);
234        let id: Ident = parse_quote!(T);
235
236        let mut exp = HashSet::new();
237        exp.insert(id);
238        let generics = exp.clone();
239
240        assert_eq!(exp, FindSubscriptionParams::new(generics).visit(&[t]));
241    }
242
243    #[test]
244    fn several_type_params() {
245        let t: Type = parse_quote!(Vec<(A, B, C)>);
246
247        let mut generics: HashSet<syn::Ident> = HashSet::new();
248        let mut exp = HashSet::new();
249
250        generics.insert(parse_quote!(A));
251        generics.insert(parse_quote!(B));
252        generics.insert(parse_quote!(C));
253        generics.insert(parse_quote!(D));
254
255        exp.insert(parse_quote!(A));
256        exp.insert(parse_quote!(B));
257        exp.insert(parse_quote!(C));
258
259        assert_eq!(exp, FindSubscriptionParams::new(generics).visit(&[t]));
260    }
261
262    #[test]
263    fn nested_type() {
264        let t: Type = parse_quote!(Vec<Foo<A, B>>);
265
266        let mut generics: HashSet<syn::Ident> = HashSet::new();
267        let mut exp = HashSet::new();
268
269        generics.insert(parse_quote!(A));
270        generics.insert(parse_quote!(B));
271        generics.insert(parse_quote!(C));
272        generics.insert(parse_quote!(D));
273
274        exp.insert(parse_quote!(A));
275        exp.insert(parse_quote!(B));
276
277        assert_eq!(exp, FindSubscriptionParams::new(generics).visit(&[t]));
278    }
279}