Skip to main content

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