Fix transpilation of nested async arrow functions.

Also use `const` instead of `let` and simplify the logic a bit. ------------- Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=158856990
google · Jun 13, 2017 · 42143c3 · 42143c3
1 parent ad26a8c
commit 42143c3
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 17 deletions.
diff --git a/src/com/google/javascript/jscomp/RewriteAsyncFunctions.java b/src/com/google/javascript/jscomp/RewriteAsyncFunctions.java
@@ -58,41 +58,38 @@ public final class RewriteAsyncFunctions implements NodeTraversal.Callback, HotS
   private static final class LexicalContext {
     final Optional<Node> function; // absent for top level
     final LexicalContext thisAndArgumentsContext;
-    final Optional<LexicalContext> enclosingAsyncContext;
     boolean asyncThisReplacementWasDone = false;
     boolean asyncArgumentsReplacementWasDone = false;
 
     /** Creates root-level context. */
     LexicalContext() {
       this.function = Optional.absent();
       this.thisAndArgumentsContext = this;
-      this.enclosingAsyncContext = Optional.absent();
     }
 
     LexicalContext(LexicalContext outer, Node function) {
       this.function = Optional.of(function);
       // An arrow function shares 'this' and 'arguments' with its outer scope.
       this.thisAndArgumentsContext =
           function.isArrowFunction() ? outer.thisAndArgumentsContext : this;
-      this.enclosingAsyncContext =
-          function.isAsyncFunction() ? Optional.of(this) : outer.enclosingAsyncContext;
     }
 
-    boolean isAsyncFunction() {
+    boolean isAsyncContext() {
       return function.isPresent() && function.get().isAsyncFunction();
     }
 
     boolean mustReplaceThisAndArguments() {
-      return enclosingAsyncContext.isPresent()
-          && enclosingAsyncContext.get().thisAndArgumentsContext == thisAndArgumentsContext;
+      return thisAndArgumentsContext.isAsyncContext();
     }
 
     void recordAsyncThisReplacementWasDone() {
-      enclosingAsyncContext.get().asyncThisReplacementWasDone = true;
+      checkState(thisAndArgumentsContext.isAsyncContext());
+      thisAndArgumentsContext.asyncThisReplacementWasDone = true;
     }
 
     void recordAsyncArgumentsReplacementWasDone() {
-      enclosingAsyncContext.get().asyncArgumentsReplacementWasDone = true;
+      checkState(thisAndArgumentsContext.isAsyncContext());
+      thisAndArgumentsContext.asyncArgumentsReplacementWasDone = true;
     }
   }
 
@@ -141,7 +138,7 @@ public void visit(NodeTraversal t, Node n, Node parent) {
     }
     switch (n.getToken()) {
       case FUNCTION:
-        if (context.isAsyncFunction()) {
+        if (context.isAsyncContext()) {
           convertAsyncFunction(context);
         }
         break;
@@ -161,7 +158,7 @@ public void visit(NodeTraversal t, Node n, Node parent) {
         break;
 
       case AWAIT:
-        checkState(context.isAsyncFunction(), "await found within non-async function body");
+        checkState(context.isAsyncContext(), "await found within non-async function body");
         checkState(n.hasOneChild(), "await should have 1 operand, but has %s", n.getChildCount());
         // Awaits become yields in the converted async function's inner generator function.
         parent.replaceChild(n, IR.yield(n.removeFirstChild()).useSourceInfoIfMissingFrom(n));
@@ -180,10 +177,10 @@ private void convertAsyncFunction(LexicalContext functionContext) {
     originalFunction.replaceChild(originalBody, newBody);
 
     if (functionContext.asyncThisReplacementWasDone) {
-      newBody.addChildToBack(IR.let(IR.name(ASYNC_THIS), IR.thisNode()));
+      newBody.addChildToBack(IR.constNode(IR.name(ASYNC_THIS), IR.thisNode()));
     }
     if (functionContext.asyncArgumentsReplacementWasDone) {
-      newBody.addChildToBack(IR.let(IR.name(ASYNC_ARGUMENTS), IR.name("arguments")));
+      newBody.addChildToBack(IR.constNode(IR.name(ASYNC_ARGUMENTS), IR.name("arguments")));
     }
 
     // Normalize arrow function short body to block body

diff --git a/test/com/google/javascript/jscomp/RewriteAsyncFunctionsTest.java b/test/com/google/javascript/jscomp/RewriteAsyncFunctionsTest.java
@@ -84,7 +84,7 @@ public void testArgumentsReplacement_asyncFunction() {
         "async function f(a, b, ...rest) { return arguments.length; }",
         LINE_JOINER.join(
             "function f(a, b, ...rest) {",
-            "  let $jscomp$async$arguments = arguments;",
+            "  const $jscomp$async$arguments = arguments;",
             "  function* $jscomp$async$generator() {",
             "    return $jscomp$async$arguments.length;", // arguments replaced
             "  }",
@@ -102,7 +102,7 @@ public void testArgumentsReplacement_asyncClosure() {
         LINE_JOINER.join(
             "function outer() {",
             "  function f() {",
-            "    let $jscomp$async$arguments = arguments;",
+            "    const $jscomp$async$arguments = arguments;",
             "    function* $jscomp$async$generator() {",
             "      return $jscomp$async$arguments.length;", // arguments replaced
             "    }",
@@ -123,7 +123,7 @@ public void testArgumentsReplacement_normalClosureInAsync() {
             "}"),
         LINE_JOINER.join(
             "function a() {",
-            "  let $jscomp$async$arguments = arguments;",
+            "  const $jscomp$async$arguments = arguments;",
             "  function* $jscomp$async$generator() {",
             "    function inner() {",
             "      return arguments.length;", // unchanged
@@ -140,7 +140,7 @@ public void testClassMethod() {
         LINE_JOINER.join(
             "class A {",
             "  f() {",
-            "    let $jscomp$async$this = this;",
+            "    const $jscomp$async$this = this;",
             "    function* $jscomp$async$generator() {",
             "      return $jscomp$async$this.x;", // this replaced
             "    }",
@@ -149,6 +149,33 @@ public void testClassMethod() {
             "}"));
   }
 
+  public void testClassMethodWithAsyncArrow() {
+    test(
+        LINE_JOINER.join(
+            "class A {",
+            "  async f() {",
+            "    let g = async () => { console.log(this); };",
+            "    g();",
+            "  }",
+            "}"),
+        LINE_JOINER.join(
+            "class A {",
+            "  f() {",
+            "    const $jscomp$async$this = this;",
+            "    function *$jscomp$async$generator() {",
+            "      let g = () => {",
+            "        function *$jscomp$async$generator() {",
+            "          console.log($jscomp$async$this);",
+            "        }",
+            "        return $jscomp.executeAsyncGenerator($jscomp$async$generator());",
+            "      };",
+            "      g();",
+            "    }",
+            "    return $jscomp.executeAsyncGenerator($jscomp$async$generator());",
+            "  }",
+            "}"));
+  }
+
   public void testArrowFunctionExpressionBody() {
     test(
         "let f = async () => 1;",

diff --git a/test/com/google/javascript/jscomp/runtime_tests/async_function_test.js b/test/com/google/javascript/jscomp/runtime_tests/async_function_test.js
@@ -144,6 +144,24 @@ testSuite({
     });
   },
 
+  testMemberFunctionUsingThisInAsyncArrowFunction() {
+    class C {
+      constructor() {
+        this.value = 0;
+      }
+
+      async delayedIncrementAndReturnThis() {
+        const nestedArrow = async () => { this.value++; return this; };
+        return await nestedArrow();
+      }
+    }
+    const c = new C();
+    return c.delayedIncrementAndReturnThis().then(result => {
+      assertEquals(c, result);
+      assertEquals(1, c.value);
+    });
+  },
+
   testArgumentsHandledCorrectly() {
     const expected1 = {};
     const expected2 = 2;