support async function (#4333)

6 files changed, 158 insertions, 47 deletions

diff --git a/lib/ast.js b/lib/ast.js
index 792264ff..ecaecc5b 100644
--- a/lib/ast.js
+++ b/lib/ast.js
@@ -211,7 +211,7 @@ function must_be_expression(node, prop) {
     if (!(node[prop] instanceof AST_Node)) throw new Error(prop + " must be AST_Node");
     if (node[prop] instanceof AST_Hole) throw new Error(prop + " cannot be AST_Hole");
     if (node[prop] instanceof AST_Spread) throw new Error(prop + " cannot be AST_Spread");
-    if (node[prop] instanceof AST_Statement && !(node[prop] instanceof AST_Function)) {
+    if (node[prop] instanceof AST_Statement && !is_function(node[prop])) {
         throw new Error(prop + " cannot be AST_Statement");
     }
 }
@@ -280,7 +280,7 @@ var AST_Block = DEFNODE("Block", "body", {
     _validate: function() {
         this.body.forEach(function(node) {
             if (!(node instanceof AST_Statement)) throw new Error("body must be AST_Statement[]");
-            if (node instanceof AST_Function) throw new Error("body cannot contain AST_Function");
+            if (is_function(node)) throw new Error("body cannot contain AST_Function");
         });
     },
 }, AST_BlockScope);
@@ -296,7 +296,7 @@ var AST_StatementWithBody = DEFNODE("StatementWithBody", "body", {
     },
     _validate: function() {
         if (!(this.body instanceof AST_Statement)) throw new Error("body must be AST_Statement");
-        if (this.body instanceof AST_Function) throw new Error("body cannot be AST_Function");
+        if (is_function(this.body)) throw new Error("body cannot be AST_Function");
     },
 }, AST_BlockScope);
 
@@ -390,7 +390,7 @@ var AST_For = DEFNODE("For", "init condition step", {
         if (this.init != null) {
             if (!(this.init instanceof AST_Node)) throw new Error("init must be AST_Node");
             if (this.init instanceof AST_Statement
-                && !(this.init instanceof AST_Definitions || this.init instanceof AST_Function)) {
+                && !(this.init instanceof AST_Definitions || is_function(this.init))) {
                 throw new Error("init cannot be AST_Statement");
             }
         }
@@ -547,6 +547,19 @@ var AST_Accessor = DEFNODE("Accessor", null, {
     },
 }, AST_Lambda);
 
+function is_function(node) {
+    return node instanceof AST_AsyncFunction || node instanceof AST_Function;
+}
+
+var AST_AsyncFunction = DEFNODE("AsyncFunction", null, {
+    $documentation: "An asynchronous function expression",
+    _validate: function() {
+        if (this.name != null) {
+            if (!(this.name instanceof AST_SymbolLambda)) throw new Error("name must be AST_SymbolLambda");
+        }
+    },
+}, AST_Lambda);
+
 var AST_Function = DEFNODE("Function", "inlined", {
     $documentation: "A function expression",
     _validate: function() {
@@ -556,6 +569,13 @@ var AST_Function = DEFNODE("Function", "inlined", {
     },
 }, AST_Lambda);
 
+var AST_AsyncDefun = DEFNODE("AsyncDefun", null, {
+    $documentation: "An asynchronous function definition",
+    _validate: function() {
+        if (!(this.name instanceof AST_SymbolDefun)) throw new Error("name must be AST_SymbolDefun");
+    },
+}, AST_Lambda);
+
 var AST_Defun = DEFNODE("Defun", "inlined", {
     $documentation: "A function definition",
     _validate: function() {
@@ -642,7 +662,7 @@ var AST_If = DEFNODE("If", "condition alternative", {
         must_be_expression(this, "condition");
         if (this.alternative != null) {
             if (!(this.alternative instanceof AST_Statement)) throw new Error("alternative must be AST_Statement");
-            if (this.alternative instanceof AST_Function) throw new error("alternative cannot be AST_Function");
+            if (is_function(this.alternative)) throw new error("alternative cannot be AST_Function");
         }
     },
 }, AST_StatementWithBody);
@@ -824,7 +844,7 @@ function must_be_expressions(node, prop, allow_spread, allow_hole) {
         if (!(node instanceof AST_Node)) throw new Error(prop + " must be AST_Node[]");
         if (!allow_hole && node instanceof AST_Hole) throw new Error(prop + " cannot be AST_Hole");
         if (!allow_spread && node instanceof AST_Spread) throw new Error(prop + " cannot be AST_Spread");
-        if (node instanceof AST_Statement && !(node instanceof AST_Function)) {
+        if (node instanceof AST_Statement && !is_function(node)) {
             throw new Error(prop + " cannot contain AST_Statement");
         }
     });
@@ -1024,6 +1044,22 @@ var AST_Assign = DEFNODE("Assign", null, {
     },
 }, AST_Binary);
 
+var AST_Await = DEFNODE("Await", "expression", {
+    $documentation: "An await expression",
+    $propdoc: {
+        expression: "[AST_Node] expression with Promise to resolve on",
+    },
+    walk: function(visitor) {
+        var node = this;
+        visitor.visit(node, function() {
+            node.expression.walk(visitor);
+        });
+    },
+    _validate: function() {
+        must_be_expression(this, "expression");
+    },
+});
+
 /* -----[ LITERALS ]----- */
 
 var AST_Array = DEFNODE("Array", "elements", {
diff --git a/lib/compress.js b/lib/compress.js
index e84465f6..cf138377 100644
--- a/lib/compress.js
+++ b/lib/compress.js
@@ -322,7 +322,7 @@ merge(Compressor.prototype, {
         if (value instanceof AST_String) return native_fns.String[name];
         if (name == "valueOf") return false;
         if (value instanceof AST_Array) return native_fns.Array[name];
-        if (value instanceof AST_Function) return native_fns.Function[name];
+        if (value instanceof AST_Lambda) return native_fns.Function[name];
         if (value instanceof AST_Object) return native_fns.Object[name];
         if (value instanceof AST_RegExp) return native_fns.RegExp[name] && !value.value.global;
     }
@@ -650,14 +650,6 @@ merge(Compressor.prototype, {
             lhs.walk(scanner);
         }
 
-        def(AST_Accessor, function(tw, descend, compressor) {
-            push(tw);
-            reset_variables(tw, compressor, this);
-            descend();
-            pop(tw);
-            walk_defuns(tw, this);
-            return true;
-        });
         def(AST_Assign, function(tw, descend, compressor) {
             var node = this;
             var left = node.left;
@@ -935,6 +927,14 @@ merge(Compressor.prototype, {
             pop(tw);
             return true;
         });
+        def(AST_Lambda, function(tw, descend, compressor) {
+            push(tw);
+            reset_variables(tw, compressor, this);
+            descend();
+            pop(tw);
+            walk_defuns(tw, this);
+            return true;
+        });
         def(AST_Switch, function(tw, descend, compressor) {
             this.variables.each(function(def) {
                 reset_def(tw, compressor, def);
@@ -1365,7 +1365,7 @@ merge(Compressor.prototype, {
 
     function is_iife_call(node) {
         if (node.TYPE != "Call") return false;
-        return node.expression instanceof AST_Function || is_iife_call(node.expression);
+        return is_function(node.expression) || is_iife_call(node.expression);
     }
 
     function is_undeclared_ref(node) {
@@ -1744,6 +1744,7 @@ merge(Compressor.prototype, {
             }
 
             function is_last_node(node, parent) {
+                if (node instanceof AST_Await) return true;
                 if (node.TYPE == "Binary") return node.operator == "in" && !is_object(node.right);
                 if (node instanceof AST_Call) {
                     var def, fn = node.expression;
@@ -1887,6 +1888,8 @@ merge(Compressor.prototype, {
                     if (expr.left instanceof AST_SymbolRef) {
                         assignments[expr.left.name] = (assignments[expr.left.name] || 0) + 1;
                     }
+                } else if (expr instanceof AST_Await) {
+                    extract_candidates(expr.expression);
                 } else if (expr instanceof AST_Binary) {
                     extract_candidates(expr.left);
                     extract_candidates(expr.right);
@@ -1975,6 +1978,7 @@ merge(Compressor.prototype, {
                 var parent = scanner.parent(level);
                 if (parent instanceof AST_Array) return node;
                 if (parent instanceof AST_Assign) return node;
+                if (parent instanceof AST_Await) return node;
                 if (parent instanceof AST_Binary) return node;
                 if (parent instanceof AST_Call) return node;
                 if (parent instanceof AST_Case) return node;
@@ -2074,6 +2078,7 @@ merge(Compressor.prototype, {
                 if (parent instanceof AST_Assign) {
                     return may_throw(parent) ? node : find_stop_unused(parent, level + 1);
                 }
+                if (parent instanceof AST_Await) return node;
                 if (parent instanceof AST_Binary) return find_stop_unused(parent, level + 1);
                 if (parent instanceof AST_Call) return find_stop_unused(parent, level + 1);
                 if (parent instanceof AST_Case) return find_stop_unused(parent, level + 1);
@@ -4015,7 +4020,7 @@ merge(Compressor.prototype, {
         def(AST_Call, function(compressor, ignore_side_effects, cached, depth) {
             var exp = this.expression;
             var fn = exp instanceof AST_SymbolRef ? exp.fixed_value() : exp;
-            if (fn instanceof AST_Lambda) {
+            if (fn instanceof AST_Defun || fn instanceof AST_Function) {
                 if (fn.evaluating) return this;
                 if (fn.name && fn.name.definition().recursive_refs > 0) return this;
                 if (this.is_expr_pure(compressor)) return this;
@@ -4142,6 +4147,9 @@ merge(Compressor.prototype, {
         def(AST_Statement, function() {
             throw new Error("Cannot negate a statement");
         });
+        def(AST_AsyncFunction, function() {
+            return basic_negation(this);
+        });
         def(AST_Function, function() {
             return basic_negation(this);
         });
@@ -5107,7 +5115,7 @@ merge(Compressor.prototype, {
             }
             if (node === self) return;
             if (scope === self) {
-                if (node instanceof AST_Defun) {
+                if (node instanceof AST_AsyncDefun || node instanceof AST_Defun) {
                     var def = node.name.definition();
                     if (!drop_funcs && !(def.id in in_use_ids)) {
                         in_use_ids[def.id] = true;
@@ -5258,7 +5266,7 @@ merge(Compressor.prototype, {
             if (node instanceof AST_Call) calls_to_drop_args.push(node);
             if (scope !== self) return;
             if (node instanceof AST_Lambda) {
-                if (drop_funcs && node !== self && node instanceof AST_Defun) {
+                if (drop_funcs && node !== self && (node instanceof AST_AsyncDefun || node instanceof AST_Defun)) {
                     var def = node.name.definition();
                     if (!(def.id in in_use_ids)) {
                         log(node.name, "Dropping unused function {name}");
@@ -5266,7 +5274,7 @@ merge(Compressor.prototype, {
                         return in_list ? List.skip : make_node(AST_EmptyStatement, node);
                     }
                 }
-                if (node instanceof AST_Function && node.name && drop_fn_name(node.name.definition())) {
+                if (is_function(node) && node.name && drop_fn_name(node.name.definition())) {
                     unused_fn_names.push(node);
                 }
                 if (!(node instanceof AST_Accessor)) {
@@ -7642,7 +7650,7 @@ merge(Compressor.prototype, {
             }
         }
         var fn = exp instanceof AST_SymbolRef ? exp.fixed_value() : exp;
-        var is_func = fn instanceof AST_Lambda;
+        var is_func = fn instanceof AST_Defun || fn instanceof AST_Function;
         var stat = is_func && fn.first_statement();
         var can_inline = is_func
             && compressor.option("inline")
@@ -8967,7 +8975,10 @@ merge(Compressor.prototype, {
                 def.single_use = false;
                 fixed._squeezed = true;
                 fixed.single_use = true;
-                if (fixed instanceof AST_Defun) {
+                if (fixed instanceof AST_AsyncDefun) {
+                    fixed = make_node(AST_AsyncFunction, fixed, fixed);
+                    fixed.name = make_node(AST_SymbolLambda, fixed.name, fixed.name);
+                } else if (fixed instanceof AST_Defun) {
                     fixed = make_node(AST_Function, fixed, fixed);
                     fixed.name = make_node(AST_SymbolLambda, fixed.name, fixed.name);
                 }
diff --git a/lib/output.js b/lib/output.js
index 5830f2e7..bf243862 100644
--- a/lib/output.js
+++ b/lib/output.js
@@ -661,7 +661,7 @@ function OutputStream(options) {
 
     // a function expression needs parens around it when it's provably
     // the first token to appear in a statement.
-    PARENS(AST_Function, function(output) {
+    function needs_parens_function(output) {
         if (!output.has_parens() && first_in_statement(output)) return true;
         if (output.option("webkit")) {
             var p = output.parent();
@@ -671,7 +671,9 @@ function OutputStream(options) {
             var p = output.parent();
             if (p instanceof AST_Call && p.expression === this) return true;
         }
-    });
+    }
+    PARENS(AST_AsyncFunction, needs_parens_function);
+    PARENS(AST_Function, needs_parens_function);
 
     // same goes for an object literal, because otherwise it would be
     // interpreted as a block of code.
@@ -689,6 +691,8 @@ function OutputStream(options) {
         var p = output.parent();
             // [ 1, (2, 3), 4 ] ==> [ 1, 3, 4 ]
         return p instanceof AST_Array
+            // await (foo, bar)
+            || p instanceof AST_Await
             // 1 + (2, 3) + 4 ==> 8
             || p instanceof AST_Binary
             // new (foo, bar) or foo(1, (2, 3), 4)
@@ -712,6 +716,8 @@ function OutputStream(options) {
 
     PARENS(AST_Binary, function(output) {
         var p = output.parent();
+        // await (foo && bar)
+        if (p instanceof AST_Await) return true;
         // this deals with precedence: 3 * (2 + 1)
         if (p instanceof AST_Binary) {
             var po = p.operator, pp = PRECEDENCE[po];
@@ -779,6 +785,8 @@ function OutputStream(options) {
 
     function needs_parens_assign_cond(self, output) {
         var p = output.parent();
+        // await (a = foo)
+        if (p instanceof AST_Await) return true;
         // 1 + (a = 2) + 3 → 6, side effect setting a = 2
         if (p instanceof AST_Binary) return !(p instanceof AST_Assign);
         // (a = func)() —or— new (a = Object)()
@@ -967,11 +975,7 @@ function OutputStream(options) {
     });
 
     /* -----[ functions ]----- */
-    DEFPRINT(AST_Lambda, function(output, nokeyword) {
-        var self = this;
-        if (!nokeyword) {
-            output.print("function");
-        }
+    function print_lambda(self, output) {
         if (self.name) {
             output.space();
             self.name.print(output);
@@ -984,7 +988,19 @@ function OutputStream(options) {
         });
         output.space();
         print_braced(self, output, true);
+    }
+    DEFPRINT(AST_Lambda, function(output) {
+        output.print("function");
+        print_lambda(this, output);
     });
+    function print_async(output) {
+        output.print("async");
+        output.space();
+        output.print("function");
+        print_lambda(this, output);
+    }
+    DEFPRINT(AST_AsyncDefun, print_async);
+    DEFPRINT(AST_AsyncFunction, print_async);
 
     /* -----[ jumps ]----- */
     function print_jump(kind, prop) {
@@ -1272,6 +1288,11 @@ function OutputStream(options) {
         output.colon();
         self.alternative.print(output);
     });
+    DEFPRINT(AST_Await, function(output) {
+        output.print("await");
+        output.space();
+        this.expression.print(output);
+    });
 
     /* -----[ literals ]----- */
     DEFPRINT(AST_Array, function(output) {
@@ -1375,7 +1396,7 @@ function OutputStream(options) {
             output.print(type);
             output.space();
             print_property_key(self, output);
-            self.value._codegen(output, true);
+            print_lambda(self.value, output);
         };
     }
     DEFPRINT(AST_ObjectGetter, print_accessor("get"));
diff --git a/lib/parse.js b/lib/parse.js
index 872561ff..b3c3f977 100644
--- a/lib/parse.js
+++ b/lib/parse.js
@@ -47,7 +47,7 @@
 var KEYWORDS = "break case catch const continue debugger default delete do else finally for function if in instanceof let new return switch throw try typeof var void while with";
 var KEYWORDS_ATOM = "false null true";
 var RESERVED_WORDS = [
-    "abstract boolean byte char class double enum export extends final float goto implements import int interface let long native package private protected public short static super synchronized this throws transient volatile yield",
+    "await abstract boolean byte char class double enum export extends final float goto implements import int interface long native package private protected public short static super synchronized this throws transient volatile yield",
     KEYWORDS_ATOM,
     KEYWORDS,
 ].join(" ");
@@ -656,6 +656,7 @@ function parse($TEXT, options) {
         token         : null,
         prev          : null,
         peeked        : null,
+        in_async      : false,
         in_function   : 0,
         in_directives : true,
         in_loop       : 0,
@@ -786,9 +787,20 @@ function parse($TEXT, options) {
             return simple_statement();
 
           case "name":
-            return is_token(peek(), "punc", ":")
-                ? labeled_statement()
-                : simple_statement();
+            switch (S.token.value) {
+              case "async":
+                if (is_token(peek(), "keyword", "function")) {
+                    next();
+                    next();
+                    return function_(AST_AsyncDefun);
+                }
+              case "await":
+                if (S.in_async) return simple_statement();
+              default:
+                return is_token(peek(), "punc", ":")
+                    ? labeled_statement()
+                    : simple_statement();
+            }
 
           case "punc":
             switch (S.token.value) {
@@ -1026,10 +1038,18 @@ function parse($TEXT, options) {
     }
 
     var function_ = function(ctor) {
-        var in_statement = ctor === AST_Defun;
-        var name = is("name") ? as_symbol(in_statement ? AST_SymbolDefun : AST_SymbolLambda) : null;
-        if (in_statement && !name)
-            expect_token("name");
+        var was_async = S.in_async;
+        var name;
+        if (ctor === AST_AsyncDefun) {
+            name = as_symbol(AST_SymbolDefun);
+            S.in_async = true;
+        } else if (ctor === AST_Defun) {
+            name = as_symbol(AST_SymbolDefun);
+            S.in_async = false;
+        } else {
+            S.in_async = ctor === AST_AsyncFunction;
+            name = as_symbol(AST_SymbolLambda, true);
+        }
         if (name && ctor !== AST_Accessor && !(name instanceof AST_SymbolDeclaration))
             unexpected(prev());
         expect("(");
@@ -1055,6 +1075,7 @@ function parse($TEXT, options) {
         --S.in_function;
         S.in_loop = loop;
         S.labels = labels;
+        S.in_async = was_async;
         return new ctor({
             name: name,
             argnames: argnames,
@@ -1304,9 +1325,16 @@ function parse($TEXT, options) {
             }
             unexpected();
         }
-        if (is("keyword", "function")) {
+        var ctor;
+        if (is("name", "async") && is_token(peek(), "keyword", "function")) {
             next();
-            var func = function_(AST_Function);
+            ctor = AST_AsyncFunction;
+        } else if (is("keyword", "function")) {
+            ctor = AST_Function;
+        }
+        if (ctor) {
+            next();
+            var func = function_(ctor);
             func.start = start;
             func.end = prev();
             return subscripts(func, allow_calls);
@@ -1451,6 +1479,7 @@ function parse($TEXT, options) {
 
     function _make_symbol(type, token) {
         var name = token.value;
+        if (name === "await" && S.in_async) unexpected(token);
         return new (name === "this" ? AST_This : type)({
             name: "" + name,
             start: token,
@@ -1573,17 +1602,17 @@ function parse($TEXT, options) {
         return expr;
     };
 
-    var maybe_unary = function(allow_calls) {
+    function maybe_unary() {
         var start = S.token;
         if (is("operator") && UNARY_PREFIX[start.value]) {
             next();
             handle_regexp();
-            var ex = make_unary(AST_UnaryPrefix, start, maybe_unary(allow_calls));
+            var ex = make_unary(AST_UnaryPrefix, start, maybe_await());
             ex.start = start;
             ex.end = prev();
             return ex;
         }
-        var val = expr_atom(allow_calls);
+        var val = expr_atom(true);
         while (is("operator") && UNARY_POSTFIX[S.token.value] && !has_newline_before(S.token)) {
             val = make_unary(AST_UnaryPostfix, S.token, val);
             val.start = start;
@@ -1591,7 +1620,7 @@ function parse($TEXT, options) {
             next();
         }
         return val;
-    };
+    }
 
     function make_unary(ctor, token, expr) {
         var op = token.value;
@@ -1609,13 +1638,24 @@ function parse($TEXT, options) {
         return new ctor({ operator: op, expression: expr });
     }
 
+    function maybe_await() {
+        var start = S.token;
+        if (!(S.in_async && is("name", "await"))) return maybe_unary();
+        next();
+        return new AST_Await({
+            start: start,
+            expression: maybe_await(),
+            end: prev(),
+        });
+    }
+
     var expr_op = function(left, min_prec, no_in) {
         var op = is("operator") ? S.token.value : null;
         if (op == "in" && no_in) op = null;
         var prec = op != null ? PRECEDENCE[op] : null;
         if (prec != null && prec > min_prec) {
             next();
-            var right = expr_op(maybe_unary(true), prec, no_in);
+            var right = expr_op(maybe_await(), prec, no_in);
             return expr_op(new AST_Binary({
                 start    : left.start,
                 left     : left,
@@ -1628,7 +1668,7 @@ function parse($TEXT, options) {
     };
 
     function expr_ops(no_in) {
-        return expr_op(maybe_unary(true), 0, no_in);
+        return expr_op(maybe_await(), 0, no_in);
     }
 
     var maybe_conditional = function(no_in) {
diff --git a/lib/scope.js b/lib/scope.js
index 4d080fee..c1bbc086 100644
--- a/lib/scope.js
+++ b/lib/scope.js
@@ -112,7 +112,7 @@ AST_Toplevel.DEFMETHOD("figure_out_scope", function(options) {
     var next_def_id = 0;
     var scope = self.parent_scope = null;
     var tw = new TreeWalker(function(node, descend) {
-        if (node instanceof AST_Defun) {
+        if (node instanceof AST_AsyncDefun || node instanceof AST_Defun) {
             node.name.walk(tw);
             walk_scope(function() {
                 node.argnames.forEach(function(argname) {
diff --git a/lib/transform.js b/lib/transform.js
index 65e6cd60..398cf73d 100644
--- a/lib/transform.js
+++ b/lib/transform.js
@@ -138,6 +138,9 @@ TreeTransformer.prototype = new TreeWalker;
     DEF(AST_Sequence, function(self, tw) {
         self.expressions = do_list(self.expressions, tw);
     });
+    DEF(AST_Await, function(self, tw) {
+        self.expression = self.expression.transform(tw);
+    });
     DEF(AST_Dot, function(self, tw) {
         self.expression = self.expression.transform(tw);
     });