next_custom_transforms/transforms/
optimize_server_react.rs

1// This transform optimizes React code for the server bundle, in particular:
2// - Removes `useEffect` and `useLayoutEffect` calls
3// - Refactors `useState` calls (under the `optimize_use_state` flag)
4
5use serde::Deserialize;
6use swc_core::{
7    atoms::atom,
8    common::DUMMY_SP,
9    ecma::{
10        ast::*,
11        visit::{fold_pass, Fold, FoldWith},
12    },
13};
14
15#[derive(Clone, Debug, Deserialize)]
16pub struct Config {
17    pub optimize_use_state: bool,
18}
19
20pub fn optimize_server_react(config: Config) -> impl Pass {
21    fold_pass(OptimizeServerReact {
22        optimize_use_state: config.optimize_use_state,
23        ..Default::default()
24    })
25}
26
27#[derive(Debug, Default)]
28struct OptimizeServerReact {
29    optimize_use_state: bool,
30    react_ident: Option<Id>,
31    use_state_ident: Option<Id>,
32    use_effect_ident: Option<Id>,
33    use_layout_effect_ident: Option<Id>,
34}
35
36fn effect_has_side_effect_deps(call: &CallExpr) -> bool {
37    if call.args.len() != 2 {
38        return false;
39    }
40
41    // We can't optimize if the effect has a function call as a dependency:
42    // useEffect(() => {}, x())
43    if let box Expr::Call(_) = &call.args[1].expr {
44        return true;
45    }
46
47    // As well as:
48    // useEffect(() => {}, [x()])
49    if let box Expr::Array(arr) = &call.args[1].expr {
50        for elem in arr.elems.iter().flatten() {
51            if let ExprOrSpread {
52                expr: box Expr::Call(_),
53                ..
54            } = elem
55            {
56                return true;
57            }
58        }
59    }
60
61    false
62}
63
64fn wrap_expr_with_env_prod_condition(call: CallExpr) -> Expr {
65    // Wrap the call expression with the condition
66    // turn it into `process.env.__NEXT_PRIVATE_MINIMIZE_MACRO_FALSE && <call>`.
67    // And `process.env.__NEXT_PRIVATE_MINIMIZE_MACRO_FALSE` will be treated as `false` in
68    // minification. In this way the expression and dependencies are still available in
69    // compilation during bundling, but will be removed in the final DEC.
70    Expr::Bin(BinExpr {
71        span: DUMMY_SP,
72        left: Box::new(Expr::Member(MemberExpr {
73            obj: (Box::new(Expr::Member(MemberExpr {
74                obj: (Box::new(Expr::Ident(Ident {
75                    sym: atom!("process"),
76                    span: DUMMY_SP,
77                    ..Default::default()
78                }))),
79                prop: MemberProp::Ident(IdentName {
80                    sym: atom!("env"),
81                    span: DUMMY_SP,
82                }),
83                span: DUMMY_SP,
84            }))),
85            prop: (MemberProp::Ident(IdentName {
86                sym: atom!("__NEXT_PRIVATE_MINIMIZE_MACRO_FALSE"),
87                span: DUMMY_SP,
88            })),
89            span: DUMMY_SP,
90        })),
91        op: op!("&&"),
92        right: Box::new(Expr::Call(call)),
93    })
94}
95
96impl Fold for OptimizeServerReact {
97    fn fold_module_items(&mut self, items: Vec<ModuleItem>) -> Vec<ModuleItem> {
98        let mut new_items = vec![];
99
100        for item in items {
101            new_items.push(item.clone().fold_with(self));
102
103            if let ModuleItem::ModuleDecl(ModuleDecl::Import(import_decl)) = &item {
104                if import_decl.src.value != "react" {
105                    continue;
106                }
107                for specifier in &import_decl.specifiers {
108                    if let ImportSpecifier::Named(named_import) = specifier {
109                        let name = match &named_import.imported {
110                            Some(n) => match &n {
111                                ModuleExportName::Ident(n) => n.sym.to_string(),
112                                ModuleExportName::Str(n) => n.value.to_string(),
113                            },
114                            None => named_import.local.sym.to_string(),
115                        };
116
117                        if name == "useState" {
118                            self.use_state_ident = Some(named_import.local.to_id());
119                        } else if name == "useEffect" {
120                            self.use_effect_ident = Some(named_import.local.to_id());
121                        } else if name == "useLayoutEffect" {
122                            self.use_layout_effect_ident = Some(named_import.local.to_id());
123                        }
124                    } else if let ImportSpecifier::Default(default_import) = specifier {
125                        self.react_ident = Some(default_import.local.to_id());
126                    }
127                }
128            }
129        }
130
131        new_items
132    }
133
134    fn fold_expr(&mut self, expr: Expr) -> Expr {
135        if let Expr::Call(call) = &expr {
136            if let Callee::Expr(box Expr::Ident(f)) = &call.callee {
137                // Mark `useEffect` as DCE'able
138                if let Some(use_effect_ident) = &self.use_effect_ident {
139                    if &f.to_id() == use_effect_ident && !effect_has_side_effect_deps(call) {
140                        // return Expr::Lit(Lit::Null(Null { span: DUMMY_SP }));
141                        return wrap_expr_with_env_prod_condition(call.clone());
142                    }
143                }
144                // Mark `useLayoutEffect` as DCE'able
145                if let Some(use_layout_effect_ident) = &self.use_layout_effect_ident {
146                    if &f.to_id() == use_layout_effect_ident && !effect_has_side_effect_deps(call) {
147                        return wrap_expr_with_env_prod_condition(call.clone());
148                    }
149                }
150            } else if let Some(react_ident) = &self.react_ident {
151                if let Callee::Expr(box Expr::Member(member)) = &call.callee {
152                    if let box Expr::Ident(f) = &member.obj {
153                        if &f.to_id() == react_ident {
154                            if let MemberProp::Ident(i) = &member.prop {
155                                // Mark `React.useEffect` and `React.useLayoutEffect` as DCE'able
156                                // calls in production
157                                if i.sym == "useEffect" || i.sym == "useLayoutEffect" {
158                                    return wrap_expr_with_env_prod_condition(call.clone());
159                                }
160                            }
161                        }
162                    }
163                }
164            }
165        }
166
167        expr.fold_children_with(self)
168    }
169
170    // const [state, setState] = useState(x);
171    // const [state, setState] = React.useState(x);
172    fn fold_var_declarator(&mut self, decl: VarDeclarator) -> VarDeclarator {
173        if !self.optimize_use_state {
174            return decl;
175        }
176
177        if let Pat::Array(array_pat) = &decl.name {
178            if array_pat.elems.len() == 2 {
179                if let Some(box Expr::Call(call)) = &decl.init {
180                    if let Callee::Expr(box Expr::Ident(f)) = &call.callee {
181                        if let Some(use_state_ident) = &self.use_state_ident {
182                            if &f.to_id() == use_state_ident && call.args.len() == 1 {
183                                // We do the optimization only if the arg is a literal or a
184                                // type that we can
185                                // be sure is not a function (e.g. {} or [] lit).
186                                // This is because useState allows a function as the
187                                // initialiser.
188                                match &call.args[0].expr {
189                                    box Expr::Lit(_) | box Expr::Object(_) | box Expr::Array(_) => {
190                                        // const [state, setState] = [x, () => {}];
191                                        return VarDeclarator {
192                                            definite: false,
193                                            name: decl.name.clone(),
194                                            init: Some(Box::new(Expr::Array(ArrayLit {
195                                                elems: vec![
196                                                    Some(call.args[0].expr.clone().into()),
197                                                    Some(
198                                                        Expr::Arrow(ArrowExpr {
199                                                            span: DUMMY_SP,
200                                                            body: Box::new(BlockStmtOrExpr::Expr(
201                                                                Box::new(Expr::Lit(Lit::Null(
202                                                                    Null { span: DUMMY_SP },
203                                                                ))),
204                                                            )),
205                                                            is_async: false,
206                                                            is_generator: false,
207                                                            params: vec![],
208                                                            ..Default::default()
209                                                        })
210                                                        .into(),
211                                                    ),
212                                                ],
213                                                span: DUMMY_SP,
214                                            }))),
215                                            span: DUMMY_SP,
216                                        };
217                                    }
218                                    _ => {}
219                                }
220                            }
221                        }
222                    }
223                }
224            }
225        }
226
227        decl.fold_children_with(self)
228    }
229}