swc_ecma_compat_es2015/block_scoping/
mod.rs

1use std::{iter::once, mem::take};
2
3use indexmap::IndexMap;
4use rustc_hash::{FxHashMap, FxHashSet};
5use smallvec::SmallVec;
6use swc_atoms::{atom, Atom};
7use swc_common::{util::take::Take, Mark, Spanned, SyntaxContext, DUMMY_SP};
8use swc_ecma_ast::*;
9use swc_ecma_transforms_base::helper;
10use swc_ecma_utils::{
11    find_pat_ids, function::FnEnvHoister, prepend_stmt, private_ident, quote_ident, quote_str,
12    ExprFactory, StmtLike,
13};
14use swc_ecma_visit::{
15    noop_visit_mut_type, visit_mut_obj_and_computed, visit_mut_pass, VisitMut, VisitMutWith,
16};
17use swc_trace_macro::swc_trace;
18
19mod vars;
20
21///
22///
23/// TODO(kdy1): Optimization
24///
25/// ```js
26/// let functions = [];
27/// for (let i = 0; i < 10; i++) {
28///    functions.push(function() {
29///        let i = 1;
30///        console.log(i);
31///    });
32/// }
33/// ```
34pub fn block_scoping(unresolved_mark: Mark) -> impl Pass {
35    (
36        visit_mut_pass(self::vars::block_scoped_vars()),
37        visit_mut_pass(BlockScoping {
38            unresolved_mark,
39            scope: Default::default(),
40            vars: Vec::new(),
41            var_decl_kind: VarDeclKind::Var,
42        }),
43    )
44}
45
46type ScopeStack = SmallVec<[ScopeKind; 8]>;
47
48#[derive(Debug, PartialEq, Eq)]
49enum ScopeKind {
50    Loop {
51        lexical_var: Vec<Id>,
52        args: Vec<Id>,
53        /// Produced by identifier reference and consumed by for-of/in loop.
54        used: Vec<Id>,
55        /// Map of original identifier to modified syntax context
56        mutated: FxHashMap<Id, SyntaxContext>,
57    },
58    Fn,
59    Block,
60}
61
62impl ScopeKind {
63    fn new_loop() -> Self {
64        ScopeKind::Loop {
65            lexical_var: Vec::new(),
66            args: Vec::new(),
67            used: Vec::new(),
68            mutated: Default::default(),
69        }
70    }
71}
72
73struct BlockScoping {
74    unresolved_mark: Mark,
75    scope: ScopeStack,
76    vars: Vec<VarDeclarator>,
77    var_decl_kind: VarDeclKind,
78}
79
80impl BlockScoping {
81    /// This methods remove [ScopeKind::Loop] and [ScopeKind::Fn], but not
82    /// [ScopeKind::ForLetLoop]
83    fn visit_mut_with_scope<T>(&mut self, kind: ScopeKind, node: &mut T)
84    where
85        T: VisitMutWith<Self>,
86    {
87        let remove = !matches!(kind, ScopeKind::Loop { .. });
88        self.scope.push(kind);
89
90        node.visit_mut_with(self);
91
92        if remove {
93            self.scope.pop();
94        }
95    }
96
97    fn mark_as_used(&mut self, i: Id) {
98        // Only consider the variable used in a non-ScopeKind::Loop, which means it is
99        // captured in a closure
100        for scope in self
101            .scope
102            .iter_mut()
103            .rev()
104            .skip_while(|scope| matches!(scope, ScopeKind::Loop { .. }))
105        {
106            if let ScopeKind::Loop {
107                lexical_var, used, ..
108            } = scope
109            {
110                if lexical_var.contains(&i) {
111                    used.push(i);
112                    return;
113                }
114            }
115        }
116    }
117
118    fn in_loop_body(&self) -> bool {
119        self.scope
120            .last()
121            .map(|scope| matches!(scope, ScopeKind::Loop { .. }))
122            .unwrap_or(false)
123    }
124
125    fn handle_capture_of_vars(&mut self, body: &mut Box<Stmt>) {
126        let body_stmt = &mut **body;
127
128        if let Some(ScopeKind::Loop {
129            args,
130            used,
131            mutated,
132            ..
133        }) = self.scope.pop()
134        {
135            if used.is_empty() {
136                return;
137            }
138
139            let mut env_hoister =
140                FnEnvHoister::new(SyntaxContext::empty().apply_mark(self.unresolved_mark));
141            body_stmt.visit_mut_with(&mut env_hoister);
142            let mut inits: Vec<Box<Expr>> = Vec::new();
143
144            for mut var in env_hoister.to_decl() {
145                if let Some(init) = var.init.take() {
146                    inits.push(
147                        AssignExpr {
148                            span: DUMMY_SP,
149                            op: op!("="),
150                            left: var.name.clone().try_into().unwrap(),
151                            right: init,
152                        }
153                        .into(),
154                    );
155                }
156
157                self.vars.push(var);
158            }
159
160            let mut flow_helper = FlowHelper {
161                all: &args,
162                has_break: false,
163                has_return: false,
164                has_yield: false,
165                has_await: false,
166                label: IndexMap::new(),
167                inner_label: FxHashSet::default(),
168                mutated,
169                in_switch_case: false,
170                in_nested_loop: false,
171            };
172
173            body_stmt.visit_mut_with(&mut flow_helper);
174
175            let mut body_stmt = match &mut body_stmt.take() {
176                Stmt::Block(bs) => bs.take(),
177                body => BlockStmt {
178                    span: DUMMY_SP,
179                    stmts: vec![body.take()],
180                    ..Default::default()
181                },
182            };
183
184            if !flow_helper.mutated.is_empty() {
185                let no_modification = flow_helper.mutated.is_empty();
186                let mut v = MutationHandler {
187                    map: &mut flow_helper.mutated,
188                    in_function: false,
189                };
190
191                // Modifies identifiers, and add reassignments to break / continue / return
192                body_stmt.visit_mut_with(&mut v);
193
194                if !no_modification
195                    && body_stmt
196                        .stmts
197                        .last()
198                        .map(|s| !matches!(s, Stmt::Return(..)))
199                        .unwrap_or(true)
200                {
201                    body_stmt.stmts.push(v.make_reassignment(None).into_stmt());
202                }
203            }
204
205            let var_name = private_ident!("_loop");
206
207            self.vars.push(VarDeclarator {
208                span: DUMMY_SP,
209                name: var_name.clone().into(),
210                init: Some(
211                    Function {
212                        span: DUMMY_SP,
213                        params: args
214                            .iter()
215                            .map(|i| {
216                                let ctxt = flow_helper.mutated.get(i).copied().unwrap_or(i.1);
217
218                                Param {
219                                    span: DUMMY_SP,
220                                    decorators: Default::default(),
221                                    pat: Ident::new(i.0.clone(), DUMMY_SP, ctxt).into(),
222                                }
223                            })
224                            .collect(),
225                        decorators: Default::default(),
226                        body: Some(body_stmt),
227                        is_generator: flow_helper.has_yield,
228                        is_async: flow_helper.has_await,
229                        ..Default::default()
230                    }
231                    .into(),
232                ),
233                definite: false,
234            });
235
236            let mut call: Expr = CallExpr {
237                span: DUMMY_SP,
238                callee: var_name.as_callee(),
239                args: args
240                    .iter()
241                    .cloned()
242                    .map(|i| Ident::new(i.0, DUMMY_SP, i.1).as_arg())
243                    .collect(),
244                ..Default::default()
245            }
246            .into();
247
248            if flow_helper.has_await {
249                call = AwaitExpr {
250                    span: DUMMY_SP,
251                    arg: call.into(),
252                }
253                .into();
254            }
255
256            if flow_helper.has_yield {
257                call = YieldExpr {
258                    span: DUMMY_SP,
259                    arg: Some(call.into()),
260                    delegate: true,
261                }
262                .into();
263            }
264
265            if !inits.is_empty() {
266                call = SeqExpr {
267                    span: DUMMY_SP,
268                    exprs: inits.into_iter().chain(once(Box::new(call))).collect(),
269                }
270                .into()
271            }
272
273            if flow_helper.has_return || flow_helper.has_break || !flow_helper.label.is_empty() {
274                let ret = private_ident!("_ret");
275
276                let mut stmts = vec![
277                    // var _ret = _loop(i);
278                    VarDecl {
279                        span: DUMMY_SP,
280                        kind: VarDeclKind::Var,
281                        decls: vec![VarDeclarator {
282                            span: DUMMY_SP,
283                            name: ret.clone().into(),
284                            init: Some(Box::new(call.take())),
285                            definite: false,
286                        }],
287                        ..Default::default()
288                    }
289                    .into(),
290                ];
291
292                if flow_helper.has_return {
293                    // if (_type_of(_ret) === "object") return _ret.v;
294                    stmts.push(
295                        IfStmt {
296                            span: DUMMY_SP,
297                            test: BinExpr {
298                                span: DUMMY_SP,
299                                op: op!("==="),
300                                left: {
301                                    // _type_of(_ret)
302                                    let callee = helper!(type_of);
303
304                                    CallExpr {
305                                        span: Default::default(),
306                                        callee,
307                                        args: vec![ret.clone().as_arg()],
308                                        ..Default::default()
309                                    }
310                                    .into()
311                                },
312                                //"object"
313                                right: "object".into(),
314                            }
315                            .into(),
316                            cons: Box::new(
317                                ReturnStmt {
318                                    span: DUMMY_SP,
319                                    arg: Some(ret.clone().make_member(quote_ident!("v")).into()),
320                                }
321                                .into(),
322                            ),
323                            alt: None,
324                        }
325                        .into(),
326                    )
327                }
328
329                if flow_helper.has_break {
330                    stmts.push(
331                        IfStmt {
332                            span: DUMMY_SP,
333                            test: ret.clone().make_eq(quote_str!("break")).into(),
334                            cons: BreakStmt {
335                                span: DUMMY_SP,
336                                label: None,
337                            }
338                            .into(),
339                            alt: None,
340                        }
341                        .into(),
342                    );
343                }
344
345                if !flow_helper.label.is_empty() {
346                    stmts.push(
347                        SwitchStmt {
348                            span: DUMMY_SP,
349                            discriminant: Box::new(ret.into()),
350                            cases: flow_helper
351                                .label
352                                .into_iter()
353                                .map(|(key, label)| SwitchCase {
354                                    span: DUMMY_SP,
355                                    test: Some(Box::new(key.into())),
356                                    cons: vec![match label {
357                                        Label::Break(id) => Stmt::Break(BreakStmt {
358                                            span: DUMMY_SP,
359                                            label: Some(id),
360                                        }),
361
362                                        Label::Continue(id) => Stmt::Continue(ContinueStmt {
363                                            span: DUMMY_SP,
364                                            label: Some(id),
365                                        }),
366                                    }],
367                                })
368                                .collect(),
369                        }
370                        .into(),
371                    );
372                }
373
374                *body = Box::new(
375                    BlockStmt {
376                        span: DUMMY_SP,
377                        stmts,
378                        ..Default::default()
379                    }
380                    .into(),
381                );
382                return;
383            }
384
385            *body = Box::new(call.take().into_stmt());
386        }
387    }
388
389    /// This method will turn stmt like
390    /// ```js
391    /// for (let i in [1, 2])
392    ///   for (let j in [1, 2])
393    ///     console.log(i, j)
394    /// ```
395    /// into
396    /// ```js
397    /// for (let i in [1, 2]) {
398    ///   for (let j in [1, 2]) {
399    ///     console.log(i, j)
400    ///   }
401    /// }
402    /// ```
403    /// which fixes https://github.com/swc-project/swc/issues/6573
404    fn blockify_for_stmt_body(&self, body: &mut Box<Stmt>) -> bool {
405        if !body.is_block() {
406            *body = Box::new(
407                BlockStmt {
408                    span: Default::default(),
409                    stmts: vec![*body.take()],
410                    ..Default::default()
411                }
412                .into(),
413            );
414            true
415        } else {
416            false
417        }
418    }
419
420    fn undo_blockify_for_stmt_body(&self, body: &mut Box<Stmt>, blockifyed: bool) {
421        if blockifyed {
422            let stmt = body
423                .as_mut_block()
424                .and_then(|block| (block.stmts.len() == 1).then(|| block.stmts[0].take()));
425            if let Some(stmt) = stmt {
426                *body = Box::new(stmt)
427            }
428        }
429    }
430}
431
432#[swc_trace]
433impl VisitMut for BlockScoping {
434    noop_visit_mut_type!(fail);
435
436    fn visit_mut_arrow_expr(&mut self, n: &mut ArrowExpr) {
437        n.params.visit_mut_with(self);
438        self.visit_mut_with_scope(ScopeKind::Fn, &mut n.body);
439    }
440
441    fn visit_mut_block_stmt(&mut self, n: &mut BlockStmt) {
442        let vars = take(&mut self.vars);
443        n.visit_mut_children_with(self);
444        debug_assert_eq!(self.vars, Vec::new());
445        self.vars = vars;
446    }
447
448    fn visit_mut_constructor(&mut self, f: &mut Constructor) {
449        f.key.visit_mut_with(self);
450        f.params.visit_mut_with(self);
451        self.visit_mut_with_scope(ScopeKind::Fn, &mut f.body);
452    }
453
454    fn visit_mut_do_while_stmt(&mut self, node: &mut DoWhileStmt) {
455        self.visit_mut_with_scope(ScopeKind::new_loop(), &mut node.body);
456
457        node.test.visit_mut_with(self);
458        self.handle_capture_of_vars(&mut node.body);
459    }
460
461    fn visit_mut_for_in_stmt(&mut self, node: &mut ForInStmt) {
462        let blockifyed = self.blockify_for_stmt_body(&mut node.body);
463        let lexical_var = if let ForHead::VarDecl(decl) = &node.left {
464            find_lexical_vars(decl)
465        } else {
466            Vec::new()
467        };
468        let args = lexical_var.clone();
469
470        self.visit_mut_with_scope(ScopeKind::Block, &mut node.left);
471
472        node.right.visit_mut_with(self);
473
474        let kind = ScopeKind::Loop {
475            lexical_var,
476            args,
477            used: Vec::new(),
478            mutated: Default::default(),
479        };
480
481        self.visit_mut_with_scope(kind, &mut node.body);
482        self.handle_capture_of_vars(&mut node.body);
483        self.undo_blockify_for_stmt_body(&mut node.body, blockifyed);
484    }
485
486    fn visit_mut_for_of_stmt(&mut self, node: &mut ForOfStmt) {
487        let blockifyed = self.blockify_for_stmt_body(&mut node.body);
488        let vars = if let ForHead::VarDecl(decl) = &node.left {
489            find_lexical_vars(decl)
490        } else {
491            Vec::new()
492        };
493
494        self.visit_mut_with_scope(ScopeKind::Block, &mut node.left);
495
496        let args = vars.clone();
497
498        node.right.visit_mut_with(self);
499
500        let kind = ScopeKind::Loop {
501            lexical_var: vars,
502            args,
503            used: Vec::new(),
504            mutated: Default::default(),
505        };
506
507        self.visit_mut_with_scope(kind, &mut node.body);
508        self.handle_capture_of_vars(&mut node.body);
509        self.undo_blockify_for_stmt_body(&mut node.body, blockifyed);
510    }
511
512    fn visit_mut_for_stmt(&mut self, node: &mut ForStmt) {
513        let blockifyed = self.blockify_for_stmt_body(&mut node.body);
514        let lexical_var = if let Some(VarDeclOrExpr::VarDecl(decl)) = &node.init {
515            find_lexical_vars(decl)
516        } else {
517            Vec::new()
518        };
519
520        node.init.visit_mut_with(self);
521        let args = lexical_var.clone();
522
523        node.test.visit_mut_with(self);
524        node.update.visit_mut_with(self);
525
526        let kind = ScopeKind::Loop {
527            lexical_var,
528            args,
529            used: Vec::new(),
530            mutated: Default::default(),
531        };
532        self.visit_mut_with_scope(kind, &mut node.body);
533        self.handle_capture_of_vars(&mut node.body);
534        self.undo_blockify_for_stmt_body(&mut node.body, blockifyed);
535    }
536
537    fn visit_mut_function(&mut self, f: &mut Function) {
538        f.params.visit_mut_with(self);
539        f.decorators.visit_mut_with(self);
540        self.visit_mut_with_scope(ScopeKind::Fn, &mut f.body);
541    }
542
543    fn visit_mut_getter_prop(&mut self, f: &mut GetterProp) {
544        f.key.visit_mut_with(self);
545        self.visit_mut_with_scope(ScopeKind::Fn, &mut f.body);
546    }
547
548    fn visit_mut_ident(&mut self, node: &mut Ident) {
549        let id = node.to_id();
550        self.mark_as_used(id);
551    }
552
553    fn visit_mut_module_items(&mut self, stmts: &mut Vec<ModuleItem>) {
554        self.visit_mut_stmt_like(stmts);
555    }
556
557    fn visit_mut_setter_prop(&mut self, f: &mut SetterProp) {
558        f.key.visit_mut_with(self);
559        f.param.visit_mut_with(self);
560        self.visit_mut_with_scope(ScopeKind::Fn, &mut f.body);
561    }
562
563    fn visit_mut_stmts(&mut self, n: &mut Vec<Stmt>) {
564        self.visit_mut_stmt_like(n);
565    }
566
567    fn visit_mut_switch_case(&mut self, n: &mut SwitchCase) {
568        let old_vars = self.vars.take();
569
570        n.visit_mut_children_with(self);
571
572        self.vars = old_vars;
573    }
574
575    fn visit_mut_var_decl(&mut self, var: &mut VarDecl) {
576        let old = self.var_decl_kind;
577        self.var_decl_kind = var.kind;
578        if let Some(ScopeKind::Loop { lexical_var, .. }) = self.scope.last_mut() {
579            lexical_var.extend(find_lexical_vars(var));
580        }
581
582        var.visit_mut_children_with(self);
583
584        self.var_decl_kind = old;
585
586        var.kind = VarDeclKind::Var;
587    }
588
589    fn visit_mut_var_declarator(&mut self, var: &mut VarDeclarator) {
590        var.visit_mut_children_with(self);
591
592        if self.in_loop_body() && var.init.is_none() {
593            if self.var_decl_kind == VarDeclKind::Var {
594                var.init = None
595            } else {
596                var.init = Some(Expr::undefined(var.span()))
597            }
598        }
599    }
600
601    fn visit_mut_while_stmt(&mut self, node: &mut WhileStmt) {
602        self.visit_mut_with_scope(ScopeKind::new_loop(), &mut node.body);
603
604        node.test.visit_mut_with(self);
605        self.handle_capture_of_vars(&mut node.body);
606    }
607}
608
609impl BlockScoping {
610    fn visit_mut_stmt_like<T>(&mut self, stmts: &mut Vec<T>)
611    where
612        T: StmtLike,
613        Vec<T>: VisitMutWith<Self>,
614    {
615        stmts.visit_mut_children_with(self);
616
617        if !self.vars.is_empty() {
618            prepend_stmt(
619                stmts,
620                T::from(
621                    VarDecl {
622                        span: DUMMY_SP,
623                        kind: VarDeclKind::Var,
624                        declare: false,
625                        decls: take(&mut self.vars),
626                        ..Default::default()
627                    }
628                    .into(),
629                ),
630            );
631        }
632    }
633}
634
635fn find_lexical_vars(node: &VarDecl) -> Vec<Id> {
636    if node.kind == VarDeclKind::Var {
637        return Vec::new();
638    }
639
640    find_pat_ids(&node.decls)
641}
642
643struct FlowHelper<'a> {
644    has_break: bool,
645    has_return: bool,
646    has_yield: bool,
647    has_await: bool,
648
649    // label cannot be shadowed, so it's pretty safe to use Atom
650    label: IndexMap<Atom, Label>,
651    inner_label: FxHashSet<Atom>,
652    all: &'a Vec<Id>,
653    mutated: FxHashMap<Id, SyntaxContext>,
654    in_switch_case: bool,
655
656    in_nested_loop: bool,
657}
658
659enum Label {
660    Break(Ident),
661    Continue(Ident),
662}
663
664impl FlowHelper<'_> {
665    fn check(&mut self, i: Id) {
666        if self.all.contains(&i) {
667            self.mutated.insert(
668                i,
669                SyntaxContext::empty().apply_mark(Mark::fresh(Mark::root())),
670            );
671        }
672    }
673
674    fn has_outer_label(&self, label: &Option<Ident>) -> bool {
675        match label {
676            Some(l) => !self.inner_label.contains(&l.sym),
677            None => false,
678        }
679    }
680}
681
682#[swc_trace]
683impl VisitMut for FlowHelper<'_> {
684    noop_visit_mut_type!(fail);
685
686    /// noop
687    fn visit_mut_arrow_expr(&mut self, _n: &mut ArrowExpr) {}
688
689    fn visit_mut_assign_expr(&mut self, n: &mut AssignExpr) {
690        match &n.left {
691            AssignTarget::Simple(e) => {
692                if let SimpleAssignTarget::Ident(i) = e {
693                    self.check(i.to_id());
694                }
695            }
696            AssignTarget::Pat(p) => {
697                let ids: Vec<Id> = find_pat_ids(p);
698
699                for id in ids {
700                    self.check(id);
701                }
702            }
703            #[cfg(swc_ast_unknown)]
704            _ => panic!("unable to access unknown nodes"),
705        }
706
707        n.visit_mut_children_with(self);
708    }
709
710    fn visit_mut_await_expr(&mut self, e: &mut AwaitExpr) {
711        e.visit_mut_children_with(self);
712
713        self.has_await = true;
714    }
715
716    /// https://github.com/swc-project/swc/pull/2916
717    fn visit_mut_do_while_stmt(&mut self, s: &mut DoWhileStmt) {
718        let old = self.in_nested_loop;
719        self.in_nested_loop = true;
720        s.visit_mut_children_with(self);
721        self.in_nested_loop = old;
722    }
723
724    /// https://github.com/swc-project/swc/pull/2916
725    fn visit_mut_for_in_stmt(&mut self, s: &mut ForInStmt) {
726        let old = self.in_nested_loop;
727        self.in_nested_loop = true;
728        s.visit_mut_children_with(self);
729        self.in_nested_loop = old;
730    }
731
732    /// https://github.com/swc-project/swc/pull/2916
733    fn visit_mut_for_of_stmt(&mut self, s: &mut ForOfStmt) {
734        let old = self.in_nested_loop;
735        self.in_nested_loop = true;
736        s.visit_mut_children_with(self);
737        self.in_nested_loop = old;
738    }
739
740    /// https://github.com/swc-project/swc/pull/2916
741    fn visit_mut_for_stmt(&mut self, s: &mut ForStmt) {
742        let old = self.in_nested_loop;
743        self.in_nested_loop = true;
744        s.visit_mut_children_with(self);
745        self.in_nested_loop = old;
746    }
747
748    /// noop
749    fn visit_mut_function(&mut self, _f: &mut Function) {}
750
751    /// noop
752    fn visit_mut_getter_prop(&mut self, _f: &mut GetterProp) {}
753
754    /// noop
755    fn visit_mut_setter_prop(&mut self, _f: &mut SetterProp) {}
756
757    fn visit_mut_labeled_stmt(&mut self, l: &mut LabeledStmt) {
758        self.inner_label.insert(l.label.sym.clone());
759
760        l.visit_mut_children_with(self);
761    }
762
763    fn visit_mut_stmt(&mut self, node: &mut Stmt) {
764        let span = node.span();
765
766        match node {
767            Stmt::Continue(ContinueStmt { label, .. }) => {
768                if self.in_nested_loop && !self.has_outer_label(label) {
769                    return;
770                }
771                let value = if let Some(label) = label {
772                    let value: Atom = format!("continue|{}", label.sym).into();
773                    self.label
774                        .insert(value.clone(), Label::Continue(label.clone()));
775                    value
776                } else {
777                    atom!("continue")
778                };
779
780                *node = ReturnStmt {
781                    span,
782                    arg: Some(
783                        Lit::Str(Str {
784                            span,
785                            value: value.into(),
786                            raw: None,
787                        })
788                        .into(),
789                    ),
790                }
791                .into();
792            }
793            Stmt::Break(BreakStmt { label, .. }) => {
794                if (self.in_switch_case || self.in_nested_loop) && !self.has_outer_label(label) {
795                    return;
796                }
797                let value = if let Some(label) = label {
798                    let value: Atom = format!("break|{}", label.sym).into();
799                    self.label
800                        .insert(value.clone(), Label::Break(label.clone()));
801                    value
802                } else {
803                    self.has_break = true;
804                    atom!("break")
805                };
806                *node = ReturnStmt {
807                    span,
808                    arg: Some(
809                        Lit::Str(Str {
810                            span,
811                            value: value.into(),
812                            raw: None,
813                        })
814                        .into(),
815                    ),
816                }
817                .into();
818            }
819            Stmt::Return(s) => {
820                self.has_return = true;
821                s.visit_mut_with(self);
822
823                *node = ReturnStmt {
824                    span,
825                    arg: Some(
826                        ObjectLit {
827                            span,
828                            props: vec![PropOrSpread::Prop(Box::new(Prop::KeyValue(
829                                KeyValueProp {
830                                    key: PropName::Ident(IdentName::new(atom!("v"), DUMMY_SP)),
831                                    value: s.arg.take().unwrap_or_else(|| {
832                                        Box::new(Expr::Unary(UnaryExpr {
833                                            span: DUMMY_SP,
834                                            op: op!("void"),
835                                            arg: Expr::undefined(DUMMY_SP),
836                                        }))
837                                    }),
838                                },
839                            )))],
840                        }
841                        .into(),
842                    ),
843                }
844                .into();
845            }
846            _ => node.visit_mut_children_with(self),
847        }
848    }
849
850    fn visit_mut_switch_case(&mut self, n: &mut SwitchCase) {
851        let old = self.in_switch_case;
852        self.in_switch_case = true;
853
854        n.visit_mut_children_with(self);
855
856        self.in_switch_case = old;
857    }
858
859    fn visit_mut_update_expr(&mut self, n: &mut UpdateExpr) {
860        if let Expr::Ident(ref i) = *n.arg {
861            self.check(i.to_id())
862        }
863        n.visit_mut_children_with(self);
864    }
865
866    /// https://github.com/swc-project/swc/pull/2916
867    fn visit_mut_while_stmt(&mut self, s: &mut WhileStmt) {
868        let old = self.in_nested_loop;
869        self.in_nested_loop = true;
870        s.visit_mut_children_with(self);
871        self.in_nested_loop = old;
872    }
873
874    fn visit_mut_yield_expr(&mut self, e: &mut YieldExpr) {
875        e.visit_mut_children_with(self);
876
877        self.has_yield = true;
878    }
879}
880
881struct MutationHandler<'a> {
882    map: &'a mut FxHashMap<Id, SyntaxContext>,
883    in_function: bool,
884}
885
886impl MutationHandler<'_> {
887    fn make_reassignment(&self, orig: Option<Box<Expr>>) -> Expr {
888        if self.map.is_empty() {
889            return *orig.unwrap_or_else(|| Expr::undefined(DUMMY_SP));
890        }
891
892        let mut exprs = Vec::with_capacity(self.map.len() + 1);
893
894        for (id, ctxt) in &*self.map {
895            exprs.push(
896                AssignExpr {
897                    span: DUMMY_SP,
898                    left: Ident::new(id.0.clone(), DUMMY_SP, id.1).into(),
899                    op: op!("="),
900                    right: Box::new(Ident::new(id.0.clone(), DUMMY_SP, *ctxt).into()),
901                }
902                .into(),
903            );
904        }
905        exprs.push(orig.unwrap_or_else(|| Expr::undefined(DUMMY_SP)));
906
907        SeqExpr {
908            span: DUMMY_SP,
909            exprs,
910        }
911        .into()
912    }
913}
914
915#[swc_trace]
916impl VisitMut for MutationHandler<'_> {
917    noop_visit_mut_type!(fail);
918
919    visit_mut_obj_and_computed!();
920
921    fn visit_mut_arrow_expr(&mut self, n: &mut ArrowExpr) {
922        let old = self.in_function;
923        self.in_function = true;
924
925        n.visit_mut_children_with(self);
926
927        self.in_function = old;
928    }
929
930    fn visit_mut_function(&mut self, n: &mut Function) {
931        let old = self.in_function;
932        self.in_function = true;
933
934        n.visit_mut_children_with(self);
935
936        self.in_function = old;
937    }
938
939    fn visit_mut_ident(&mut self, n: &mut Ident) {
940        if let Some(&ctxt) = self.map.get(&n.to_id()) {
941            n.ctxt = ctxt;
942        }
943    }
944
945    fn visit_mut_return_stmt(&mut self, n: &mut ReturnStmt) {
946        n.visit_mut_children_with(self);
947        if self.in_function || self.map.is_empty() {
948            return;
949        }
950
951        let val = n.arg.take();
952
953        n.arg = Some(Box::new(self.make_reassignment(val)))
954    }
955}