Identify AssertionFunctionSpecs by ClosurePrimitive as well as global…

… name

This creates two new ClosurePrimitives:
  - ASSERTS_TRUTHY: equivalent to goog.asserts.assert
  - ASSERTS_MATCHES_RETURN: equivalent to goog.asserts.assertString/Element/Array/etc.

Creates a new structure 'AssertionFunctionMap' to simplify doing lookups by both ClosurePrimitive and string name, and makes TypeInference use that structure

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

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

@@ -66,8 +66,8 @@ public ImmutableCollection<String> getIndirectlyDeclaredProperties() {
   @Override
   public ImmutableCollection<AssertionFunctionSpec> getAssertionFunctions() {
     return ImmutableList.of(
-        AssertionFunctionSpec.makeTruthyAssertion("assert"),
-        AssertionFunctionSpec.makeReturnTypeAssertion("cr.ui.decorate"));
+        AssertionFunctionSpec.forTruthy().setFunctionName("assert").build(),
+        AssertionFunctionSpec.forMatchesReturn().setFunctionName("cr.ui.decorate").build());
   }
 
   @Override

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

@@ -23,6 +23,7 @@
 import com.google.javascript.rhino.Node;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 
 /**
@@ -176,15 +177,16 @@ private class FindAssertionCalls extends AbstractPostOrderCallback {
     final Set<String> assertionNames;
 
     FindAssertionCalls() {
-      ImmutableSet.Builder<String> assertionNamesBuilder = ImmutableSet.builder();
-      for (AssertionFunctionSpec spec :
-               compiler.getCodingConvention().getAssertionFunctions()) {
-        assertionNamesBuilder.add(spec.getFunctionName());
-      }
-      assertionNames = assertionNamesBuilder.build();
+      // TODO(b/126254920): make this use ClosurePrimitive instead,
+      assertionNames =
+          compiler.getCodingConvention().getAssertionFunctions().stream()
+              .map(AssertionFunctionSpec::getFunctionName)
+              // Filter out assertion functions with a null functionName, which are identified only
+              // by ClosurePrimitive id.
+              .filter(Objects::nonNull)
+              .collect(ImmutableSet.toImmutableSet());
     }
 
-
     @Override
     public void visit(NodeTraversal t, Node n, Node parent) {
       if (n.isCall()) {

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

@@ -382,16 +382,26 @@ public boolean isPrivate(String name) {
 
   @Override
   public ImmutableCollection<AssertionFunctionSpec> getAssertionFunctions() {
-    return ImmutableList.of(
-        AssertionFunctionSpec.makeTruthyAssertion("goog.asserts.assert"),
-        AssertionFunctionSpec.makeReturnTypeAssertion("goog.asserts.assertArray"),
-        AssertionFunctionSpec.makeReturnTypeAssertion("goog.asserts.assertBoolean"),
-        AssertionFunctionSpec.makeReturnTypeAssertion("goog.asserts.assertElement"),
-        AssertionFunctionSpec.makeReturnTypeAssertion("goog.asserts.assertFunction"),
-        AssertionFunctionSpec.makeReturnTypeAssertion("goog.asserts.assertInstanceof"),
-        AssertionFunctionSpec.makeReturnTypeAssertion("goog.asserts.assertNumber"),
-        AssertionFunctionSpec.makeReturnTypeAssertion("goog.asserts.assertObject"),
-        AssertionFunctionSpec.makeReturnTypeAssertion("goog.asserts.assertString"));
+    return ImmutableSet.<AssertionFunctionSpec>builder()
+        .addAll(super.getAssertionFunctions())
+        .add(
+            AssertionFunctionSpec.forTruthy().setFunctionName("goog.asserts.assert").build(),
+            createGoogAssertOnReturn("Array"),
+            createGoogAssertOnReturn("Boolean"),
+            createGoogAssertOnReturn("Element"),
+            createGoogAssertOnReturn("Function"),
+            createGoogAssertOnReturn("Instanceof"),
+            createGoogAssertOnReturn("Number"),
+            createGoogAssertOnReturn("Object"),
+            createGoogAssertOnReturn("String"))
+        .build();
+  }
+
+  /** Returns a new assertion function goog.asserts.assert[assertedTypeName] */
+  private static AssertionFunctionSpec createGoogAssertOnReturn(String assertedTypeName) {
+    return AssertionFunctionSpec.forMatchesReturn()
+        .setFunctionName("goog.asserts.assert" + assertedTypeName)
+        .build();
   }
 
   @Override

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

