1use crate::operand::{Borrowed, Operand, Owned};
2use crate::{file, full, gen};
3use anyhow::Result;
4use proc_macro2::{Ident, Span, TokenStream};
5use quote::{format_ident, quote};
6use syn::Index;
7use syn_codegen::{Data, Definitions, Features, Node, Type};
8
9const VISIT_MUT_SRC: &str = "src/gen/visit_mut.rs";
10
11fn simple_visit(item: &str, name: &Operand) -> TokenStream {
12    let ident = gen::under_name(item);
13    let method = format_ident!("visit_{}_mut", ident);
14    let name = name.ref_mut_tokens();
15    quote! {
16        v.#method(#name)
17    }
18}
19
20fn noop_visit(name: &Operand) -> TokenStream {
21    let name = name.tokens();
22    quote! {
23        skip!(#name)
24    }
25}
26
27fn visit(
28    ty: &Type,
29    features: &Features,
30    defs: &Definitions,
31    name: &Operand,
32) -> Option<TokenStream> {
33    match ty {
34        Type::Box(t) => {
35            let name = name.owned_tokens();
36            visit(t, features, defs, &Owned(quote!(*#name)))
37        }
38        Type::Vec(t) => {
39            let operand = Borrowed(quote!(it));
40            let val = visit(t, features, defs, &operand)?;
41            let name = name.ref_mut_tokens();
42            Some(quote! {
43                for it in #name {
44                    #val;
45                }
46            })
47        }
48        Type::Punctuated(p) => {
49            let operand = Borrowed(quote!(it));
50            let val = visit(&p.element, features, defs, &operand)?;
51            let name = name.ref_mut_tokens();
52            Some(quote! {
53                for mut el in Punctuated::pairs_mut(#name) {
54                    let it = el.value_mut();
55                    #val;
56                }
57            })
58        }
59        Type::Option(t) => {
60            let it = Borrowed(quote!(it));
61            let val = visit(t, features, defs, &it)?;
62            let name = name.ref_mut_tokens();
63            Some(quote! {
64                if let Some(it) = #name {
65                    #val;
66                }
67            })
68        }
69        Type::Tuple(t) => {
70            let mut code = TokenStream::new();
71            for (i, elem) in t.iter().enumerate() {
72                let name = name.tokens();
73                let i = Index::from(i);
74                let it = Owned(quote!((#name).#i));
75                let val = visit(elem, features, defs, &it).unwrap_or_else(|| noop_visit(&it));
76                code.extend(val);
77                code.extend(quote!(;));
78            }
79            Some(code)
80        }
81        Type::Syn(t) => {
82            fn requires_full(features: &Features) -> bool {
83                features.any.contains("full") && features.any.len() == 1
84            }
85            let mut res = simple_visit(t, name);
86            let target = defs.types.iter().find(|ty| ty.ident == *t).unwrap();
87            if requires_full(&target.features) && !requires_full(features) {
88                res = quote!(full!(#res));
89            }
90            Some(res)
91        }
92        Type::Ext(t) if gen::TERMINAL_TYPES.contains(&&t[..]) => Some(simple_visit(t, name)),
93        Type::Ext(_) | Type::Std(_) | Type::Token(_) | Type::Group(_) => None,
94    }
95}
96
97fn node(traits: &mut TokenStream, impls: &mut TokenStream, s: &Node, defs: &Definitions) {
98    let under_name = gen::under_name(&s.ident);
99    let ty = Ident::new(&s.ident, Span::call_site());
100    let visit_mut_fn = format_ident!("visit_{}_mut", under_name);
101
102    let mut visit_mut_impl = TokenStream::new();
103
104    match &s.data {
105        Data::Enum(variants) if variants.is_empty() => {
106            visit_mut_impl.extend(quote! {
107                match *node {}
108            });
109        }
110        Data::Enum(variants) => {
111            let mut visit_mut_variants = TokenStream::new();
112
113            for (variant, fields) in variants {
114                let variant_ident = Ident::new(variant, Span::call_site());
115
116                if fields.is_empty() {
117                    visit_mut_variants.extend(quote! {
118                        #ty::#variant_ident => {}
119                    });
120                } else {
121                    let mut bind_visit_mut_fields = TokenStream::new();
122                    let mut visit_mut_fields = TokenStream::new();
123
124                    for (idx, ty) in fields.iter().enumerate() {
125                        let binding = format_ident!("_binding_{}", idx);
126
127                        bind_visit_mut_fields.extend(quote! {
128                            #binding,
129                        });
130
131                        let borrowed_binding = Borrowed(quote!(#binding));
132
133                        visit_mut_fields.extend(
134                            visit(ty, &s.features, defs, &borrowed_binding)
135                                .unwrap_or_else(|| noop_visit(&borrowed_binding)),
136                        );
137
138                        visit_mut_fields.extend(quote!(;));
139                    }
140
141                    visit_mut_variants.extend(quote! {
142                        #ty::#variant_ident(#bind_visit_mut_fields) => {
143                            #visit_mut_fields
144                        }
145                    });
146                }
147            }
148
149            visit_mut_impl.extend(quote! {
150                match node {
151                    #visit_mut_variants
152                }
153            });
154        }
155        Data::Struct(fields) => {
156            for (field, ty) in fields {
157                let id = Ident::new(field, Span::call_site());
158                let ref_toks = Owned(quote!(node.#id));
159                let visit_mut_field = visit(ty, &s.features, defs, &ref_toks)
160                    .unwrap_or_else(|| noop_visit(&ref_toks));
161                visit_mut_impl.extend(quote! {
162                    #visit_mut_field;
163                });
164            }
165        }
166        Data::Private => {
167            if ty == "Ident" {
168                visit_mut_impl.extend(quote! {
169                    let mut span = node.span();
170                    v.visit_span_mut(&mut span);
171                    node.set_span(span);
172                });
173            }
174        }
175    }
176
177    traits.extend(quote! {
178        fn #visit_mut_fn(&mut self, i: &mut #ty) {
179            #visit_mut_fn(self, i);
180        }
181    });
182
183    impls.extend(quote! {
184        pub fn #visit_mut_fn<V>(v: &mut V, node: &mut #ty)
185        where
186            V: VisitMut + ?Sized,
187        {
188            #visit_mut_impl
189        }
190    });
191}
192
193pub fn generate(defs: &Definitions) -> Result<()> {
194    let (traits, impls) = gen::traverse(defs, node);
195    let full_macro = full::get_macro();
196    file::write(
197        VISIT_MUT_SRC,
198        quote! {
199            #![allow(unused_variables)]
200            #![allow(clippy::needless_pass_by_ref_mut)]
201
202            #[cfg(any(feature = "full", feature = "derive"))]
203            use crate::punctuated::Punctuated;
204            use crate::*;
205            use proc_macro2::Span;
206
207            #full_macro
208
209            macro_rules! skip {
210                ($($tt:tt)*) => {};
211            }
212
213            /// Syntax tree traversal to mutate an exclusive borrow of a syntax tree in
214            /// place.
215            ///
216            /// See the [module documentation] for details.
217            ///
218            /// [module documentation]: self
219            pub trait VisitMut {
220                #traits
221            }
222
223            #impls
224        },
225    )?;
226    Ok(())
227}
228