fix(es/transforms/compat): Handle references to arguments inside arrow functions and block scoped loops (#1585)

Co-authored-by: 강동윤 <kdy1997.dev@gmail.com>

ecmascript/transforms/compat/Cargo.toml

@@ -6,7 +6,7 @@ edition = "2018"
 license = "Apache-2.0/MIT"
 name = "swc_ecma_transforms_compat"
 repository = "https://github.com/swc-project/swc.git"
-version = "0.13.3"
+version = "0.13.4"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]

ecmascript/transforms/compat/src/es2015/arrow.rs

@@ -1,10 +1,11 @@
+use swc_atoms::js_word;
 use swc_common::{Spanned, DUMMY_SP};
 use swc_ecma_ast::*;
 use swc_ecma_transforms_base::perf::Check;
 use swc_ecma_transforms_macros::fast_path;
-use swc_ecma_utils::contains_this_expr;
 use swc_ecma_utils::quote_ident;
 use swc_ecma_utils::ExprFactory;
+use swc_ecma_utils::{contains_this_expr, private_ident};
 use swc_ecma_visit::noop_visit_type;
 use swc_ecma_visit::{noop_fold_type, Fold, FoldWith, Node, Visit};
 
@@ -61,15 +62,15 @@ pub fn arrow() -> impl Fold {
 }
 
 #[derive(Default)]
