enhance `awaits` (#4971)

3 files changed, 113 insertions, 35 deletions

diff --git a/lib/compress.js b/lib/compress.js
index b07e85ad..4c85c51b 100644
--- a/lib/compress.js
+++ b/lib/compress.js
@@ -7519,11 +7519,13 @@ merge(Compressor.prototype, {
             return make_sequence(this, [ lhs, rhs ]);
         });
         function drop_returns(compressor, exp) {
+            var arrow = is_arrow(exp);
+            var async = is_async(exp);
             var drop_body = false;
-            if (compressor.option("arrows") && is_arrow(exp)) {
+            if (arrow && compressor.option("arrows")) {
                 if (!exp.value) {
                     drop_body = true;
-                } else if (!is_async(exp) || is_primitive(compressor, exp.value)) {
+                } else if (!async || is_primitive(compressor, exp.value)) {
                     exp.value = exp.value.drop_side_effect_free(compressor);
                 }
             } else if (exp instanceof AST_AsyncFunction || exp instanceof AST_Function) {
@@ -7535,7 +7537,6 @@ merge(Compressor.prototype, {
                 }
             }
             if (drop_body) {
-                var async = is_async(exp);
                 exp.process_expression(false, function(node) {
                     var value = node.value;
                     if (value) {
@@ -7552,43 +7553,55 @@ merge(Compressor.prototype, {
                         node.value = value.drop_side_effect_free(compressor);
                     }
                 });
-                if (async && compressor.option("awaits")) {
-                    exp.process_expression(true, function(node) {
-                        var body = node.body;
-                        if (body instanceof AST_Await) {
-                            if (is_primitive(compressor, body.expression)) {
-                                body = body.expression.drop_side_effect_free(compressor, true);
-                                if (!body) return make_node(AST_EmptyStatement, node);
-                                node.body = body;
-                            }
-                        } else if (body instanceof AST_Sequence) {
-                            var tail = body.tail_node();
-                            if (tail instanceof AST_Await && is_primitive(compressor, tail.expression)) {
-                                tail = tail.expression.drop_side_effect_free(compressor);
-                                body.expressions.pop();
-                                if (tail) body.expressions.push(tail);
-                            }
+            }
+            if (async && compressor.option("awaits")) {
+                if (drop_body) exp.process_expression(true, function(node) {
+                    var body = node.body;
+                    if (body instanceof AST_Await) {
+                        if (is_primitive(compressor, body.expression)) {
+                            body = body.expression.drop_side_effect_free(compressor, true);
+                            if (!body) return make_node(AST_EmptyStatement, node);
+                            node.body = body;
                         }
-                        return node;
-                    });
-                    if (all(exp.body, is_empty) && !(is_arrow(exp) && exp.value)) {
-                        var ctor;
-                        switch (exp.CTOR) {
-                          case AST_AsyncArrow:
-                            ctor = AST_Arrow;
-                            break;
-                          case AST_AsyncFunction:
-                            ctor = AST_Function;
-                            break;
-                          case AST_AsyncGeneratorFunction:
-                            ctor = AST_GeneratorFunction;
-                            break;
+                    } else if (body instanceof AST_Sequence) {
+                        var tail = body.tail_node();
+                        if (tail instanceof AST_Await && is_primitive(compressor, tail.expression)) {
+                            tail = tail.expression.drop_side_effect_free(compressor);
+                            body.expressions.pop();
+                            if (tail) body.expressions.push(tail);
                         }
-                        return make_node(ctor, exp, exp);
                     }
+                    return node;
+                });
+                var abort = arrow && exp.value && !is_primitive(compressor, exp.value);
+                var tw = new TreeWalker(function(node) {
+                    if (abort) return true;
+                    if (node instanceof AST_Await) return abort = true;
+                    if (node instanceof AST_Return) {
+                        if (node.value && !is_primitive(compressor, node.value)) return abort = true;
+                        return;
+                    }
+                    if (node instanceof AST_Scope && node !== exp) return true;
+                    if (tw.parent() === exp && node.may_throw(compressor)) return abort = true;
+                });
+                exp.walk(tw);
+                if (!abort) {
+                    var ctor;
+                    switch (exp.CTOR) {
+                      case AST_AsyncArrow:
+                        ctor = AST_Arrow;
+                        break;
+                      case AST_AsyncFunction:
+                        ctor = AST_Function;
+                        break;
+                      case AST_AsyncGeneratorFunction:
+                        ctor = AST_GeneratorFunction;
+                        break;
+                    }
+                    return make_node(ctor, exp, exp);
                 }
-                return exp.clone();
             }
+            return drop_body && exp.clone();
         }
         def(AST_Call, function(compressor, first_in_statement) {
             var self = this;
diff --git a/test/compress/awaits.js b/test/compress/awaits.js
index f794e9cf..d5b60789 100644
--- a/test/compress/awaits.js
+++ b/test/compress/awaits.js
@@ -555,6 +555,51 @@ collapse_property_lambda: {
     node_version: ">=8"
 }
 
+drop_async_1: {
+    options = {
+        awaits: true,
+        inline: true,
+        side_effects: true,
+    }
+    input: {
+        console.log(function(a) {
+            (async function() {
+                a *= 7;
+            })();
+            return a;
+        }(6));
+    }
+    expect: {
+        console.log(function(a) {
+            void (a *= 7);
+            return a;
+        }(6));
+    }
+    expect_stdout: "42"
+    node_version: ">=8"
+}
+
+drop_async_2: {
+    options = {
+        awaits: true,
+        collapse_vars: true,
+        evaluate: true,
+        inline: true,
+        side_effects: true,
+    }
+    input: {
+        console.log(function(a) {
+            (async b => await (a *= b))(7);
+            return a;
+        }(6));
+    }
+    expect: {
+        console.log(42);
+    }
+    expect_stdout: "42"
+    node_version: ">=8"
+}
+
 drop_return: {
     options = {
         side_effects: true,
diff --git a/test/compress/functions.js b/test/compress/functions.js
index 5713a37b..8ab7c1fc 100644
--- a/test/compress/functions.js
+++ b/test/compress/functions.js
@@ -6208,3 +6208,23 @@ reduce_cross_reference_4_toplevel: {
     expect: {}
     expect_stdout: true
 }
+
+recursive_collapse: {
+    options = {
+        collapse_vars: true,
+        reduce_vars: true,
+    }
+    input: {
+        console.log(function f(a) {
+            var b = a && f();
+            return b;
+        }("FAIL") || "PASS");
+    }
+    expect: {
+        console.log(function f(a) {
+            var b;
+            return a && f();
+        }("FAIL") || "PASS");
+    }
+    expect_stdout: "PASS"
+}