generate_code/generators/
visitor.rs

1use std::collections::HashSet;
2
3use inflector::Inflector;
4use proc_macro2::{Span, TokenStream};
5use quote::{quote, ToTokens};
6use swc_config::regex::CachedRegex;
7use syn::{
8    parse_quote, Arm, Attribute, Expr, Field, Fields, File, GenericArgument, Ident, Item, Lit,
9    LitInt, Path, PathArguments, Stmt, TraitItem, Type,
10};
11
12pub fn generate(crate_name_str: &str, node_types: &[&Item], excluded_types: &[String]) -> File {
13    let crate_name = Ident::new(crate_name_str, Span::call_site());
14    let crate_name = &crate_name;
15    let mut output = File {
16        shebang: None,
17        attrs: Vec::new(),
18        items: Vec::new(),
19    };
20    let mut all_types = all_field_types(node_types).into_iter().collect::<Vec<_>>();
21
22    if !excluded_types.is_empty() {
23        all_types.retain(|ty| {
24            !excluded_types
25                .iter()
26                .any(|type_name| ty.contains_type(type_name))
27        });
28    }
29
30    all_types.sort_by_cached_key(|v| v.method_name());
31
32    let mut typedefs = HashSet::new();
33
34    for node_type in node_types {
35        match node_type {
36            Item::Enum(data) => {
37                typedefs.insert(FieldType::Normal(data.ident.to_string()));
38            }
39            Item::Struct(data) => {
40                typedefs.insert(FieldType::Normal(data.ident.to_string()));
41            }
42            _ => {}
43        }
44    }
45
46    let field_only_types = {
47        let mut all = all_types.clone();
48
49        all.retain(|ty| !typedefs.contains(ty));
50
51        all
52    };
53
54    output.attrs.push(parse_quote!(
55        //! This file is generated by `tools/generate-code`. DO NOT MODIFY.
56    ));
57    output.attrs.push(parse_quote!(
58        #![allow(unused_variables)]
59    ));
60    output.attrs.push(parse_quote!(
61        #![allow(clippy::all)]
62    ));
63
64    output.items.push(parse_quote!(
65        use #crate_name::*;
66    ));
67
68    output.items.push(parse_quote!(
69        pub use ::swc_visit::All;
70    ));
71
72    for &kind in [TraitKind::Visit, TraitKind::VisitMut, TraitKind::Fold].iter() {
73        for &variant in [Variant::Normal, Variant::AstPath].iter() {
74            let g = Generator {
75                has_unknown: crate_name_str == "swc_ecma_ast",
76                kind,
77                variant,
78                excluded_types,
79            };
80
81            output.items.extend(g.declare_visit_trait(&all_types));
82
83            output.items.extend(g.declare_visit_with_trait());
84
85            output
86                .items
87                .extend(g.implement_visit_with_for_node_types(node_types));
88
89            output
90                .items
91                .extend(g.implement_visit_with_for_non_node_types(&field_only_types));
92
93            output
94                .items
95                .extend(g.implement_visit_with_for_generic_types());
96        }
97    }
98
99    output.items.push(parse_quote!(
100        #[cfg(any(docsrs, feature = "path"))]
101        pub type AstKindPath = swc_visit::AstKindPath<AstParentKind>;
102    ));
103    output.items.push(parse_quote!(
104        #[cfg(any(docsrs, feature = "path"))]
105        pub type AstNodePath<'ast> = swc_visit::AstNodePath<AstParentNodeRef<'ast>>;
106    ));
107    output.items.extend(define_fields(crate_name, node_types));
108
109    output
110}
111
112#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
113enum FieldType {
114    Normal(String),
115    Generic(String, Box<FieldType>),
116}
117
118impl ToTokens for FieldType {
119    fn to_tokens(&self, tokens: &mut TokenStream) {
120        match self {
121            FieldType::Normal(name) => {
122                let parsed: Path = syn::parse_str(name).expect("failed to parse path");
123                parsed.to_tokens(tokens);
124            }
125            FieldType::Generic(name, ty) => {
126                let name = Ident::new(name, Span::call_site());
127                let ty = &**ty;
128                quote!(#name<#ty>).to_tokens(tokens);
129            }
130        }
131    }
132}
133
134impl FieldType {
135    pub fn method_name(&self) -> String {
136        match self {
137            FieldType::Normal(name) => name.split("::").last().unwrap().to_snake_case(),
138            FieldType::Generic(name, ty) => match &**name {
139                "Option" => format!("opt_{}", ty.method_name()),
140                "Vec" => {
141                    // Vec<Option<Foo>> => opt_vec_foo
142                    match &**ty {
143                        FieldType::Generic(name, ty) if name == "Option" => {
144                            return format!("opt_vec_{}s", ty.method_name())
145                        }
146                        _ => {}
147                    }
148                    format!("{}s", ty.method_name())
149                }
150                "Box" => ty.method_name(),
151                _ => todo!("method_name for generic type: {}", name),
152            },
153        }
154    }
155
156    fn contains_type(&self, type_name: &str) -> bool {
157        let regex = CachedRegex::new(type_name).expect("failed to create regex");
158
159        match self {
160            FieldType::Normal(name) => regex.is_match(name),
161            FieldType::Generic(name, ty) => regex.is_match(name) || ty.contains_type(type_name),
162        }
163    }
164}
165
166fn all_field_types(node_types: &[&Item]) -> HashSet<FieldType> {
167    let mut all_types = HashSet::new();
168
169    for ty in node_types {
170        let type_name = match ty {
171            Item::Enum(data) => data.ident.to_string(),
172            Item::Struct(data) => data.ident.to_string(),
173            _ => continue,
174        };
175        all_types.insert(FieldType::Normal(type_name));
176
177        match ty {
178            Item::Enum(data) => {
179                for variant in &data.variants {
180                    for field in &variant.fields {
181                        let ty = &field.ty;
182                        all_types.extend(all_types_in_ty(ty));
183                    }
184                }
185            }
186            Item::Struct(data) => {
187                for field in &data.fields {
188                    let ty = &field.ty;
189                    all_types.extend(all_types_in_ty(ty));
190                }
191            }
192            _ => continue,
193        }
194    }
195
196    all_types
197}
198
199fn to_field_ty(ty: &Type) -> Option<FieldType> {
200    if let Some(ty) = extract_vec(ty) {
201        return to_field_ty(ty).map(|ty| FieldType::Generic("Vec".into(), Box::new(ty)));
202    }
203
204    if let Some(ty) = extract_generic("Box", ty) {
205        return to_field_ty(ty).map(|ty| FieldType::Generic("Box".into(), Box::new(ty)));
206    }
207
208    if let Some(ty) = extract_generic("Option", ty) {
209        return to_field_ty(ty).map(|ty| FieldType::Generic("Option".into(), Box::new(ty)));
210    }
211
212    match ty {
213        Type::Path(p) => {
214            let last = p.path.segments.last().unwrap();
215
216            if last.arguments.is_empty() {
217                let i = &last.ident;
218
219                if i == "bool"
220                    || i == "char"
221                    || i == "f32"
222                    || i == "f64"
223                    || i == "i8"
224                    || i == "i16"
225                    || i == "i32"
226                    || i == "i64"
227                    || i == "i128"
228                    || i == "isize"
229                    || i == "str"
230                    || i == "u8"
231                    || i == "u16"
232                    || i == "u32"
233                    || i == "u64"
234                    || i == "u128"
235                    || i == "usize"
236                {
237                    return None;
238                }
239
240                return Some(FieldType::Normal(quote!(#p).to_string()));
241            }
242
243            todo!("to_field_ty: {:?}", ty)
244        }
245        _ => todo!("to_field_ty"),
246    }
247}
248
249fn all_types_in_ty(ty: &Type) -> Vec<FieldType> {
250    if let Some(ty) = extract_vec(ty) {
251        let mut types = all_types_in_ty(ty);
252        types.extend(to_field_ty(ty).map(|ty| FieldType::Generic("Vec".into(), Box::new(ty))));
253        return types;
254    }
255
256    if let Some(ty) = extract_generic("Box", ty) {
257        let mut types = all_types_in_ty(ty);
258        types.extend(to_field_ty(ty).map(|ty| FieldType::Generic("Box".into(), Box::new(ty))));
259        return types;
260    }
261
262    if let Some(ty) = extract_generic("Option", ty) {
263        let mut types = all_types_in_ty(ty);
264        types.extend(to_field_ty(ty).map(|ty| FieldType::Generic("Option".into(), Box::new(ty))));
265        return types;
266    }
267
268    to_field_ty(ty).into_iter().collect()
269}
270
271#[derive(Debug, Clone, Copy, PartialEq, Eq)]
272enum TraitKind {
273    Visit,
274    VisitMut,
275    Fold,
276}
277
278impl TraitKind {
279    pub fn method_prefix(self) -> &'static str {
280        match self {
281            TraitKind::Visit => "visit",
282            TraitKind::VisitMut => "visit_mut",
283            TraitKind::Fold => "fold",
284        }
285    }
286
287    pub fn trait_prefix(self) -> &'static str {
288        match self {
289            TraitKind::Visit => "Visit",
290            TraitKind::VisitMut => "VisitMut",
291            TraitKind::Fold => "Fold",
292        }
293    }
294}
295
296#[derive(Debug, Clone, Copy, PartialEq, Eq)]
297enum Variant {
298    Normal,
299    AstPath,
300}
301
302impl Variant {
303    pub fn method_suffix(self, is_visitor_method: bool) -> &'static str {
304        if self == Variant::Normal || is_visitor_method {
305            ""
306        } else {
307            "_ast_path"
308        }
309    }
310}
311
312struct Generator<'a> {
313    has_unknown: bool,
314    kind: TraitKind,
315    variant: Variant,
316
317    excluded_types: &'a [String],
318}
319
320impl Generator<'_> {
321    fn should_skip(&self, ty: &Type) -> bool {
322        if let Some(ty) = extract_generic("Box", ty) {
323            return self.should_skip(ty);
324        }
325
326        if let Some(ty) = extract_generic("Vec", ty) {
327            return self.should_skip(ty);
328        }
329
330        if let Some(ty) = extract_generic("Option", ty) {
331            return self.should_skip(ty);
332        }
333
334        let ty = to_field_ty(ty);
335        match ty {
336            Some(ty) => {
337                for excluded_type in self.excluded_types {
338                    if ty.contains_type(excluded_type) {
339                        return true;
340                    }
341                }
342            }
343            None => return true,
344        }
345
346        false
347    }
348
349    fn method_lifetime(&self) -> TokenStream {
350        match self.kind {
351            TraitKind::Visit => match self.variant {
352                Variant::Normal => quote!(),
353                Variant::AstPath => quote!(<'ast: 'r, 'r>),
354            },
355            TraitKind::VisitMut => quote!(),
356            TraitKind::Fold => quote!(),
357        }
358    }
359
360    fn parameter_type_token(&self, ty: TokenStream) -> TokenStream {
361        match self.kind {
362            TraitKind::Visit => match self.variant {
363                Variant::Normal => quote!(&#ty),
364                Variant::AstPath => quote!(&'ast #ty),
365            },
366            TraitKind::VisitMut => quote!(&mut #ty),
367            TraitKind::Fold => ty,
368        }
369    }
370
371    /// This includes `->`
372    fn return_type_token(&self, ty: TokenStream) -> TokenStream {
373        match self.kind {
374            TraitKind::Visit => quote!(),
375            TraitKind::VisitMut => quote!(),
376            TraitKind::Fold => quote!(-> #ty),
377        }
378    }
379
380    fn arg_extra_token(&self) -> TokenStream {
381        match self.variant {
382            Variant::Normal => quote!(),
383            Variant::AstPath => quote!(, __ast_path),
384        }
385    }
386
387    fn param_extra_token(&self) -> TokenStream {
388        match self.variant {
389            Variant::Normal => quote!(),
390            Variant::AstPath => match self.kind {
391                TraitKind::Visit => {
392                    quote!(, __ast_path: &mut AstNodePath<'r>)
393                }
394                TraitKind::VisitMut | TraitKind::Fold => quote!(, __ast_path: &mut AstKindPath),
395            },
396        }
397    }
398
399    fn trait_name(&self, with: bool) -> Ident {
400        let name = self.kind.trait_prefix();
401
402        let name = if with {
403            format!("{name}With")
404        } else {
405            name.to_string()
406        };
407
408        match self.variant {
409            Variant::Normal => Ident::new(&name, Span::call_site()),
410            Variant::AstPath => Ident::new(&format!("{name}AstPath"), Span::call_site()),
411        }
412    }
413
414    fn base_trait_attrs(&self) -> Vec<Attribute> {
415        let mut attrs = Vec::new();
416
417        if self.variant == Variant::AstPath {
418            attrs.push(parse_quote!(#[cfg(any(docsrs, feature = "path"))]));
419            attrs.push(parse_quote!(#[cfg_attr(docsrs, doc(cfg(feature = "path")))]));
420        }
421
422        attrs
423    }
424
425    fn declare_visit_trait(&self, all_types: &[FieldType]) -> Vec<Item> {
426        let mut items = Vec::<Item>::new();
427        let lifetime = self.method_lifetime();
428        let ast_path_arg = self.arg_extra_token();
429        let ast_path_params = self.param_extra_token();
430        let with_trait_name = self.trait_name(true);
431        let trait_name = self.trait_name(false);
432        let attrs = self.base_trait_attrs();
433        let mut trait_methods = Vec::<TraitItem>::new();
434        let mut either_impl_methods = Vec::<TraitItem>::new();
435        let mut optional_impl_methods = Vec::<TraitItem>::new();
436        let mut ptr_impl_methods = Vec::<TraitItem>::new();
437
438        for ty in all_types {
439            if let FieldType::Generic(name, ..) = &ty {
440                if name == "Box" {
441                    continue;
442                }
443            }
444            let type_name = quote!(#ty);
445            let return_type = self.return_type_token(quote!(#type_name));
446            let node_type = self.node_type_for_visitor_method(ty);
447            let type_param = self.parameter_type_token(quote!(#node_type));
448
449            let visit_method_name = Ident::new(
450                &format!(
451                    "{}_{}{}",
452                    self.kind.method_prefix(),
453                    ty.method_name(),
454                    self.variant.method_suffix(true)
455                ),
456                Span::call_site(),
457            );
458            let visit_with_children_name = Ident::new(
459                &format!(
460                    "{}_children_with{}",
461                    self.kind.method_prefix(),
462                    self.variant.method_suffix(false)
463                ),
464                Span::call_site(),
465            );
466
467            let recurse_doc = "If you want to recurse, you need to call it manually.";
468
469            let method_doc = doc(&format!(
470                "Visit a node of type `{type_name}`.\n\nBy default, this method calls \
471                 [`{type_name}::{visit_with_children_name}`]. {recurse_doc}"
472            ));
473
474            trait_methods.push(parse_quote!(
475                #method_doc
476                #[inline]
477                fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
478                    <#node_type as #with_trait_name<Self>>::#visit_with_children_name(node, self #ast_path_arg)
479                }
480            ));
481
482            either_impl_methods.push(parse_quote!(
483                #[inline]
484                fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
485                    match self {
486                        swc_visit::Either::Left(visitor) => {
487                            #trait_name::#visit_method_name(visitor, node #ast_path_arg)
488                        }
489                        swc_visit::Either::Right(visitor) => {
490                            #trait_name::#visit_method_name(visitor, node #ast_path_arg)
491                        }
492                    }
493                }
494            ));
495
496            let else_block = if self.kind == TraitKind::Fold {
497                quote!(node)
498            } else {
499                quote!()
500            };
501            optional_impl_methods.push(parse_quote!(
502                #[inline]
503                fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
504                    if self.enabled {
505                        <V as #trait_name>::#visit_method_name(&mut self.visitor, node #ast_path_arg)
506                    } else {
507                        #else_block
508                    }
509                }
510            ));
511
512            ptr_impl_methods.push(parse_quote!(
513                #[inline]
514                fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
515                    <V as #trait_name>::#visit_method_name(&mut **self, node #ast_path_arg)
516                }
517            ));
518        }
519
520        items.push(parse_quote! {
521            /// A visitor trait for traversing the AST.
522            #(#attrs)*
523            pub trait #trait_name {
524                #(#trait_methods)*
525            }
526        });
527
528        // &mut V
529        items.push(parse_quote! {
530            #(#attrs)*
531            impl<V> #trait_name for &mut V where V: ?Sized + #trait_name {
532                #(#ptr_impl_methods)*
533            }
534        });
535
536        // Box<V>
537        items.push(parse_quote! {
538            #(#attrs)*
539            impl<V> #trait_name for Box<V> where V: ?Sized + #trait_name {
540                #(#ptr_impl_methods)*
541            }
542        });
543
544        // ::swc_visit::Either<A, B>
545
546        items.push(parse_quote! {
547            #(#attrs)*
548            impl<A, B> #trait_name for ::swc_visit::Either<A, B>
549            where
550                A: #trait_name,
551                B: #trait_name,
552            {
553                #(#either_impl_methods)*
554            }
555        });
556
557        // ::swc_visit::Optional<V>
558
559        items.push(parse_quote! {
560            #(#attrs)*
561            impl<V> #trait_name for ::swc_visit::Optional<V>
562            where
563                V: #trait_name,
564            {
565                #(#optional_impl_methods)*
566            }
567        });
568
569        items
570    }
571
572    fn declare_visit_with_trait(&self) -> Vec<Item> {
573        let visitor_trait_name = self.trait_name(false);
574        let trait_name = self.trait_name(true);
575        let attrs = self.base_trait_attrs();
576        let mut visit_with_trait_methods: Vec<TraitItem> = Vec::new();
577
578        {
579            let lifetime = self.method_lifetime();
580            let ast_path_extra = self.param_extra_token();
581            let return_type = self.return_type_token(quote!(Self));
582            let receiver = self.parameter_type_token(quote!(self));
583
584            let visit_with_name = Ident::new(
585                &format!(
586                    "{}_with{}",
587                    self.kind.method_prefix(),
588                    self.variant.method_suffix(false)
589                ),
590                Span::call_site(),
591            );
592            let visit_with_children_name = Ident::new(
593                &format!(
594                    "{}_children_with{}",
595                    self.kind.method_prefix(),
596                    self.variant.method_suffix(false)
597                ),
598                Span::call_site(),
599            );
600
601            visit_with_trait_methods.push(parse_quote!(
602                /// Calls a visitor method (visitor.fold_xxx) with self.
603                fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_extra) #return_type;
604            ));
605
606            visit_with_trait_methods.push(parse_quote!(
607            /// Visit children nodes of `self`` with `visitor`.
608            fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_extra) #return_type;
609        ));
610        }
611
612        let mut items: Vec<Item> = Vec::new();
613        items.push(parse_quote!(
614            /// A trait implemented for types that can be visited using a visitor.
615            #(#attrs)*
616            pub trait #trait_name<V: ?Sized + #visitor_trait_name> {
617                #(#visit_with_trait_methods)*
618            }
619        ));
620
621        items
622    }
623
624    fn implement_visit_with_for_node_types(&self, node_types: &[&Item]) -> Vec<Item> {
625        let visitor_trait_name = self.trait_name(false);
626        let trait_name = self.trait_name(true);
627        let attrs = self.base_trait_attrs();
628
629        let mut items: Vec<Item> = Vec::new();
630
631        for node_type in node_types {
632            let type_name = match node_type {
633                Item::Enum(data) => data.ident.clone(),
634                Item::Struct(data) => data.ident.clone(),
635                _ => continue,
636            };
637
638            let lifetime = self.method_lifetime();
639            let ast_path_arg = self.arg_extra_token();
640            let ast_path_param = self.param_extra_token();
641            let return_type = self.return_type_token(quote!(Self));
642
643            let receiver = self.parameter_type_token(quote!(self));
644            let visit_with_name = Ident::new(
645                &format!(
646                    "{}_with{}",
647                    self.kind.method_prefix(),
648                    self.variant.method_suffix(false)
649                ),
650                Span::call_site(),
651            );
652            let visit_with_children_name = Ident::new(
653                &format!(
654                    "{}_children_with{}",
655                    self.kind.method_prefix(),
656                    self.variant.method_suffix(false)
657                ),
658                Span::call_site(),
659            );
660
661            let visit_method_name = Ident::new(
662                &format!(
663                    "{}_{}{}",
664                    self.kind.method_prefix(),
665                    type_name.to_string().to_snake_case(),
666                    self.variant.method_suffix(true)
667                ),
668                Span::call_site(),
669            );
670
671            let visit_with_doc = doc(&format!(
672                "Calls [{visitor_trait_name}`::{visit_method_name}`] with `self`."
673            ));
674
675            let default_body: Expr = match node_type {
676                Item::Enum(data) => {
677                    let name = &data.ident;
678                    let mut match_arms = Vec::new();
679
680                    for v in &data.variants {
681                        let variant_name = &v.ident;
682
683                        match_arms.push(self.default_visit_body(
684                            quote!(#name::#variant_name),
685                            name,
686                            Some(variant_name),
687                            &v.fields,
688                        ));
689                    }
690
691                    match (self.has_unknown, self.kind) {
692                        (false, _) => parse_quote!(match self {
693                            #(#match_arms)*
694                        }),
695                        (true, TraitKind::Visit | TraitKind::VisitMut) => parse_quote!(match self {
696                            #(#match_arms)*
697                            #[cfg(swc_ast_unknown)]
698                            _ => ()
699                        }),
700                        (true, TraitKind::Fold) => parse_quote!(match self {
701                            #(#match_arms)*
702                            #[cfg(swc_ast_unknown)]
703                            _ => self
704                        }),
705                    }
706                }
707                Item::Struct(data) => {
708                    let name = &data.ident;
709
710                    let arm = self.default_visit_body(quote!(#name), name, None, &data.fields);
711
712                    parse_quote!(match self { #arm })
713                }
714                _ => continue,
715            };
716
717            items.push(parse_quote!(
718                #(#attrs)*
719                impl<V: ?Sized + #visitor_trait_name> #trait_name<V> for #type_name {
720                    #visit_with_doc
721                    fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
722                        <V as #visitor_trait_name>::#visit_method_name(visitor, self #ast_path_arg)
723                    }
724
725                    fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
726                        #default_body
727                    }
728                }
729            ));
730        }
731
732        items
733    }
734
735    fn default_visit_body(
736        &self,
737        path: TokenStream,
738        type_name: &Ident,
739        enum_variant_name: Option<&Ident>,
740        fields: &Fields,
741    ) -> Arm {
742        let ast_path_arg = match self.variant {
743            Variant::Normal => quote!(),
744            Variant::AstPath => quote!(, &mut *__ast_path),
745        };
746
747        let with_visitor_trait_name = self.trait_name(true);
748        let visit_with_name = Ident::new(
749            &format!(
750                "{}_with{}",
751                self.kind.method_prefix(),
752                self.variant.method_suffix(false)
753            ),
754            Span::call_site(),
755        );
756
757        let fields_enum_name = Ident::new(&format!("{type_name}Field"), Span::call_site());
758
759        let enum_ast_path = match enum_variant_name {
760            Some(variant_name) if self.variant == Variant::AstPath => {
761                let field_variant = Ident::new(
762                    &variant_name.to_string().to_pascal_case(),
763                    Span::call_site(),
764                );
765
766                match self.kind {
767                    TraitKind::Visit => Some(quote!(
768                        let mut __ast_path = __ast_path
769                            .with_guard(
770                                AstParentNodeRef::#type_name(self, self::fields::#fields_enum_name::#field_variant),
771                            );
772                    )),
773                    _ => Some(quote!(
774                        let mut __ast_path = __ast_path
775                            .with_guard(
776                                AstParentKind::#type_name(self::fields::#fields_enum_name::#field_variant),
777                            );
778                    )),
779                }
780            }
781            _ => None,
782        };
783
784        match fields {
785            Fields::Named(n) => {
786                let mut stmts: Vec<Stmt> = Vec::new();
787                let mut bindings = Vec::new();
788                let mut reconstruct = match self.kind {
789                    TraitKind::Visit | TraitKind::VisitMut => None,
790                    TraitKind::Fold => Some(Vec::<TokenStream>::new()),
791                };
792
793                for field in &n.named {
794                    let field_name = field.ident.as_ref().unwrap();
795                    let ty = &field.ty;
796
797                    bindings.push(field_name.clone());
798
799                    let field_variant =
800                        Ident::new(&field_name.to_string().to_pascal_case(), Span::call_site());
801
802                    let mut ast_path_guard_expr: Option<Stmt> = None;
803
804                    if self.variant == Variant::AstPath && !self.should_skip(ty) {
805                        let mut kind = quote!(self::fields::#fields_enum_name::#field_variant);
806
807                        if extract_vec(extract_generic("Option", ty).unwrap_or(ty)).is_some() {
808                            kind = quote!(#kind(usize::MAX));
809                        }
810
811                        match self.kind {
812                            TraitKind::Visit => {
813                                ast_path_guard_expr = Some(parse_quote!(
814                                    let mut __ast_path = __ast_path
815                                        .with_guard(
816                                            AstParentNodeRef::#type_name(self, #kind),
817                                        );
818                                ));
819                            }
820                            _ => {
821                                ast_path_guard_expr = Some(parse_quote!(
822                                    let mut __ast_path = __ast_path
823                                        .with_guard(
824                                            AstParentKind::#type_name(#kind),
825                                        );
826                                ));
827                            }
828                        }
829                    }
830
831                    if let Some(reconstructor) = &mut reconstruct {
832                        if !self.should_skip(ty) {
833                            stmts.push(parse_quote!(
834                                let #field_name = {
835                                    #ast_path_guard_expr
836                                    <#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg)
837                                };
838                            ));
839                        }
840
841                        reconstructor.push(parse_quote!(#field_name));
842                    } else if !self.should_skip(ty) {
843                        stmts.push(parse_quote!(
844                            {
845                                #ast_path_guard_expr
846                                <#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg)
847                            };
848                        ));
849                    }
850                }
851
852                match self.kind {
853                    TraitKind::Visit | TraitKind::VisitMut => {
854                        parse_quote!(#path { #(#bindings),* } => {
855                            #enum_ast_path;
856
857                            #(#stmts)*
858                        })
859                    }
860                    TraitKind::Fold => {
861                        let reconstruct = reconstruct.unwrap();
862
863                        parse_quote!(#path { #(#bindings),* } => {
864                            #enum_ast_path;
865
866
867                            #(#stmts)*
868
869                            #path {
870                                #(#reconstruct),*
871                            }
872                        })
873                    }
874                }
875            }
876            Fields::Unnamed(u) => {
877                let mut stmts: Vec<Stmt> = Vec::new();
878                let mut bindings = Vec::<TokenStream>::new();
879                let mut reconstruct = match self.kind {
880                    TraitKind::Visit | TraitKind::VisitMut => None,
881                    TraitKind::Fold => Some(Vec::<TokenStream>::new()),
882                };
883
884                for (idx, field) in u.unnamed.iter().enumerate() {
885                    let field_name = Ident::new(&format!("_field_{idx}"), Span::call_site());
886                    let ty = &field.ty;
887                    let binding_idx = Lit::Int(LitInt::new(&idx.to_string(), Span::call_site()));
888                    bindings.push(parse_quote!(#binding_idx: #field_name));
889
890                    if let Some(reconstructor) = &mut reconstruct {
891                        if !self.should_skip(ty) {
892                            stmts.push(parse_quote!(
893                                let #field_name = <#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg);
894                            ));
895                        }
896
897                        reconstructor.push(parse_quote!(#binding_idx: #field_name));
898                    } else if !self.should_skip(ty) {
899                        stmts.push(parse_quote!(
900                            <#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg);
901                        ));
902                    }
903                }
904
905                match self.kind {
906                    TraitKind::Visit | TraitKind::VisitMut => {
907                        parse_quote!(#path { #(#bindings),* }=> {
908                            #enum_ast_path;
909
910                            #(#stmts)*
911                        })
912                    }
913                    TraitKind::Fold => {
914                        let reconstruct = reconstruct.unwrap();
915
916                        parse_quote!(#path { #(#bindings),* } => {
917                            #enum_ast_path;
918
919                            #(#stmts)*
920
921                            #path{#(#reconstruct),*}
922                        })
923                    }
924                }
925            }
926            Fields::Unit => match self.kind {
927                TraitKind::Visit | TraitKind::VisitMut => {
928                    parse_quote!(#path => {})
929                }
930                TraitKind::Fold => parse_quote!(#path => #path,),
931            },
932        }
933    }
934
935    fn implement_visit_with_for_non_node_types(&self, non_leaf_types: &[FieldType]) -> Vec<Item> {
936        let visitor_trait_name = self.trait_name(false);
937        let visit_with_trait_name = self.trait_name(true);
938        let attrs = self.base_trait_attrs();
939        let lifetime = self.method_lifetime();
940        let ast_path_arg = self.arg_extra_token();
941        let ast_path_param = self.param_extra_token();
942        let return_type = self.return_type_token(quote!(Self));
943
944        let receiver = self.parameter_type_token(quote!(self));
945
946        let mut items: Vec<Item> = Vec::new();
947
948        for node_type in non_leaf_types {
949            let visit_with_name = Ident::new(
950                &format!(
951                    "{}_with{}",
952                    self.kind.method_prefix(),
953                    self.variant.method_suffix(false)
954                ),
955                Span::call_site(),
956            );
957            let visit_with_children_name = Ident::new(
958                &format!(
959                    "{}_children_with{}",
960                    self.kind.method_prefix(),
961                    self.variant.method_suffix(false)
962                ),
963                Span::call_site(),
964            );
965
966            let visit_method_name = Ident::new(
967                &format!(
968                    "{}_{}{}",
969                    self.kind.method_prefix(),
970                    node_type.method_name(),
971                    self.variant.method_suffix(true)
972                ),
973                Span::call_site(),
974            );
975
976            let visit_with_doc = doc(&format!(
977                "Calls [{visitor_trait_name}`::{visit_method_name}`] with `self`. (Extra impl)"
978            ));
979
980            let default_body: Expr = match node_type {
981                FieldType::Normal(..) => match self.kind {
982                    TraitKind::Visit => {
983                        parse_quote!({})
984                    }
985                    TraitKind::VisitMut => {
986                        parse_quote!({})
987                    }
988                    TraitKind::Fold => {
989                        parse_quote!(self)
990                    }
991                },
992
993                FieldType::Generic(name, inner) => match &**name {
994                    "Vec" => {
995                        let inner = inner.as_ref();
996                        let inner_ty = quote!(#inner);
997
998                        match (self.kind, self.variant) {
999                            (TraitKind::Visit, Variant::Normal) => {
1000                                parse_quote!(self.iter().for_each(|item| {
1001                                    <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor #ast_path_arg)
1002                                }))
1003                            }
1004                            (TraitKind::Visit, Variant::AstPath) => {
1005                                parse_quote!(self.iter().enumerate().for_each(|(__idx, item)| {
1006                                    let mut __ast_path = __ast_path.with_index_guard(__idx);
1007                                    <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor, &mut *__ast_path)
1008                                }))
1009                            }
1010                            (TraitKind::VisitMut, Variant::Normal) => {
1011                                parse_quote!(
1012                                    self.iter_mut().for_each(|item| {
1013                                        <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor #ast_path_arg)
1014                                    })
1015                                )
1016                            }
1017                            (TraitKind::VisitMut, Variant::AstPath) => {
1018                                parse_quote!(
1019                                    self.iter_mut().enumerate().for_each(|(__idx, item)| {
1020                                        let mut __ast_path = __ast_path.with_index_guard(__idx);
1021                                        <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor, &mut *__ast_path)
1022                                    })
1023                                )
1024                            }
1025                            (TraitKind::Fold, Variant::Normal) => {
1026                                parse_quote!(
1027                                    swc_visit::util::move_map::MoveMap::move_map(self, |item| {
1028                                        <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor #ast_path_arg)
1029                                    })
1030                                )
1031                            }
1032                            (TraitKind::Fold, Variant::AstPath) => {
1033                                parse_quote!(
1034                                    self.into_iter().enumerate().map(|(__idx, item)| {
1035                                        let mut __ast_path = __ast_path.with_index_guard(__idx);
1036                                        <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor, &mut *__ast_path)
1037                                    }).collect()
1038                                )
1039                            }
1040                        }
1041                    }
1042                    "Option" => {
1043                        let inner = inner.as_ref();
1044                        let inner_ty = quote!(#inner);
1045
1046                        match self.kind {
1047                            TraitKind::Visit => {
1048                                parse_quote!(
1049                                    match self {
1050                                        Some(inner) => {
1051                                            <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
1052                                        }
1053                                        None => {}
1054                                    }
1055                                )
1056                            }
1057                            TraitKind::VisitMut => {
1058                                parse_quote!(
1059                                    match self {
1060                                        Some(inner) => {
1061                                            <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
1062                                        }
1063                                        None => {}
1064                                    }
1065                                )
1066                            }
1067                            TraitKind::Fold => {
1068                                parse_quote!(
1069                                    self.map(|inner| {
1070                                        <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
1071                                    })
1072                                )
1073                            }
1074                        }
1075                    }
1076                    "Box" => continue,
1077                    _ => unreachable!("unexpected generic type: {}", name),
1078                },
1079            };
1080
1081            let target_type = self.node_type_for_visitor_method(node_type);
1082
1083            items.push(parse_quote!(
1084                #(#attrs)*
1085                impl<V: ?Sized + #visitor_trait_name> #visit_with_trait_name<V> for #target_type {
1086                    #visit_with_doc
1087                    #[inline]
1088                    fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
1089                        <V as #visitor_trait_name>::#visit_method_name(visitor, self #ast_path_arg)
1090                    }
1091
1092                    #[inline]
1093                    fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
1094                        #default_body
1095                    }
1096                }
1097            ));
1098        }
1099
1100        items
1101    }
1102
1103    fn implement_visit_with_for_generic_types(&self) -> Vec<Item> {
1104        let visit_trait_name = self.trait_name(false);
1105        let visit_with_trait_name = self.trait_name(true);
1106        let lifetime = self.method_lifetime();
1107        let ast_path_arg = self.arg_extra_token();
1108        let ast_path_param = self.param_extra_token();
1109        let return_type = self.return_type_token(quote!(Self));
1110        let attrs = self.base_trait_attrs();
1111        let receiver = self.parameter_type_token(quote!(self));
1112
1113        let visit_with_name = Ident::new(
1114            &format!(
1115                "{}_with{}",
1116                self.kind.method_prefix(),
1117                self.variant.method_suffix(false)
1118            ),
1119            Span::call_site(),
1120        );
1121        let visit_with_children_name = Ident::new(
1122            &format!(
1123                "{}_children_with{}",
1124                self.kind.method_prefix(),
1125                self.variant.method_suffix(false)
1126            ),
1127            Span::call_site(),
1128        );
1129
1130        let mut items = Vec::<Item>::new();
1131
1132        {
1133            // Box<T> => T
1134            match self.kind {
1135                TraitKind::Fold => {
1136                    items.push(parse_quote!(
1137                        #(#attrs)*
1138                        impl<V, T> #visit_with_trait_name<V> for std::boxed::Box<T>
1139                            where V: ?Sized + #visit_trait_name,
1140                                T: #visit_with_trait_name<V> {
1141                            #[inline]
1142                            fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
1143                                swc_visit::util::map::Map::map(self, |inner| {
1144                                    <T as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
1145                                })
1146                            }
1147                            #[inline]
1148                            fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
1149                                swc_visit::util::map::Map::map(self, |inner| {
1150                                    <T as #visit_with_trait_name<V>>::#visit_with_children_name(inner, visitor #ast_path_arg)
1151                                })
1152                            }
1153                        }
1154                    ));
1155                }
1156
1157                _ => {
1158                    let deref_expr = match self.kind {
1159                        TraitKind::Visit => {
1160                            quote!(&**self)
1161                        }
1162                        TraitKind::VisitMut => {
1163                            quote!(&mut **self)
1164                        }
1165                        TraitKind::Fold => {
1166                            unreachable!()
1167                        }
1168                    };
1169
1170                    let restore_expr = match self.kind {
1171                        TraitKind::Visit => {
1172                            quote!()
1173                        }
1174                        TraitKind::VisitMut => {
1175                            quote!()
1176                        }
1177                        TraitKind::Fold => {
1178                            unreachable!()
1179                        }
1180                    };
1181
1182                    items.push(parse_quote!(
1183                        #(#attrs)*
1184                        impl<V, T> #visit_with_trait_name<V> for std::boxed::Box<T>
1185                            where V: ?Sized + #visit_trait_name,
1186                                T: #visit_with_trait_name<V> {
1187                            #[inline]
1188                            fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
1189                                let v = <T as #visit_with_trait_name<V>>::#visit_with_name(#deref_expr, visitor #ast_path_arg);
1190                                #restore_expr
1191                                v
1192                            }
1193                            #[inline]
1194                            fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
1195                                let v = <T as #visit_with_trait_name<V>>::#visit_with_children_name(#deref_expr, visitor #ast_path_arg);
1196                                #restore_expr
1197                                v
1198                            }
1199                        }
1200                    ));
1201                }
1202            }
1203        }
1204
1205        if self.kind == TraitKind::Visit {
1206            // Vec<T> => [T]
1207            items.push(parse_quote!(
1208                #(#attrs)*
1209                impl<V, T> #visit_with_trait_name<V> for std::vec::Vec<T>
1210                    where V: ?Sized + #visit_trait_name,
1211                        [T]: #visit_with_trait_name<V> {
1212                    #[inline]
1213                    fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
1214                        let v = <[T] as #visit_with_trait_name<V>>::#visit_with_name(self, visitor #ast_path_arg);
1215                        v
1216                    }
1217
1218                    #[inline]
1219                    fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
1220                        let v = <[T] as #visit_with_trait_name<V>>::#visit_with_children_name(self, visitor #ast_path_arg);
1221                        v
1222                    }
1223                }
1224            ));
1225        }
1226
1227        items
1228    }
1229
1230    fn node_type_for_visitor_method(&self, node_type: &FieldType) -> TokenStream {
1231        match self.kind {
1232            TraitKind::Visit => match node_type {
1233                FieldType::Generic(name, inner) if name == "Vec" => {
1234                    let inner_ty = quote!(#inner);
1235
1236                    quote!([#inner_ty])
1237                }
1238
1239                _ => quote!(#node_type),
1240            },
1241            _ => quote!(#node_type),
1242        }
1243    }
1244}
1245
1246fn doc(s: &str) -> Attribute {
1247    parse_quote!(#[doc = #s])
1248}
1249
1250fn field_variant(type_name: &Ident, field: &Field) -> Option<(TokenStream, Option<Arm>)> {
1251    if let Some(field_name) = &field.ident {
1252        let variant_name = Ident::new(&field_name.to_string().to_pascal_case(), Span::call_site());
1253        let variant_doc = doc(&format!("Represents [`{type_name}::{field_name}`]"));
1254
1255        if extract_vec(extract_generic("Option", &field.ty).unwrap_or(&field.ty)).is_some() {
1256            let v = quote!(
1257                #variant_doc
1258                #variant_name(usize)
1259            );
1260            let arg = parse_quote!(
1261                Self::#variant_name(idx) => {
1262                    assert_initial_index(*idx, index);
1263
1264                    *idx = index;
1265                },
1266            );
1267            return Some((v, Some(arg)));
1268        }
1269
1270        return Some((
1271            quote!(
1272                #variant_doc
1273                #variant_name
1274            ),
1275            None,
1276        ));
1277    }
1278
1279    None
1280}
1281
1282fn extract_vec(ty: &Type) -> Option<&Type> {
1283    extract_generic("Vec", ty)
1284}
1285
1286fn extract_generic<'a>(name: &str, ty: &'a Type) -> Option<&'a Type> {
1287    if let Type::Path(p) = ty {
1288        let last = p.path.segments.last().unwrap();
1289
1290        if !last.arguments.is_empty() && last.ident == name {
1291            match &last.arguments {
1292                PathArguments::AngleBracketed(tps) => {
1293                    let arg = tps.args.first().unwrap();
1294
1295                    match arg {
1296                        GenericArgument::Type(arg) => return Some(arg),
1297                        _ => unimplemented!("generic parameter other than type"),
1298                    }
1299                }
1300                _ => unimplemented!("Box() -> T or Box without a type parameter"),
1301            }
1302        }
1303    }
1304
1305    if let Type::Reference(r) = ty {
1306        return extract_generic(name, &r.elem);
1307    }
1308
1309    None
1310}
1311
1312fn to_iter(e: TokenStream, ty: &Type, node_names: &[Ident]) -> Option<Expr> {
1313    if let Some(ty) = extract_vec(ty) {
1314        let inner_expr = to_iter(quote!(item), ty, node_names)?;
1315        return Some(parse_quote!(#e.iter().flat_map(|item| #inner_expr)));
1316    }
1317
1318    if let Some(ty) = extract_generic("Option", ty) {
1319        let inner_expr = to_iter(quote!(item), ty, node_names)?;
1320        return Some(parse_quote!(#e.iter().flat_map(|item| #inner_expr)));
1321    }
1322
1323    if let Some(ty) = extract_generic("Box", ty) {
1324        let inner_expr = to_iter(quote!(item), ty, node_names)?;
1325        return Some(parse_quote!({
1326            let item = &*#e;
1327            #inner_expr
1328        }));
1329    }
1330
1331    if let Type::Path(p) = ty {
1332        let ty = &p.path.segments.last().unwrap().ident;
1333
1334        if node_names.contains(ty) {
1335            return Some(parse_quote!(::std::iter::once(NodeRef::#ty(&#e))));
1336        }
1337
1338        None
1339    } else {
1340        todo!("to_iter for {:?}", ty);
1341    }
1342}
1343
1344fn define_fields(crate_name: &Ident, node_types: &[&Item]) -> Vec<Item> {
1345    let mut items = Vec::<Item>::new();
1346    let mut kind_enum_members = Vec::new();
1347    let mut parent_enum_members = Vec::new();
1348    let mut node_ref_enum_members = Vec::new();
1349
1350    let mut kind_set_index_arms = Vec::<Arm>::new();
1351    let mut node_ref_set_index_arms = Vec::<Arm>::new();
1352    let mut node_ref_kind_arms = Vec::<Arm>::new();
1353    let mut node_ref_iter_next_arms = Vec::<Arm>::new();
1354
1355    let node_names = node_types
1356        .iter()
1357        .filter_map(|ty| match ty {
1358            Item::Enum(data) => Some(data.ident.clone()),
1359            Item::Struct(data) => Some(data.ident.clone()),
1360            _ => None,
1361        })
1362        .collect::<Vec<_>>();
1363
1364    let is_node_ref_raw = |ty: &Type| match ty {
1365        Type::Path(p) => node_names.contains(&p.path.segments.last().unwrap().ident),
1366        _ => false,
1367    };
1368
1369    let is_node_ref = |ty: &Type| {
1370        if let Some(ty) = extract_generic("Box", ty) {
1371            return is_node_ref_raw(ty);
1372        }
1373
1374        is_node_ref_raw(ty)
1375    };
1376
1377    {
1378        let mut defs = Vec::<Item>::new();
1379
1380        defs.push(parse_quote!(
1381            use #crate_name::*;
1382        ));
1383
1384        defs.push(parse_quote!(
1385            #[inline(always)]
1386            fn assert_initial_index(idx: usize, index: usize) {
1387                #[cfg(debug_assertions)]
1388                if !(idx == usize::MAX || index == usize::MAX) {
1389                    {
1390                        panic!("Should be usize::MAX");
1391                    }
1392                }
1393            }
1394        ));
1395
1396        for ty in node_types {
1397            let type_name = match ty {
1398                Item::Enum(data) => data.ident.clone(),
1399                Item::Struct(data) => data.ident.clone(),
1400                _ => continue,
1401            };
1402
1403            let fields_enum_name = Ident::new(&format!("{type_name}Field"), Span::call_site());
1404
1405            let mut variants = Vec::new();
1406
1407            kind_set_index_arms.push(parse_quote!(
1408                Self::#type_name(v) => v.set_index(index),
1409            ));
1410            node_ref_kind_arms.push(parse_quote!(
1411                Self::#type_name(_, __field_kind) => AstParentKind::#type_name(*__field_kind),
1412            ));
1413            node_ref_set_index_arms.push(parse_quote!(
1414                Self::#type_name(_, __field_kind) => __field_kind.set_index(index),
1415            ));
1416
1417            match ty {
1418                Item::Enum(data) => {
1419                    for variant in &data.variants {
1420                        let orig_ident = &variant.ident;
1421                        let variant_name = Ident::new(
1422                            &variant.ident.to_string().to_pascal_case(),
1423                            Span::call_site(),
1424                        );
1425
1426                        let variant_doc = doc(&format!("Represents [`{type_name}::{orig_ident}`]"));
1427                        variants.push(quote!(
1428                            #variant_doc
1429                            #variant_name
1430                        ));
1431                    }
1432
1433                    kind_enum_members.push(quote!(
1434                        #type_name(#fields_enum_name)
1435                    ));
1436
1437                    parent_enum_members.push(quote!(
1438                        #type_name(&'ast #type_name, #fields_enum_name)
1439                    ));
1440
1441                    node_ref_enum_members.push(quote!(
1442                        #type_name(&'ast #type_name)
1443                    ));
1444
1445                    items.push(parse_quote!(
1446                        impl<'ast> From<&'ast #type_name> for NodeRef<'ast> {
1447                            fn from(node: &'ast #type_name) -> Self {
1448                                NodeRef::#type_name(node)
1449                            }
1450                        }
1451                    ));
1452
1453                    {
1454                        let mut arms = Vec::<Arm>::new();
1455
1456                        for variant in &data.variants {
1457                            let variant_name = &variant.ident;
1458
1459                            // TODO: Support all kinds of fields
1460                            if variant.fields.len() != 1 {
1461                                continue;
1462                            }
1463
1464                            for f in variant.fields.iter().filter(|f| is_node_ref(&f.ty)) {
1465                                let mut ty = &f.ty;
1466                                if let Some(inner) = extract_generic("Box", ty) {
1467                                    ty = inner;
1468                                }
1469
1470                                if let Some(value) = f.ident.as_ref() {
1471                                    arms.push(parse_quote!(
1472                                        #type_name::#variant_name { #value: v0 } => {
1473                                            Box::new(::std::iter::once(NodeRef::#ty(v0)))
1474                                        },
1475                                    ));
1476                                } else {
1477                                    arms.push(parse_quote!(
1478                                        #type_name::#variant_name(v0) => {
1479                                            Box::new(::std::iter::once(NodeRef::#ty(v0)))
1480                                        },
1481                                    ));
1482                                }
1483                            }
1484                        }
1485
1486                        node_ref_iter_next_arms.push(parse_quote!(
1487                            NodeRef::#type_name(node) => {
1488                                match node {
1489                                    #(#arms)*
1490
1491                                    _ => Box::new(::std::iter::empty::<NodeRef<'ast>>())
1492                                }
1493                            }
1494                        ));
1495                    }
1496
1497                    defs.push(parse_quote!(
1498                        impl #fields_enum_name {
1499                            #[inline(always)]
1500                            pub(crate) fn set_index(&mut self, _: usize) {
1501                                swc_visit::wrong_ast_path();
1502                            }
1503                        }
1504                    ));
1505                }
1506                Item::Struct(data) => {
1507                    let mut set_index_arms = Vec::<Arm>::new();
1508
1509                    for field in &data.fields {
1510                        let opt = field_variant(&type_name, field);
1511                        let opt = match opt {
1512                            Some(v) => v,
1513                            None => continue,
1514                        };
1515                        variants.push(opt.0);
1516
1517                        set_index_arms.extend(opt.1);
1518                    }
1519
1520                    kind_enum_members.push(quote!(
1521                        #type_name(#fields_enum_name)
1522                    ));
1523
1524                    parent_enum_members.push(quote!(
1525                        #type_name(&'ast #type_name, #fields_enum_name)
1526                    ));
1527
1528                    node_ref_enum_members.push(quote!(
1529                        #type_name(&'ast #type_name)
1530                    ));
1531
1532                    items.push(parse_quote!(
1533                        impl<'ast> From<&'ast #type_name> for NodeRef<'ast> {
1534                            fn from(node: &'ast #type_name) -> Self {
1535                                NodeRef::#type_name(node)
1536                            }
1537                        }
1538                    ));
1539
1540                    {
1541                        let mut iter: Expr = parse_quote!(::std::iter::empty::<NodeRef<'ast>>());
1542
1543                        match &data.fields {
1544                            Fields::Named(fields) => {
1545                                for f in fields.named.iter() {
1546                                    let ident = &f.ident;
1547                                    let iter_expr =
1548                                        to_iter(quote!(node.#ident), &f.ty, &node_names);
1549                                    if let Some(iter_expr) = iter_expr {
1550                                        iter = parse_quote!(#iter.chain(#iter_expr));
1551                                    }
1552                                }
1553                            }
1554
1555                            Fields::Unnamed(_fields) => {
1556                                // TODO: Support unnamed fields
1557                            }
1558                            Fields::Unit => {}
1559                        }
1560
1561                        node_ref_iter_next_arms.push(parse_quote!(
1562                            NodeRef::#type_name(node) => {
1563                                let iterator = #iter;
1564                                Box::new(iterator)
1565                            }
1566                        ));
1567                    }
1568
1569                    defs.push(parse_quote!(
1570                        impl #fields_enum_name {
1571                            pub(crate) fn set_index(&mut self, index: usize) {
1572                                match self {
1573                                    #(#set_index_arms)*
1574
1575                                    _ => {
1576                                        swc_visit::wrong_ast_path()
1577                                    }
1578                                }
1579                            }
1580                        }
1581                    ));
1582                }
1583                _ => continue,
1584            }
1585
1586            defs.push(parse_quote!(
1587                #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
1588                #[cfg_attr(feature = "serde-impl", derive(serde::Serialize, serde::Deserialize))]
1589                pub enum #fields_enum_name {
1590                    #(#variants),*
1591                }
1592            ))
1593        }
1594
1595        {
1596            defs.push(parse_quote!(
1597                #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
1598                #[cfg_attr(feature = "serde-impl", derive(serde::Serialize, serde::Deserialize))]
1599                pub enum AstParentKind {
1600                    #(#kind_enum_members),*
1601                }
1602            ));
1603
1604            defs.push(parse_quote!(
1605                impl ::swc_visit::ParentKind for AstParentKind {
1606                    #[inline]
1607                    fn set_index(&mut self, index: usize) {
1608                        match self {
1609                            #(#kind_set_index_arms)*
1610                        }
1611                    }
1612                }
1613            ));
1614        }
1615
1616        {
1617            defs.push(parse_quote!(
1618                #[derive(Debug, Clone, Copy)]
1619                pub enum AstParentNodeRef<'ast> {
1620                    #(#parent_enum_members),*
1621                }
1622            ));
1623            items.push(parse_quote!(
1624                #[derive(Debug, Clone, Copy)]
1625                pub enum NodeRef<'ast> {
1626                    #(#node_ref_enum_members),*
1627                }
1628            ));
1629
1630            defs.push(parse_quote!(
1631                impl<'ast> ::swc_visit::NodeRef for AstParentNodeRef<'ast> {
1632                    type ParentKind = AstParentKind;
1633
1634                    #[inline(always)]
1635                    fn kind(&self) -> AstParentKind {
1636                        self.kind()
1637                    }
1638
1639                    fn set_index(&mut self, index: usize) {
1640                        match self {
1641                            #(#node_ref_set_index_arms)*
1642                        }
1643                    }
1644                }
1645            ));
1646            defs.push(parse_quote!(
1647                #[cfg(any(docsrs, feature = "path"))]
1648                impl<'ast> AstParentNodeRef<'ast> {
1649                    #[inline]
1650                    pub fn kind(&self) -> AstParentKind {
1651                        match self {
1652                            #(#node_ref_kind_arms)*
1653                        }
1654                    }
1655                }
1656            ));
1657            items.push(parse_quote!(
1658                impl<'ast> NodeRef<'ast> {
1659                    /// This is not a part of semver-stable API. It is experimental and subject to change.
1660                    #[allow(unreachable_patterns)]
1661                    pub fn experimental_raw_children<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = NodeRef<'ast>>> {
1662                        match self {
1663                            #(#node_ref_iter_next_arms)*
1664                        }
1665                    }
1666                }
1667            ));
1668
1669            items.push(parse_quote!(
1670                impl<'ast> NodeRef<'ast> {
1671                    /// Visit all nodes in self in preorder.
1672                    ///
1673                    /// This is not a part of semver-stable API. It is
1674                    /// experimental and subject to change.
1675                    pub fn experimental_traverse(
1676                        &'ast self,
1677                    ) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
1678                        let mut queue = std::collections::VecDeque::<NodeRef<'ast>>::new();
1679                        queue.push_back(*self);
1680
1681                        Box::new(std::iter::from_fn(move || {
1682                            let node: NodeRef<'ast> = queue.pop_front()?;
1683                            {
1684                                let children = node.experimental_raw_children();
1685                                queue.extend(children);
1686                            }
1687                            Some(node)
1688                        }))
1689                    }
1690                }
1691            ));
1692        }
1693
1694        items.insert(
1695            0,
1696            parse_quote!(
1697                #[cfg(any(docsrs, feature = "path"))]
1698                pub mod fields {
1699                    #(#defs)*
1700                }
1701            ),
1702        );
1703
1704        items.push(parse_quote!(
1705            #[cfg(any(docsrs, feature = "path"))]
1706            pub use self::fields::{AstParentKind, AstParentNodeRef};
1707        ));
1708    }
1709
1710    items
1711}