swc_ecma_transforms_macros/
parallel.rs

1#![allow(non_snake_case)]
2
3use proc_macro2::{Span, TokenStream};
4use syn::{parse_quote, Expr, Ident, ImplItem, ImplItemFn, ItemImpl, Meta, Type};
5
6use crate::common::Mode;
7
8pub fn expand(attr: TokenStream, mut item: ItemImpl) -> ItemImpl {
9    let mode = {
10        let p = &item.trait_.as_ref().unwrap().1;
11        if p.is_ident("Fold") {
12            Mode::Fold
13        } else if p.is_ident("VisitMut") {
14            Mode::VisitMut
15        } else {
16            unimplemented!("Unknown visitor type: {:?}", p)
17        }
18    };
19    let meta = if attr.is_empty() {
20        None
21    } else {
22        Some(syn::parse2::<Meta>(attr).expect("failed to parse meta"))
23    };
24    let explode = meta
25        .as_ref()
26        .map(|v| v.path().is_ident("explode"))
27        .unwrap_or(false);
28
29    item.items.push(ImplItem::Fn(make_par_visit_method(
30        mode,
31        "module_items",
32        explode,
33    )));
34    item.items
35        .push(ImplItem::Fn(make_par_visit_method(mode, "stmts", explode)));
36
37    item
38}
39
40fn node_type(suffix: &str) -> Type {
41    match suffix {
42        "module_items" => parse_quote!(ModuleItem),
43        "stmts" => parse_quote!(Stmt),
44        _ => {
45            unimplemented!("Unknown suffix `{}`", suffix)
46        }
47    }
48}
49
50fn post_visit_hook(mode: Mode, suffix: &str) -> Option<Expr> {
51    match suffix {
52        "module_items" => Some(match mode {
53            Mode::Fold => parse_quote!(
54                swc_ecma_transforms_base::perf::Parallel::after_module_items(self, &mut nodes)
55            ),
56
57            Mode::VisitMut => parse_quote!(
58                swc_ecma_transforms_base::perf::Parallel::after_module_items(self, nodes)
59            ),
60        }),
61        "stmts" => Some(match mode {
62            Mode::Fold => parse_quote!(swc_ecma_transforms_base::perf::Parallel::after_stmts(
63                self, &mut nodes
64            )),
65
66            Mode::VisitMut => parse_quote!(swc_ecma_transforms_base::perf::Parallel::after_stmts(
67                self, nodes
68            )),
69        }),
70        _ => None,
71    }
72}
73
74fn explode_hook_method_name(explode: bool, suffix: &str) -> Option<Ident> {
75    if !explode {
76        return None;
77    }
78    match suffix {
79        "module_items" => Some(Ident::new("after_one_module_item", Span::call_site())),
80        "stmts" => Some(Ident::new("after_one_stmt", Span::call_site())),
81        _ => None,
82    }
83}
84
85fn make_par_visit_method(mode: Mode, suffix: &str, explode: bool) -> ImplItemFn {
86    let method_name = Ident::new(&format!("{}_{}", mode.prefix(), suffix), Span::call_site());
87    let hook = post_visit_hook(mode, suffix);
88    let explode_method_name = explode_hook_method_name(explode, suffix);
89    let node_type = node_type(suffix);
90
91    match (mode, explode_method_name) {
92        (Mode::Fold, Some(explode_method_name)) => parse_quote!(
93            fn #method_name(&mut self, mut nodes: Vec<#node_type>) -> Vec<#node_type> {
94                use swc_common::errors::HANDLER;
95                use swc_ecma_transforms_base::perf::{ParExplode, Parallel};
96                use swc_ecma_visit::FoldWith;
97
98                let mut buf = Vec::with_capacity(nodes.len());
99
100                for node in nodes {
101                    let mut visitor = Parallel::create(&*self);
102                    let node = node.fold_with(&mut visitor);
103                    ParExplode::#explode_method_name(&mut visitor, &mut buf);
104                    buf.push(node);
105                }
106
107                let mut nodes = buf;
108                {
109                    #hook;
110                }
111
112                nodes
113            }
114        ),
115        (Mode::Fold, None) => parse_quote!(
116            fn #method_name(&mut self, nodes: Vec<#node_type>) -> Vec<#node_type> {
117                use swc_common::errors::HANDLER;
118                use swc_ecma_transforms_base::perf::Parallel;
119                use swc_ecma_visit::FoldWith;
120
121                let mut nodes = nodes.fold_children_with(self);
122                {
123                    #hook;
124                }
125
126                nodes
127            }
128        ),
129        (Mode::VisitMut, Some(explode_method_name)) => parse_quote!(
130            fn #method_name(&mut self, nodes: &mut Vec<#node_type>) {
131                use std::mem::take;
132
133                use swc_common::errors::HANDLER;
134                use swc_ecma_transforms_base::perf::{ParExplode, Parallel};
135                use swc_ecma_visit::VisitMutWith;
136
137                let mut buf = Vec::with_capacity(nodes.len());
138
139                for mut node in take(nodes) {
140                    let mut visitor = Parallel::create(&*self);
141                    node.visit_mut_with(&mut visitor);
142                    ParExplode::#explode_method_name(&mut visitor, &mut buf);
143                    buf.push(node);
144                }
145
146                *nodes = buf;
147
148                {
149                    #hook;
150                }
151            }
152        ),
153        (Mode::VisitMut, None) => parse_quote!(
154            fn #method_name(&mut self, nodes: &mut Vec<#node_type>) {
155                use swc_common::errors::HANDLER;
156                use swc_ecma_transforms_base::perf::Parallel;
157                use swc_ecma_visit::VisitMutWith;
158
159                nodes.visit_mut_children_with(self);
160                {
161                    #hook;
162                }
163            }
164        ),
165    }
166}
swc_ecma_transforms_macros/parallel.rs

swc_ecma_transforms_macros/
parallel.rs