-struct Arrow;
+struct Arrow {
+    in_arrow: bool,
+}
 
 #[fast_path(ArrowVisitor)]
 impl Fold for Arrow {
     noop_fold_type!();
 
     fn fold_expr(&mut self, e: Expr) -> Expr {
-        let e = e.fold_children_with(self);
-
         match e {
             Expr::Arrow(ArrowExpr {
                 span,
@@ -82,19 +83,54 @@ impl Fold for Arrow {
             }) => {
                 let used_this = contains_this_expr(&body);
 
+                let mut params: Vec<Param> = params
+                    .fold_with(self)
+                    .into_iter()
+                    .map(|pat| Param {
+                        span: DUMMY_SP,
+                        decorators: Default::default(),
+                        pat,
+                    })
+                    .collect();
+
+                // If we aren't already in an arrow expression, check if there is
+                // any arguments references within this function. If so, we need to
+                // replace them so that they refer to the parent function environment.
+                let (body, used_arguments) = if !self.in_arrow {
+                    let mut arguments_replacer = ArgumentsReplacer {
+                        arguments: private_ident!("_arguments"),
+                        found: false,
+                    };
+                    let body = body.fold_with(&mut arguments_replacer);
+
+                    if arguments_replacer.found {
+                        params.insert(
+                            0,
+                            Param {
+                                span: DUMMY_SP,
+                                decorators: Default::default(),
+                                pat: Pat::Ident(BindingIdent::from(
+                                    arguments_replacer.arguments.clone(),
+                                )),
+                            },
+                        );
+                    }
+                    (body, arguments_replacer.found)
+                } else {
+                    (body, false)
+                };
+
+                let in_arrow = self.in_arrow;
+                self.in_arrow = true;
+                let body = body.fold_with(self);
+                self.in_arrow = in_arrow;
+
                 let fn_expr = Expr::Fn(FnExpr {
                     ident: None,
                     function: Function {
                         decorators: vec![],
                         span,
-                        params: params
-                            .into_iter()
-                            .map(|pat| Param {
-                                span: DUMMY_SP,
-                                decorators: Default::default(),
-                                pat,
-                            })
-                            .collect(),
+                        params,
                         is_async,
                         is_generator,
                         body: Some(match body {
@@ -112,20 +148,85 @@ impl Fold for Arrow {
                     },
                 });
 
-                if !used_this {
+                if !used_this && !used_arguments {
                     return fn_expr;
                 }
 
+                let mut args = vec![ThisExpr { span: DUMMY_SP }.as_arg()];
+                if used_arguments {
+                    args.push(quote_ident!("arguments").as_arg());
+                }
+
                 Expr::Call(CallExpr {
                     span,
                     callee: fn_expr.make_member(quote_ident!("bind")).as_callee(),
-                    args: vec![ThisExpr { span: DUMMY_SP }.as_arg()],
+                    args,
                     type_args: Default::default(),
                 })
             }
-            _ => e,
+            _ => e.fold_children_with(self),
         }
     }
+
+    fn fold_constructor(&mut self, c: Constructor) -> Constructor {
+        let in_arrow = self.in_arrow;
+        self.in_arrow = false;
+        let res = c.fold_children_with(self);
+        self.in_arrow = in_arrow;
+        res
+    }
+
+    fn fold_function(&mut self, f: Function) -> Function {
+        let in_arrow = self.in_arrow;
+        self.in_arrow = false;
+        let res = f.fold_children_with(self);
+        self.in_arrow = in_arrow;
+        res
+    }
+}
+
+struct ArgumentsReplacer {
+    arguments: Ident,
+    found: bool,
+}
+
+impl Fold for ArgumentsReplacer {
+    noop_fold_type!();
+
+    fn fold_expr(&mut self, e: Expr) -> Expr {
+        match e {
+            Expr::Ident(id) if id.sym == js_word!("arguments") => {
+                self.found = true;
+                Expr::Ident(self.arguments.clone())
+            }
+            _ => e.fold_children_with(self),
+        }
+    }
+
+    /// Don't recurse into prop of member expression unless computed
+    fn fold_member_expr(&mut self, m: MemberExpr) -> MemberExpr {
+        let prop = match &*m.prop {
+            Expr::Ident(id) if !m.computed => Expr::Ident(id.clone()),
+            v => v.clone().fold_with(self),
+        };
+
+        MemberExpr {
+            span: m.span,
+            obj: m.obj.fold_with(self),
+            prop: Box::new(prop),
+            computed: m.computed,
+        }
+    }
+
+    /// Don't recurse into constructor
+    fn fold_constructor(&mut self, c: Constructor) -> Constructor {
+        c
+    }
+
+    /// Don't recurse into fn
+    fn fold_function(&mut self, f: Function) -> Function {
+        f
+    }
 }
 
 #[derive(Default)]

ecmascript/transforms/compat/src/es2015/block_scoping.rs

@@ -10,8 +10,8 @@ use swc_ecma_utils::quote_ident;
 use swc_ecma_utils::quote_str;
 use swc_ecma_utils::undefined;
 use swc_ecma_utils::{
-    contains_this_expr, find_ids, ident::IdentLike, prepend, var::VarCollector, ExprFactory, Id,
-    StmtLike,
+    contains_arguments, contains_this_expr, find_ids, ident::IdentLike, prepend, var::VarCollector,
+    ExprFactory, Id, StmtLike,
 };
 use swc_ecma_visit::noop_visit_type;
 use swc_ecma_visit::{
@@ -147,6 +147,19 @@ impl BlockScoping {
                     None
                 };
 
+                let arguments = if contains_arguments(&body) {
+                    let ident = private_ident!("_arguments");
+                    self.vars.push(VarDeclarator {
+                        span: DUMMY_SP,
+                        name: Pat::Ident(ident.clone().into()),
+                        init: Some(Box::new(Expr::Ident(quote_ident!("arguments")))),
+                        definite: false,
+                    });
+                    Some(ident)
+                } else {
+                    None
+                };
+
                 let mut flow_helper = FlowHelper {
                     all: &args,
                     has_continue: false,
@@ -164,12 +177,13 @@ impl BlockScoping {
                     },
                 };
 
-                if !flow_helper.mutated.is_empty() || this.is_some() {
+                if !flow_helper.mutated.is_empty() || this.is_some() || arguments.is_some() {
                     let no_modification = flow_helper.mutated.is_empty();
                     let mut v = MutationHandler {
                         map: &mut flow_helper.mutated,
                         in_function: false,
                         this,
+                        arguments,
                     };
 
                     // Modifies identifiers, and add reassignments to break / continue / return
@@ -868,6 +882,7 @@ struct MutationHandler<'a> {
     map: &'a mut HashMap<Id, SyntaxContext>,
     in_function: bool,
     this: Option<Ident>,
+    arguments: Option<Ident>,
 }
 
 impl MutationHandler<'_> {
@@ -921,6 +936,11 @@ impl VisitMut for MutationHandler<'_> {
                     ))
                 }
             }
+            Expr::Ident(id) if id.sym == js_word!("arguments") => {
+                if let Some(arguments) = &self.arguments {
+                    *n = Expr::Ident(arguments.clone())
+                }
+            }
             _ => {}
         }
     }
@@ -936,14 +956,17 @@ impl VisitMut for MutationHandler<'_> {
 
     fn visit_mut_function(&mut self, n: &mut Function) {
         let old = self.in_function;
+        let arguments = self.arguments.take();
         self.in_function = true;
 
         n.visit_mut_children_with(self);
 
         self.in_function = old;
+        self.arguments = arguments;
     }
 
     fn visit_mut_return_stmt(&mut self, n: &mut ReturnStmt) {
+        n.visit_mut_children_with(self);
         if self.in_function || self.map.is_empty() {
             return;
         }
@@ -952,6 +975,15 @@ impl VisitMut for MutationHandler<'_> {
 
         n.arg = Some(Box::new(self.make_reassignment(val)))
     }
+
+    /// Don't recurse into member expression prop if not computed
+    fn visit_mut_member_expr(&mut self, m: &mut MemberExpr) {
+        m.obj.visit_mut_with(self);
+        match &*m.prop {
+            Expr::Ident(_) if !m.computed => {}
+            _ => m.prop.visit_mut_with(self),
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -1288,7 +1320,7 @@ expect(foo()).toBe(false);
             console.log(i1++, [
                 2
             ].every(function(x) {
-                return x != i;
+                return x != i1;
             }));
             i = i1, void 0;
         };
@@ -1319,7 +1351,7 @@ expect(foo()).toBe(false);
             console.log(i1++, [
                 2
             ].every(function(x) {
-                return x != i;
+                return x != i1;
             }));
             if (i1 % 2 === 0) return i = i1, "continue";
             i = i1, void 0;
@@ -1354,7 +1386,7 @@ expect(foo()).toBe(false);
             console.log(i1++, [
                 2
             ].every(function(x) {
-                return x != i;
+                return x != i1;
             }));
             if (i1 % 2 === 0) return i = i1, "break";
             i = i1, void 0;
@@ -1669,4 +1701,132 @@ expect(foo()).toBe(false);
         }
         "
     );
