Support async functions and await in the typechecker

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=202014654

src/com/google/javascript/jscomp/FunctionTypeBuilder.java

@@ -22,6 +22,7 @@
 import static com.google.javascript.jscomp.TypeCheck.BAD_IMPLEMENTED_TYPE;
 import static com.google.javascript.rhino.jstype.JSTypeNative.FUNCTION_FUNCTION_TYPE;
 import static com.google.javascript.rhino.jstype.JSTypeNative.GENERATOR_TYPE;
+import static com.google.javascript.rhino.jstype.JSTypeNative.PROMISE_TYPE;
 import static com.google.javascript.rhino.jstype.JSTypeNative.UNKNOWN_TYPE;
 import static com.google.javascript.rhino.jstype.JSTypeNative.VOID_TYPE;
 
@@ -737,36 +738,47 @@ private boolean addParameter(FunctionParamBuilder builder,
     return emittedWarning;
   }
 
+  /** Sets the returnType for this function using very basic type inference. */
+  private void provideDefaultReturnType() {
+    if (contents.getSourceNode() != null && contents.getSourceNode().isGeneratorFunction()) {
+      // Set the return type of a generator function to:
+      //   @return {!Generator<?>}
+      ObjectType generatorType = typeRegistry.getNativeObjectType(GENERATOR_TYPE);
+      returnType =
+          typeRegistry.createTemplatizedType(
+              generatorType, typeRegistry.getNativeType(UNKNOWN_TYPE));
+    } else if (contents.getSourceNode() != null && contents.getSourceNode().isAsyncFunction()) {
+      // Set the return type of an async function to:
+      //   @return {!Promise<?>}
+      ObjectType promiseType = typeRegistry.getNativeObjectType(PROMISE_TYPE);
+      returnType =
+          typeRegistry.createTemplatizedType(promiseType, typeRegistry.getNativeType(UNKNOWN_TYPE));
+    } else if (!contents.mayHaveNonEmptyReturns()
+        && !contents.mayHaveSingleThrow()
+        && !contents.mayBeFromExterns()) {
+      // Infer return types for non-generator functions.
+      // We need to be extremely conservative about this, because of two
+      // competing needs.
+      // 1) If we infer the return type of f too widely, then we won't be able
+      //    to assign f to other functions.
+      // 2) If we infer the return type of f too narrowly, then we won't be
+      //    able to override f in subclasses.
+      // So we only infer in cases where the user doesn't expect to write
+      // @return annotations--when it's very obvious that the function returns
+      // nothing.
+      returnType = typeRegistry.getNativeType(VOID_TYPE);
+      returnTypeInferred = true;
+    } else {
+      returnType = typeRegistry.getNativeType(UNKNOWN_TYPE);
+    }
+  }
+
   /**
    * Builds the function type, and puts it in the registry.
    */
   FunctionType buildAndRegister() {
     if (returnType == null) {
-      if (contents.getSourceNode() != null && contents.getSourceNode().isGeneratorFunction()) {
-        // Set the return type of a generator function to:
-        //   @return {!Generator<?>}
-        ObjectType generatorType = typeRegistry.getNativeObjectType(GENERATOR_TYPE);
-        returnType =
-            typeRegistry.createTemplatizedType(
-                generatorType, typeRegistry.getNativeType(UNKNOWN_TYPE));
-      } else if (!contents.mayHaveNonEmptyReturns()
-          && !contents.mayHaveSingleThrow()
-          && !contents.mayBeFromExterns()) {
-        // Infer return types for non-generator functions.
-        // We need to be extremely conservative about this, because of two
-        // competing needs.
-        // 1) If we infer the return type of f too widely, then we won't be able
-        //    to assign f to other functions.
-        // 2) If we infer the return type of f too narrowly, then we won't be
-        //    able to override f in subclasses.
-        // So we only infer in cases where the user doesn't expect to write
-        // @return annotations--when it's very obvious that the function returns
-        // nothing.
-        returnType = typeRegistry.getNativeType(VOID_TYPE);
-        returnTypeInferred = true;
-      } else {
-        returnType = typeRegistry.getNativeType(UNKNOWN_TYPE);
-      }
+      provideDefaultReturnType();
     }
 
     if (parametersNode == null) {

src/com/google/javascript/jscomp/Promises.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright 2018 The Closure Compiler Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.javascript.jscomp;
+
+import com.google.javascript.rhino.jstype.JSType;
+import com.google.javascript.rhino.jstype.JSTypeNative;
+import com.google.javascript.rhino.jstype.JSTypeRegistry;
+import com.google.javascript.rhino.jstype.TemplateTypeMap;
+import com.google.javascript.rhino.jstype.UnionTypeBuilder;
+
+/**
+ * Models different Javascript Promise-related operations
+ *
+ * @author lharker@google.com (Laura Harker)
+ */
+final class Promises {
+
+  private Promises() {}
+
+  /**
+   * If this object is known to be an IThenable, returns the type it resolves to.
+   *
+   * <p>Returns unknown otherwise.
+   *
+   * <p>(This is different from {@code getResolvedType}, which will attempt to model the then type
+   * of an expression after calling Promise.resolve() on it.
+   */
+  static final JSType getTemplateTypeOfThenable(JSTypeRegistry registry, JSType maybeThenable) {
+    return maybeThenable
+        // TODO(lharker): if we ban declaring async functions to return nullable promises/thenables
+        // then we should remove this .restrictByNotNullOrUndefined()
+        .restrictByNotNullOrUndefined()
+        .getInstantiatedTypeArgument(registry.getNativeType(JSTypeNative.I_THENABLE_TYPE));
+  }
+
+  /**
+   * Returns the type of `await [expr]`.
+   *
+   * <p>This is equivalent to the type of `result` in `Promise.resolve([expr]).then(result => `
+   *
+   * <p>For example:
+   *
+   * <p>{@code !Promise<number>} becomes {@code number}
+   *
+   * <p>{@code !IThenable<number>} becomes {@code number}
+   *
+   * <p>{@code string} becomes {@code string}
+   *
+   * <p>{@code (!Promise<number>|string)} becomes {@code (number|string)}
+   *
+   * <p>{@code ?Promise<number>} becomes {@code (null|number)}
+   */
+  static final JSType getResolvedType(JSTypeRegistry registry, JSType type) {
+    if (type.isUnknownType()) {
+      return type;
+    }
+
+    if (type.isUnionType()) {
+      UnionTypeBuilder unionTypeBuilder = new UnionTypeBuilder(registry);
+      for (JSType alternate : type.toMaybeUnionType().getAlternatesWithoutStructuralTyping()) {
+        unionTypeBuilder.addAlternate(getResolvedType(registry, alternate));
+      }
+      return unionTypeBuilder.build();
+    }
+
+    // If we can find the "IThenable" template key (which is true for Promise and IThenable), return
+    // the resolved value. e.g. for "!Promise<string>" return "string".
+    TemplateTypeMap templates = type.getTemplateTypeMap();
+    if (templates.hasTemplateKey(registry.getIThenableTemplate())) {
+      // Call getResolvedPromiseType again in case someone does something unusual like
+      // !Promise<!Promise<number>>
+      // TODO(lharker): we don't need to handle this case and should report an error for this in a
+      // type annotation (not here, maybe in TypeCheck). A Promise cannot resolve to another Promise
+      return getResolvedType(
+          registry, templates.getResolvedTemplateType(registry.getIThenableTemplate()));
+    }
+
+    return type;
+  }
+}

src/com/google/javascript/jscomp/TypeCheck.java

@@ -617,7 +617,10 @@ public void visit(NodeTraversal t, Node n, Node parent) {
 
       case YIELD:
         visitYield(t, n);
-        typeable = false;
+        break;
+
+      case AWAIT:
+        ensureTyped(n);
         break;
 
       case DEC:
@@ -1956,6 +1959,11 @@ private void visitFunction(NodeTraversal t, Node n) {
       JSType returnType = functionType.getReturnType();
       validator.expectGeneratorSupertype(
           t, n, returnType, "A generator function must return a (supertype of) Generator");
+
+    } else if (n.isAsyncFunction()) {
+      // An async function must return a Promise or supertype of Promise
+      JSType returnType = functionType.getReturnType();
+      validator.expectValidAsyncReturnType(t, n, returnType.restrictByNotNullOrUndefined());
     }
   }
 
@@ -2284,8 +2292,10 @@ private void checkArgumentsMatchParameters(
     }
   }
 
+  /** Visits an arrow function expression body. */
   private void visitImplicitReturnExpression(NodeTraversal t, Node exprNode) {
-    JSType jsType = getJSType(t.getEnclosingFunction());
+    Node enclosingFunction = t.getEnclosingFunction();
+    JSType jsType = getJSType(enclosingFunction);
     if (jsType.isFunctionType()) {
       FunctionType functionType = jsType.toMaybeFunctionType();
 
@@ -2294,10 +2304,16 @@ private void visitImplicitReturnExpression(NodeTraversal t, Node exprNode) {
       // (it's a void function)
       if (expectedReturnType == null) {
         expectedReturnType = getNativeType(VOID_TYPE);
+      } else if (enclosingFunction.isAsyncFunction()) {
+        // Unwrap the async function's declared return type.
+        expectedReturnType = Promises.getTemplateTypeOfThenable(typeRegistry, expectedReturnType);
       }
 
       // Fetch the returned value's type
       JSType actualReturnType = getJSType(exprNode);
+      if (enclosingFunction.isAsyncFunction()) {
+        actualReturnType = Promises.getResolvedType(typeRegistry, actualReturnType);
+      }
 
       validator.expectCanAssignTo(t, exprNode, actualReturnType, expectedReturnType,
           "inconsistent return type");
@@ -2335,6 +2351,10 @@ private void visitReturn(NodeTraversal t, Node n) {
         // Unwrap the template variable from a generator function's declared return type.
         // e.g. if returnType is "Generator<string>", make it just "string".
         returnType = getTemplateTypeOfGenerator(returnType);
+      } else if (enclosingFunction.isAsyncFunction()) {
+        // Unwrap the template variable from a async function's declared return type.
+        // e.g. if returnType is "!Promise<string>" or "!IThenable<string>", make it just "string".
+        returnType = Promises.getTemplateTypeOfThenable(typeRegistry, returnType);
       }
 
       // fetching the returned value's type
@@ -2345,11 +2365,16 @@ private void visitReturn(NodeTraversal t, Node n) {
         valueNode = n;
       } else {
         actualReturnType = getJSType(valueNode);
+        if (enclosingFunction.isAsyncFunction()) {
+          // We want to treat `return Promise.resolve(1);` as if it were `return 1;` inside an async
+          // function.
+          actualReturnType = Promises.getResolvedType(typeRegistry, actualReturnType);
+        }
       }
 
       // verifying
-      validator.expectCanAssignTo(t, valueNode, actualReturnType, returnType,
-          "inconsistent return type");
+      validator.expectCanAssignTo(
+          t, valueNode, actualReturnType, returnType, "inconsistent return type");
     }
   }
 

src/com/google/javascript/jscomp/TypeInference.java

@@ -537,6 +537,10 @@ private FlowScope traverse(Node n, FlowScope scope) {
         scope = traverseChildren(n, scope);
         break;
 
+      case AWAIT:
+        scope = traverseAwait(n, scope);
+        break;
+
       default:
         break;
     }
@@ -1865,6 +1869,16 @@ private BooleanOutcomePair traverseWithinShortCircuitingBinOp(
     }
   }
 
+  private FlowScope traverseAwait(Node await, FlowScope scope) {
+    scope = traverseChildren(await, scope);
+
+    Node expr = await.getFirstChild();
+    JSType exprType = getJSType(expr);
+    await.setJSType(Promises.getResolvedType(registry, exprType));
+
+    return scope;
+  }
+
   private static BooleanLiteralSet joinBooleanOutcomes(
       boolean isAnd, BooleanLiteralSet left, BooleanLiteralSet right) {
     // A truthy value on the lhs of an {@code &&} can never make it to the

src/com/google/javascript/jscomp/TypeValidator.java

@@ -25,6 +25,7 @@
 import static com.google.javascript.rhino.jstype.JSTypeNative.GENERATOR_TYPE;
 import static com.google.javascript.rhino.jstype.JSTypeNative.ITERABLE_TYPE;
 import static com.google.javascript.rhino.jstype.JSTypeNative.I_TEMPLATE_ARRAY_TYPE;
+import static com.google.javascript.rhino.jstype.JSTypeNative.I_THENABLE_TYPE;
 import static com.google.javascript.rhino.jstype.JSTypeNative.NO_OBJECT_TYPE;
 import static com.google.javascript.rhino.jstype.JSTypeNative.NULL_TYPE;
 import static com.google.javascript.rhino.jstype.JSTypeNative.NUMBER_STRING;
@@ -305,6 +306,28 @@ void expectGeneratorSupertype(NodeTraversal t, Node n, JSType type, String msg)
     }
   }
 