@@ -17,17 +17,23 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 
+import com.google.auto.value.AutoValue;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
 import com.google.errorprone.annotations.Immutable;
+import com.google.javascript.rhino.ClosurePrimitive;
 import com.google.javascript.rhino.Node;
 import com.google.javascript.rhino.NominalTypeBuilder;
 import com.google.javascript.rhino.StaticSourceFile;
 import com.google.javascript.rhino.jstype.FunctionType;
+import com.google.javascript.rhino.jstype.JSType;
 import com.google.javascript.rhino.jstype.JSTypeRegistry;
 import com.google.javascript.rhino.jstype.ObjectType;
 import java.io.Serializable;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 import javax.annotation.Nullable;
 
 /**
@@ -486,50 +492,66 @@ static class ObjectLiteralCast {
    *       match-return-type handling does not have.
    * </ul>
    */
-  final class AssertionFunctionSpec {
-    private final String functionName;
-    private final AssertionKind kind;
-    private final int paramIndex; // the index of the formal parameter that is actually asserted
+  @Immutable
+  @AutoValue
+  abstract class AssertionFunctionSpec {
+    // TODO(b/126254920): remove this field and always use ClosurePrimitive
+    @Nullable
+    abstract String getFunctionName();
+
+    @Nullable
+    abstract ClosurePrimitive getClosurePrimitive();
+
+    abstract AssertionKind getAssertionKind();
+
+    abstract int getParamIndex(); // the index of the formal parameter that is actually asserted
 
     public enum AssertionKind {
       TRUTHY, // an assertion that the parameter is 'truthy'
       MATCHES_RETURN_TYPE // an assertion that the parameter matches the inferred return kind
     }
 
-    private AssertionFunctionSpec(String functionName, AssertionKind kind, int paramIndex) {
-      this.functionName = functionName;
-      this.kind = kind;
-      this.paramIndex = paramIndex;
+    static Builder builder() {
+      return new AutoValue_CodingConvention_AssertionFunctionSpec.Builder().setParamIndex(0);
     }
 
-    /** Returns a truthy assertion function on the first param */
-    public static AssertionFunctionSpec makeTruthyAssertion(String functionName) {
-      return new AssertionFunctionSpec(functionName, AssertionKind.TRUTHY, /* paramIndex= */ 0);
+    public static Builder forTruthy() {
+      return builder().setAssertionKind(AssertionKind.TRUTHY);
     }
 
-    /** Returns an assertion function asserting the first param matches the return type */
-    public static AssertionFunctionSpec makeReturnTypeAssertion(String functionName) {
-      return new AssertionFunctionSpec(functionName, AssertionKind.MATCHES_RETURN_TYPE, 0);
+    public static Builder forMatchesReturn() {
+      return builder().setAssertionKind(AssertionKind.MATCHES_RETURN_TYPE);
     }
 
-    /** Returns an assertion function asserting the nth arg is truthy (0-indexed) */
-    public static AssertionFunctionSpec makeTruthyAssertion(String functionName, int paramIndex) {
-      return new AssertionFunctionSpec(functionName, AssertionKind.TRUTHY, paramIndex);
-    }
+    @AutoValue.Builder
+    public abstract static class Builder {
+      public abstract Builder setFunctionName(String name);
+
+      abstract Builder setClosurePrimitive(ClosurePrimitive primitive);
+
+      public abstract Builder setParamIndex(int paramIndex);
 
-    /** Returns the name of the function. */
-    public String getFunctionName() {
-      return functionName;
+      abstract Builder setAssertionKind(AssertionKind kind);
+
+      abstract AssertionFunctionSpec autoBuild();
+
+      public AssertionFunctionSpec build() {
+        AssertionFunctionSpec spec = autoBuild();
+        Preconditions.checkState(
+            spec.getFunctionName() != null || spec.getClosurePrimitive() != null,
+            "Must provide a function name or ClosurePrimitive for each spec");
+        return spec;
+      }
     }
 
-    AssertionKind getAssertionKind() {
-      return this.kind;
+    private Object getId() {
+      return getClosurePrimitive() != null ? getClosurePrimitive() : getFunctionName();
     }
 
     /** Returns which argument is actually being asserted, or null if fewer args than expected */
     @Nullable
-    public Node getAssertedArg(Node firstArg) {
-      for (int i = 0; i < paramIndex; i++) {
+    Node getAssertedArg(Node firstArg) {
+      for (int i = 0; i < getParamIndex(); i++) {
         if (firstArg == null) {
           // If there are fewer arguments than expected, return null instead of crashing in this
           // function.
@@ -540,4 +562,54 @@ public Node getAssertedArg(Node firstArg) {
       return firstArg;
     }
   }
+
+  /** This stores a relation from either name or Closure Primitive to assertion function */
+  @Immutable
+  final class AssertionFunctionLookup {
+    // the key type 'Object' is mutable, but at runtime is only ever a string or ClosurePrimitive
+    @SuppressWarnings("Immutable")
+    private final ImmutableMap<Object, AssertionFunctionSpec> internal;
+
+    private AssertionFunctionLookup(ImmutableMap<Object, AssertionFunctionSpec> internal) {
+      this.internal = internal;
+    }
+
+    /**
+     * Returns a new map containing all the given {@link AssertionFunctionSpec}s
+     *
+     * <p>Assumes that in the input, there is a unique mapping from string name to spec and closure
+     * primitive to spec.
+     */
+    static AssertionFunctionLookup of(Collection<AssertionFunctionSpec> specs) {
+      ImmutableMap<Object, AssertionFunctionSpec> idToSpecMap =
+          specs.stream()
+              .collect(
+                  ImmutableMap.toImmutableMap(AssertionFunctionSpec::getId, Function.identity()));
+
+      return new AssertionFunctionLookup(idToSpecMap);
+    }
+
+    /**
+     * Returns the {@link AssertionFunctionSpec} matching the given function reference.
+     *
+     * <p>This first looks up specs by their ClosurePrimitive, then falls back to qualified name
+     */
+    @Nullable
+    AssertionFunctionSpec lookupByCallee(Node callee) {
+      FunctionType fnType = JSType.toMaybeFunctionType(callee.getJSType());
+      if (fnType != null && fnType.getClosurePrimitive() != null) {
+        AssertionFunctionSpec spec = internal.get(fnType.getClosurePrimitive());
+        if (spec != null) {
+          return spec;
+        }
+      }
+
+      // TODO(b/126254920): remove this
+      if (callee.isQualifiedName()) {
+        return internal.get(callee.getQualifiedName());
+      }
+
+      return null;
+    }
+  }
 }

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

@@ -19,6 +19,7 @@
 import static com.google.common.base.Preconditions.checkState;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.Immutable;
 import com.google.javascript.rhino.ClosurePrimitive;
 import com.google.javascript.rhino.Node;
@@ -28,7 +29,6 @@
 import com.google.javascript.rhino.jstype.JSTypeRegistry;
 import com.google.javascript.rhino.jstype.ObjectType;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -551,8 +551,14 @@ public ObjectLiteralCast getObjectLiteralCast(Node callNode) {
     }
 
     @Override
-    public Collection<AssertionFunctionSpec> getAssertionFunctions() {
-      return Collections.emptySet();
+    public ImmutableSet<AssertionFunctionSpec> getAssertionFunctions() {
+      return ImmutableSet.of(
+          AssertionFunctionSpec.forTruthy()
+              .setClosurePrimitive(ClosurePrimitive.ASSERTS_TRUTHY)
+              .build(),
+          AssertionFunctionSpec.forMatchesReturn()
+              .setClosurePrimitive(ClosurePrimitive.ASSERTS_MATCHES_RETURN)
+              .build());
     }
 
     @Override

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

@@ -55,6 +55,7 @@
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.regex.Pattern;
 import java.util.regex.PatternSyntaxException;
@@ -322,11 +323,14 @@ public AbstractTypeRestrictionRule(AbstractCompiler compiler, Requirement requir
       List<String> whitelistedTypeNames = requirement.getValueList();
       whitelistedTypes = union(whitelistedTypeNames);
 
-      ImmutableList.Builder<Node> builder = ImmutableList.builder();
-      for (AssertionFunctionSpec fn : compiler.getCodingConvention().getAssertionFunctions()) {
-        builder.add(NodeUtil.newQName(compiler, fn.getFunctionName()));
-      }
-      assertionsFunctionNames = builder.build();
+      // TODO(b/126254920): make this use ClosurePrimitive instead
+      // so that we don't have to filter out the null functionNames.
+      assertionsFunctionNames =
+          compiler.getCodingConvention().getAssertionFunctions().stream()
+              .map(AssertionFunctionSpec::getFunctionName)
+              .filter(Objects::nonNull)
+              .map(name -> NodeUtil.newQName(compiler, name))
+              .collect(ImmutableList.toImmutableList());
     }
 
     protected boolean isWhitelistedType(Node n) {

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

@@ -36,6 +36,7 @@
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
+import com.google.javascript.jscomp.CodingConvention.AssertionFunctionLookup;
 import com.google.javascript.jscomp.CodingConvention.AssertionFunctionSpec;
 import com.google.javascript.jscomp.ControlFlowGraph.Branch;
 import com.google.javascript.jscomp.graph.DiGraph.DiGraphEdge;
@@ -89,18 +90,21 @@ class TypeInference
   private final FlowScope bottomScope;
   private final TypedScope containerScope;
   private final TypedScopeCreator scopeCreator;
-  private final Map<String, AssertionFunctionSpec> assertionFunctionsMap;
+  private final AssertionFunctionLookup assertionFunctionLookup;
 
   // Scopes that have had their unbound untyped vars inferred as undefined.
   private final Set<TypedScope> inferredUnboundVars = new HashSet<>();
 
   // For convenience
   private final ObjectType unknownType;
 
-  TypeInference(AbstractCompiler compiler, ControlFlowGraph<Node> cfg,
-                ReverseAbstractInterpreter reverseInterpreter,
-                TypedScope syntacticScope, TypedScopeCreator scopeCreator,
-                Map<String, AssertionFunctionSpec> assertionFunctionsMap) {
+  TypeInference(
+      AbstractCompiler compiler,
+      ControlFlowGraph<Node> cfg,
+      ReverseAbstractInterpreter reverseInterpreter,
+      TypedScope syntacticScope,
+      TypedScopeCreator scopeCreator,
+      AssertionFunctionLookup assertionFunctionLookup) {
     super(cfg, new LinkedFlowScope.FlowScopeJoinOp());
     this.compiler = compiler;
     this.registry = compiler.getTypeRegistry();
@@ -110,7 +114,7 @@ class TypeInference
     this.containerScope = syntacticScope;
 
     this.scopeCreator = scopeCreator;
-    this.assertionFunctionsMap = assertionFunctionsMap;
+    this.assertionFunctionLookup = assertionFunctionLookup;
 
     FlowScope entryScope =
         inferDeclarativelyUnboundVarsWithoutTypes(
@@ -1519,9 +1523,13 @@ private FlowScope traverseFunctionInvocation(Node n, FlowScope scope) {
   private FlowScope tightenTypesAfterAssertions(FlowScope scope, Node callNode) {
     Node left = callNode.getFirstChild();
     Node firstParam = left.getNext();
-    AssertionFunctionSpec assertionFunctionSpec =
-        assertionFunctionsMap.get(left.getQualifiedName());
-    if (assertionFunctionSpec == null || firstParam == null) {
+    if (firstParam == null) {
+      // this may be an assertion call but there are no arguments to assert
+      return scope;
+    }
+    AssertionFunctionSpec assertionFunctionSpec = assertionFunctionLookup.lookupByCallee(left);
+    if (assertionFunctionSpec == null) {
+      // this is not a recognized assertion function
       return scope;
     }
     Node assertedNode = assertionFunctionSpec.getAssertedArg(firstParam);