+
+    test!(
+        ::swc_ecma_parser::Syntax::default(),
+        |_| {
+            let mark = Mark::fresh(Mark::root());
+            es2015::es2015(
+                mark,
+                es2015::Config {
+                    ..Default::default()
+                },
+            )
+        },
+        arguments_loop,
+        "
+        function test() {
+            for (var i = 0; i < arguments.length; i++) {
+              var arg = arguments[i];
+              console.log((() => arg)());
+            }
+        }
+        ",
+        "
+        function test() {
+            var _arguments = arguments, _loop = function(i) {
+                var arg = _arguments[i];
+                console.log(function() {
+                    return arg;
+                }());
+            };
+            for(var i = 0; i < arguments.length; i++)_loop(i);
+        }
+        "
+    );
+
+    test!(
+        ::swc_ecma_parser::Syntax::default(),
+        |_| {
+            let mark = Mark::fresh(Mark::root());
+            es2015::es2015(
+                mark,
+                es2015::Config {
+                    ..Default::default()
+                },
+            )
+        },
+        arguments_loop_member,
+        "
+        function test(a) {
+            for (var i = 0; i < a.arguments.length; i++) {
+              var arg = a.arguments[i];
+              console.log((() => arg)());
+            }
+        }
+        ",
+        "
+        function test(a) {
+            var _loop = function(i) {
+                var arg = a.arguments[i];
+                console.log(function() {
+                    return arg;
+                }());
+            };
+            for(var i = 0; i < a.arguments.length; i++)_loop(i);
+        }
+        "
+    );
+
+    test!(
+        ::swc_ecma_parser::Syntax::default(),
+        |_| {
+            let mark = Mark::fresh(Mark::root());
+            es2015::es2015(
+                mark,
+                es2015::Config {
+                    ..Default::default()
+                },
+            )
+        },
+        arguments_arrow,
+        "
+        function test() {
+            for (var i = 0; i < arguments.length; i++) {
+              console.log((() => arguments[i])());
+            }
+        }
+        ",
+        "
+        function test() {
+            var _this = this, _arguments = arguments, _loop = function(i) {
+                console.log((function(_arguments1) {
+                    return _arguments1[i];
+                }).bind(_this, _arguments)());
+            };
+            for(var i = 0; i < arguments.length; i++)_loop(i);
+        }
+        "
+    );
+
+    test!(
+        ::swc_ecma_parser::Syntax::default(),
+        |_| {
+            let mark = Mark::fresh(Mark::root());
+            es2015::es2015(
+                mark,
+                es2015::Config {
+                    ..Default::default()
+                },
+            )
+        },
+        arguments_function,
+        "
+        function test() {
+            for (var i = 0; i < arguments.length; i++) {
+              console.log((function () { return arguments[i] })());
+            }
+        }
+        ",
+        "
+        function test() {
+            var _loop = function(i) {
+                console.log(function() {
+                    return arguments[i];
+                }());
+            };
+            for(var i = 0; i < arguments.length; i++)_loop(i);
+        }
+        "
+    );
 }

