swc_ecma_transforms_macros/
fast.rs

1use proc_macro2::TokenStream;
2use quote::quote;
3use swc_macros_common::call_site;
4use syn::{parse_quote, FnArg, Ident, ImplItem, ImplItemFn, ItemImpl, Pat, Path};
5
6use crate::common::Mode;
7
8pub fn expand(attr: TokenStream, item: ItemImpl) -> ItemImpl {
9    let expander = Expander {
10        handler: syn::parse2(attr).expect("Usage should be like #[fast_path(ArrowVisitor)]"),
11        mode: detect_mode(&item),
12    };
13    let items = expander.inject_default_methods(item.items);
14
15    ItemImpl {
16        items: items
17            .into_iter()
18            .map(|item| match item {
19                ImplItem::Fn(m) => ImplItem::Fn(expander.patch_method(m)),
20                _ => item,
21            })
22            .collect(),
23        ..item
24    }
25}
26
27fn detect_mode(i: &ItemImpl) -> Mode {
28    if i.items.iter().any(|item| match item {
29        ImplItem::Fn(m) => m.sig.ident.to_string().starts_with("fold"),
30        _ => false,
31    }) {
32        return Mode::Fold;
33    }
34
35    Mode::VisitMut
36}
37
38struct Expander {
39    mode: Mode,
40    handler: Path,
41}
42
43impl Expander {
44    fn inject_default_methods(&self, mut items: Vec<ImplItem>) -> Vec<ImplItem> {
45        let list = &[
46            ("stmt", quote!(swc_ecma_ast::Stmt)),
47            ("stmts", quote!(Vec<swc_ecma_ast::Stmt>)),
48            ("module_decl", quote!(swc_ecma_ast::ModuleDecl)),
49            ("module_item", quote!(swc_ecma_ast::ModuleItem)),
50            ("module_items", quote!(Vec<swc_ecma_ast::ModuleItem>)),
51            ("expr", quote!(swc_ecma_ast::Expr)),
52            ("exprs", quote!(Vec<Box<swc_ecma_ast::Expr>>)),
53            ("decl", quote!(swc_ecma_ast::Decl)),
54            ("pat", quote!(swc_ecma_ast::Pat)),
55        ];
56
57        for (name, ty) in list {
58            let has = items.iter().any(|item| match item {
59                ImplItem::Fn(i) => i.sig.ident.to_string().ends_with(name),
60                _ => false,
61            });
62            if has {
63                continue;
64            }
65            let name = Ident::new(&format!("{}_{}", self.mode.prefix(), name), call_site());
66
67            let method = match self.mode {
68                Mode::Fold => parse_quote!(
69                    fn #name(&mut self, node: #ty) -> #ty {
70                        node.fold_children_with(self)
71                    }
72                ),
73                Mode::VisitMut => parse_quote!(
74                    fn #name(&mut self, node: &mut #ty) {
75                        node.visit_mut_children_with(self)
76                    }
77                ),
78            };
79
80            items.push(method);
81        }
82
83        items
84    }
85
86    /// Add fast path to a method
87    fn patch_method(&self, mut m: ImplItemFn) -> ImplItemFn {
88        let ty_arg = m
89            .sig
90            .inputs
91            .last()
92            .expect("method of Fold / VisitMut must accept two parameters");
93        let ty_arg = match ty_arg {
94            FnArg::Receiver(_) => unreachable!(),
95            FnArg::Typed(ty) => ty,
96        };
97        if m.sig.ident == "visit_mut_ident" || m.sig.ident == "fold_ident" {
98            return m;
99        }
100        if m.block.stmts.is_empty() {
101            return m;
102        }
103
104        let arg = match &*ty_arg.pat {
105            Pat::Ident(i) => &i.ident,
106            _ => unimplemented!(
107                "Fast-path injection for Fold / VisitMut where pattern is not an ident"
108            ),
109        };
110
111        let checker = &self.handler;
112
113        let fast_path = match self.mode {
114            Mode::Fold => parse_quote!(
115                if !swc_ecma_transforms_base::perf::should_work::<#checker, _>(&#arg) {
116                    return #arg;
117                }
118            ),
119            Mode::VisitMut => parse_quote!(
120                if !swc_ecma_transforms_base::perf::should_work::<#checker, _>(&*#arg) {
121                    return;
122                }
123            ),
124        };
125        let mut stmts = vec![fast_path];
126        stmts.extend(m.block.stmts);
127
128        m.block.stmts = stmts;
129        m
130    }
131}
swc_ecma_transforms_macros/fast.rs

swc_ecma_transforms_macros/
fast.rs