+  /**
+   * Expect the type to be an IThenable, Promise, or the all type/unknown type/Object type.
+   *
+   * <p>We forbid returning a union type containing Promise/IThenable because it complicates how we
+   * typecheck returns within async functions.
+   */
+  @SuppressWarnings("ReferenceEquality")
+  void expectValidAsyncReturnType(NodeTraversal t, Node n, JSType type) {
+    // Allow returning `?`, `*`, or `Object`.
+    if (type.isUnknownType()
+        || type.isAllType()
+        || type == getNativeType(JSTypeNative.OBJECT_TYPE)) {
+      return;
+    }
+
+    // Get "Promise" from "Promise<string>" or "IThenable" from "IThenable<!Array<number>>"
+    if (!type.getTemplateTypeMap().hasTemplateKey(typeRegistry.getIThenableTemplate())) {
+      mismatch(
+          t, n, "An async function must return a (supertype of) Promise", type, I_THENABLE_TYPE);
+    }
+  }
+
   /** Expect the type to be an ITemplateArray or supertype of ITemplateArray. */
   void expectITemplateArraySupertype(NodeTraversal t, Node n, JSType type, String msg) {
     if (!getNativeType(I_TEMPLATE_ARRAY_TYPE).isSubtypeOf(type)) {

src/com/google/javascript/rhino/jstype/JSType.java

@@ -553,7 +553,11 @@ public void extendTemplateTypeMap(TemplateTypeMap otherMap) {
   public boolean isObject() {
     return false;
   }
-
+  /**
+   * Tests whether this type is an {@code Object}, or any subtype thereof.
+   *
+   * @return <code>this &lt;: Object</code>
+   */
   public final boolean isObjectType() {
     return isObject();
   }
@@ -598,6 +602,10 @@ public final boolean isNominalConstructor() {
     return false;
   }
 
+  public boolean isNativeObjectType() {
+    return false;
+  }
+
   /**
    * Whether this type is an Instance object of some constructor.
    * Does not necessarily mean this is an {@link InstanceObjectType}.

src/com/google/javascript/rhino/jstype/JSTypeRegistry.java

@@ -276,6 +276,11 @@ public TemplateType getIteratorTemplate() {
     return checkNotNull(iteratorTemplate);
   }
 
+  /** @return The template variable for the IThenable interface. */
+  public TemplateType getIThenableTemplate() {
+    return checkNotNull(iThenableTemplateKey);
+  }
+
   /** @return return an immutable list of template types of the given builtin. */
   public ImmutableList<TemplateType> maybeGetTemplateTypesOfBuiltin(String fnName) {
     JSType type = getType(null, fnName);