ecmascript/transforms/compat/tests/es2015_arrow.rs

@@ -158,3 +158,113 @@ export const getBadgeBorderRadius = function(text, color) {
 };
 "#
 );
+
+test!(
+    ::swc_ecma_parser::Syntax::default(),
+    |_| arrow(),
+    arguments,
+    r#"
+function test() {
+  return () => arguments[0];
+}"#,
+    r#"
+  function test() {
+    return (function(_arguments) {
+      return _arguments[0];
+    }).bind(this, arguments)
+  }"#
+);
+
+test!(
+    ::swc_ecma_parser::Syntax::default(),
+    |_| arrow(),
+    arguments_nested_arrow,
+    r#"
+  function test() {
+    console.log(arguments[0]);
+    return () => {
+      console.log(arguments[0]);
+      return () => {
+        console.log(arguments[0])
+      };
+    }
+  }"#,
+    r#"
+  function test() {
+    console.log(arguments[0]);
+    return (function(_arguments) {
+        console.log(_arguments[0]);
+        return function() {
+            console.log(_arguments[0]);
+        };
+    }).bind(this, arguments);
+  }"#
+);
+
+test!(
+    ::swc_ecma_parser::Syntax::default(),
+    |_| arrow(),
+    arguments_nested_fn,
+    r#"
+  function test() {
+    console.log(arguments[0]);
+    return () => {
+      console.log(arguments[0]);
+      return function() {
+        console.log(arguments[0]);
+        return () => {
+          console.log(arguments[0])
+        };
+      };
+    }
+  }"#,
+    r#"
+  function test() {
+    console.log(arguments[0]);
+    return (function(_arguments) {
+        console.log(_arguments[0]);
+        return function() {
+            console.log(arguments[0]);
+            return (function(_arguments1) {
+                console.log(_arguments1[0]);
+            }).bind(this, arguments);
+        };
+    }).bind(this, arguments);
+  }"#
+);
+
+test!(
+    ::swc_ecma_parser::Syntax::default(),
+    |_| arrow(),
+    arguments_member,
+    r#"
+  function test() {
+    return (foo) => {
+      return foo.arguments;
+    }
+  }"#,
+    r#"
+  function test() {
+    return function(foo) {
+      return foo.arguments;
+    };
+  }"#
+);
+
+test!(
+    ::swc_ecma_parser::Syntax::default(),
+    |_| arrow(),
+    arguments_fn_expr,
+    r#"
+  function test() {
+    return function() {
+      return arguments[0];
+    };
+  }"#,
+    r#"
+  function test() {
+    return function() {
+      return arguments[0];
+    };
+  }"#
+);

