c2rust_macros/
lib.rs

1#![recursion_limit = "128"]
2extern crate proc_macro;
3extern crate proc_macro2;
4extern crate quote;
5extern crate syn;
6
7use proc_macro2::{Span, TokenStream};
8use quote::quote;
9use syn::parse_macro_input;
10use syn::visit::Visit;
11use syn::{Block, FnArg, Ident, Pat, TraitItemMethod, Type, TypeReference};
12
13#[derive(Default)]
14struct VisitorImpls {
15    tokens: TokenStream,
16    count: usize,
17}
18
19impl VisitorImpls {
20    fn generate_visit(
21        &mut self,
22        method_name: &Ident,
23        arg_pat: &Pat,
24        ty: &Type,
25        walk: &Block,
26    ) {
27        self.tokens.extend(quote! {
28            impl MutVisit for #ty {
29                fn visit<F: MutVisitor>(&mut self, f: &mut F) {
30                    f.#method_name(self)
31                }
32            }
33        });
34
35        let folder_name = format!("Folder{}", self.count);
36        let folder_ident = Ident::new(&folder_name, Span::call_site());
37
38        if !walk.stmts.is_empty() {
39            let noop_fn_name = format!("noop_{}", method_name);
40            let noop_fn = Ident::new(&noop_fn_name, Span::call_site());
41            self.tokens.extend(quote! {
42                impl WalkAst for #ty {
43                    fn walk<T: MutVisitor>(&mut self, v: &mut T) {
44                        syntax::mut_visit::#noop_fn(self, v);
45                    }
46                }
47            });
48        }
49        self.tokens.extend(quote! {
50            struct #folder_ident<F>
51                where F: FnMut(&mut #ty)
52            {
53                callback: F,
54            }
55
56            impl<F> MutVisitor for #folder_ident<F>
57                where F: FnMut(&mut #ty)
58            {
59                fn #method_name(&mut self, #arg_pat: &mut #ty) {
60                    #walk
61                    (self.callback)(#arg_pat)
62                }
63            }
64
65            impl MutVisitNodes for #ty {
66                fn visit<T, F>(target: &mut T, callback: F)
67                    where T: MutVisit,
68                          F: FnMut(&mut Self)
69                {
70                    let mut f = #folder_ident { callback };
71                    target.visit(&mut f)
72                }
73            }
74        });
75
76        self.count += 1;
77    }
78
79    fn generate_flat_map(
80        &mut self,
81        method_name: &Ident,
82        arg_pat: &Pat,
83        ty: &Type,
84        walk: &Block,
85    ) {
86        self.tokens.extend(quote! {
87            impl MutVisit for #ty {
88                fn visit<F: MutVisitor>(&mut self, f: &mut F) {
89                    let new = f.#method_name(self.clone());
90                    *self = new.lone();
91                }
92
93                fn flat_map<F: MutVisitor>(self, f: &mut F) -> SmallVec<[#ty; 1]> {
94                    f.#method_name(self)
95                }
96            }
97        });
98
99        let folder_name = format!("Folder{}", self.count);
100        let folder_ident = Ident::new(&folder_name, Span::call_site());
101
102        if !walk.stmts.is_empty() {
103            let noop_fn_name = format!("noop_{}", method_name);
104            let noop_fn = Ident::new(&noop_fn_name, Span::call_site());
105            self.tokens.extend(quote! {
106                impl WalkAst for #ty {
107                    fn walk<T: MutVisitor>(&mut self, v: &mut T) {
108                        *self = syntax::mut_visit::#noop_fn(self.clone(), v).lone();
109                    }
110                }
111            })
112        }
113        self.tokens.extend(quote! {
114            struct #folder_ident<F>
115                where F: FnMut(#ty) -> SmallVec<[#ty; 1]>
116            {
117                callback: F,
118            }
119
120            impl<F> MutVisitor for #folder_ident<F>
121                where F: FnMut(#ty) -> SmallVec<[#ty; 1]>
122            {
123                fn #method_name(&mut self, #arg_pat: #ty) -> SmallVec<[#ty; 1]> {
124                    let mut v = #walk;
125                    v.flat_map_in_place(|x| (self.callback)(x));
126                    v
127                }
128            }
129
130            impl FlatMapNodes for #ty {
131                fn visit<T, F>(target: &mut T, callback: F)
132                    where T: MutVisit,
133                          F: FnMut(#ty) -> SmallVec<[#ty; 1]>
134                {
135                    let mut f = #folder_ident { callback };
136                    target.visit(&mut f)
137                }
138
139                fn flat_map<T, F>(target: T, callback: F) -> SmallVec<[T; 1]>
140                    where T: MutVisit,
141                          F: FnMut(#ty) -> SmallVec<[#ty; 1]>
142                {
143                    let mut f = #folder_ident { callback };
144                    target.flat_map(&mut f)
145                }
146            }
147        });
148
149        self.count += 1;
150    }
151}
152
153impl<'ast> Visit<'ast> for VisitorImpls {
154    fn visit_trait_item_method(&mut self, m: &TraitItemMethod) {
155        let method_name = &m.sig.ident;
156        let method_noop = m.default.as_ref().unwrap();
157        match &m.sig.inputs[1] {
158            FnArg::Typed(pat_ty) => match &*pat_ty.ty {
159                Type::Reference(TypeReference {
160                    mutability: Some(_),
161                    elem,
162                    ..
163                }) => self.generate_visit(method_name, &pat_ty.pat, &elem, method_noop),
164
165                ty => self.generate_flat_map(method_name, &pat_ty.pat, &ty, method_noop),
166            },
167
168            _ => {}
169        }
170    }
171}
172
173#[proc_macro]
174pub fn gen_visitor_impls(tokens: proc_macro::TokenStream) -> proc_macro::TokenStream {
175    let visitor_trait: syn::ItemTrait = parse_macro_input!(tokens);
176    let mut visitor = VisitorImpls::default();
177    visitor.visit_item_trait(&visitor_trait);
178
179    visitor.tokens.into()
180}
c2rust_macros/lib.rs

c2rust_macros/
lib.rs
