swc_ecma_ext_transforms/
jest.rs

1use phf::phf_set;
2use swc_common::util::take::Take;
3use swc_ecma_ast::*;
4use swc_ecma_utils::{prepend_stmts, StmtLike};
5use swc_ecma_visit::{noop_visit_mut_type, visit_mut_pass, VisitMut, VisitMutWith};
6
7static HOIST_METHODS: phf::Set<&str> = phf_set![
8    "mock",
9    "unmock",
10    "enableAutomock",
11    "disableAutomock",
12    "deepUnmock"
13];
14
15pub fn jest() -> impl Pass {
16    visit_mut_pass(Jest::default())
17}
18
19#[derive(Default)]
20struct Jest {
21    imported: Vec<Id>,
22}
23
24impl Jest {
25    fn visit_mut_stmt_like<T>(&mut self, orig: &mut Vec<T>)
26    where
27        T: StmtLike + VisitMutWith<Self>,
28    {
29        for item in &mut *orig {
30            item.visit_mut_with(self);
31        }
32
33        let items = orig.take();
34
35        let mut new = Vec::with_capacity(items.len());
36        let mut hoisted = Vec::with_capacity(8);
37        items.into_iter().for_each(|item| {
38            match item.try_into_stmt() {
39                Ok(stmt) => match &stmt {
40                    Stmt::Expr(ExprStmt { expr, .. }) => match &**expr {
41                        Expr::Call(CallExpr {
42                            callee: Callee::Expr(callee),
43                            ..
44                        }) => {
45                            if self.should_hoist(callee) {
46                                hoisted.push(T::from(stmt))
47                            } else {
48                                new.push(T::from(stmt))
49                            }
50                        }
51                        _ => new.push(T::from(stmt)),
52                    },
53
54                    _ => new.push(T::from(stmt)),
55                },
56                Err(node) => new.push(node),
57            };
58        });
59
60        prepend_stmts(&mut new, hoisted.into_iter());
61
62        *orig = new;
63    }
64
65    fn should_hoist(&self, e: &Expr) -> bool {
66        match e {
67            Expr::Ident(i) => self.imported.iter().any(|imported| *imported == i.to_id()),
68
69            Expr::Member(
70                callee @ MemberExpr {
71                    prop: MemberProp::Ident(prop),
72                    ..
73                },
74            ) => is_global_jest(&callee.obj) && HOIST_METHODS.contains(&*prop.sym),
75
76            _ => false,
77        }
78    }
79}
80
81impl VisitMut for Jest {
82    noop_visit_mut_type!();
83
84    fn visit_mut_module_items(&mut self, items: &mut Vec<ModuleItem>) {
85        for item in items.iter() {
86            if let ModuleItem::ModuleDecl(ModuleDecl::Import(ImportDecl {
87                specifiers, src, ..
88            })) = item
89            {
90                if src.value == "@jest/globals" {
91                    for s in specifiers {
92                        match s {
93                            ImportSpecifier::Named(ImportNamedSpecifier {
94                                local,
95                                imported: None,
96                                is_type_only: false,
97                                ..
98                            }) => {
99                                if HOIST_METHODS.contains(&*local.sym) {
100                                    self.imported.push(local.to_id());
101                                }
102                            }
103
104                            ImportSpecifier::Named(ImportNamedSpecifier {
105                                local,
106                                imported: Some(exported),
107                                is_type_only: false,
108                                ..
109                            }) => {
110                                if HOIST_METHODS.contains(exported.atom()) {
111                                    self.imported.push(local.to_id());
112                                }
113                            }
114                            _ => {}
115                        }
116                    }
117                }
118            }
119        }
120
121        self.visit_mut_stmt_like(items)
122    }
123
124    fn visit_mut_stmts(&mut self, stmts: &mut Vec<Stmt>) {
125        self.visit_mut_stmt_like(stmts)
126    }
127}
128
129fn is_global_jest(e: &Expr) -> bool {
130    match e {
131        Expr::Ident(i) => i.sym == *"jest",
132        Expr::Member(MemberExpr { obj, .. }) => is_global_jest(obj),
133        Expr::Call(CallExpr {
134            callee: Callee::Expr(callee),
135            ..
136        }) => is_global_jest(callee),
137        _ => false,
138    }
139}
swc_ecma_ext_transforms/jest.rs

swc_ecma_ext_transforms/
jest.rs