ecmascript/transforms/compat/tests/es2017_async_to_generator.rs

@@ -229,15 +229,15 @@ function _s() {
         for(let _len = arguments.length, args = new Array(_len > 1 ? _len - 1 : 0), _key = 1; _key < _len; _key++){
             args[_key - 1] = arguments[_key];
         }
-        let t = (function(y, a) {
-            var _ref = _asyncToGenerator((function*(y, a) {
+        let t = (function(_arguments, y, a) {
+            var _ref = _asyncToGenerator((function*(_arguments, y, a) {
                 let r = (function(z, b) {
                     var _ref1 = _asyncToGenerator((function*(z, b) {
                         for(let _len1 = arguments.length, innerArgs = new Array(_len1 > 2 ? _len1 - 2 : 0), _key1 = 2; _key1 < _len1; _key1++){
                             innerArgs[_key1 - 2] = arguments[_key1];
                         }
                         yield z;
-                        console.log(this, innerArgs, arguments);
+                        console.log(this, innerArgs, _arguments);
                         return this.x;
                     }).bind(this));
                     return function() {
@@ -245,13 +245,13 @@ function _s() {
                     };
                 })().bind(this);
                 yield r();
-                console.log(this, args, arguments);
+                console.log(this, args, _arguments);
                 return this.g(r);
             }).bind(this));
             return function () {
                 return _ref.apply(this, arguments);
             };
-        })().bind(this);
+        })().bind(this, arguments);
         yield t();
         return this.h(t);
     }).bind(this));

ecmascript/utils/src/lib.rs

@@ -99,6 +99,66 @@ impl Visit for IdentFinder<'_> {
     }
 }
 
+pub fn contains_arguments<N>(body: &N) -> bool
+where
+    N: VisitWith<ArgumentsFinder>,
+{
+    let mut visitor = ArgumentsFinder { found: false };
+    body.visit_with(&Invalid { span: DUMMY_SP } as _, &mut visitor);
+    visitor.found
+}
+
+pub struct ArgumentsFinder {
+    found: bool,
+}
+
+impl Visit for ArgumentsFinder {
+    noop_visit_type!();
+
+    /// Don't recurse into constructor
+    fn visit_constructor(&mut self, _: &Constructor, _: &dyn Node) {}
+
+    fn visit_expr(&mut self, e: &Expr, _: &dyn Node) {
+        e.visit_children_with(self);
+
+        match *e {
+            Expr::Ident(Ident {
+                sym: js_word!("arguments"),
+                ..
+            }) => {
+                self.found = true;
+            }
+            _ => {}
+        }
+    }
+
+    /// Don't recurse into fn
+    fn visit_function(&mut self, _: &Function, _: &dyn Node) {}
+
+    /// Don't recurse into member expression prop if not computed
+    fn visit_member_expr(&mut self, m: &MemberExpr, _: &dyn Node) {
+        m.obj.visit_with(m, self);
+        match &*m.prop {
+            Expr::Ident(_) if !m.computed => {}
+            _ => m.prop.visit_with(m, self),
+        }
+    }
+
+    fn visit_prop(&mut self, n: &Prop, _: &dyn Node) {
+        n.visit_children_with(self);
+
+        match n {
+            Prop::Shorthand(Ident {
+                sym: js_word!("arguments"),
+                ..
+            }) => {
+                self.found = true;
+            }
+            _ => {}
+        }
+    }
+}
+
 pub trait ModuleItemLike: StmtLike {
     fn try_into_module_decl(self) -> Result<ModuleDecl, Self> {
         Err(self)