Use @closurePrimitive to identify assertion functions in more places

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

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

@@ -16,15 +16,12 @@
 
 package com.google.javascript.jscomp;
 
-import com.google.common.collect.ImmutableSet;
-import com.google.javascript.jscomp.CodingConvention.AssertionFunctionSpec;
+import com.google.javascript.jscomp.CodingConvention.AssertionFunctionLookup;
 import com.google.javascript.jscomp.NodeTraversal.AbstractPostOrderCallback;
 import com.google.javascript.jscomp.NodeTraversal.Callback;
 import com.google.javascript.rhino.Node;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Objects;
-import java.util.Set;
 
 /**
  * <p>Compiler pass that removes Closure-specific code patterns.</p>
@@ -174,26 +171,17 @@ public void visit(NodeTraversal t, Node n, Node parent) {
    * Identifies all assertion calls.
    */
   private class FindAssertionCalls extends AbstractPostOrderCallback {
-    final Set<String> assertionNames;
+    final AssertionFunctionLookup assertionNames;
 
     FindAssertionCalls() {
-      // 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());
+          AssertionFunctionLookup.of(compiler.getCodingConvention().getAssertionFunctions());
     }
 
     @Override
     public void visit(NodeTraversal t, Node n, Node parent) {
-      if (n.isCall()) {
-        String fnName = n.getFirstChild().getQualifiedName();
-        if (assertionNames.contains(fnName)) {
-          assertionCalls.add(n);
-        }
+      if (n.isCall() && assertionNames.lookupByCallee(n.getFirstChild()) != null) {
+        assertionCalls.add(n);
       }
     }
   }

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

@@ -33,7 +33,7 @@
 import com.google.common.reflect.TypeToken;
 import com.google.javascript.jscomp.CheckConformance.InvalidRequirementSpec;
 import com.google.javascript.jscomp.CheckConformance.Rule;
-import com.google.javascript.jscomp.CodingConvention.AssertionFunctionSpec;
+import com.google.javascript.jscomp.CodingConvention.AssertionFunctionLookup;
 import com.google.javascript.jscomp.Requirement.Severity;
 import com.google.javascript.jscomp.Requirement.Type;
 import com.google.javascript.jscomp.parsing.JsDocInfoParser;
@@ -55,7 +55,6 @@
 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;
@@ -314,7 +313,7 @@ protected void report(Node n, ConformanceResult result) {
   abstract static class AbstractTypeRestrictionRule extends AbstractRule {
     private final JSType nativeObjectType;
     private final JSType whitelistedTypes;
-    private final ImmutableList<Node> assertionsFunctionNames;
+    private final AssertionFunctionLookup assertionFunctions;
 
     public AbstractTypeRestrictionRule(AbstractCompiler compiler, Requirement requirement)
         throws InvalidRequirementSpec {
@@ -323,14 +322,8 @@ public AbstractTypeRestrictionRule(AbstractCompiler compiler, Requirement requir
       List<String> whitelistedTypeNames = requirement.getValueList();
       whitelistedTypes = union(whitelistedTypeNames);
 
-      // 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());
+      assertionFunctions =
+          AssertionFunctionLookup.of(compiler.getCodingConvention().getAssertionFunctions());
     }
 
     protected boolean isWhitelistedType(Node n) {
@@ -394,11 +387,7 @@ protected JSType union(List<String> typeNames) {
     protected boolean isAssertionCall(Node n) {
       if (n.isCall() && n.getFirstChild().isQualifiedName()) {
         Node target = n.getFirstChild();
-        for (int i = 0; i < assertionsFunctionNames.size(); i++) {
-          if (target.matchesQualifiedName(assertionsFunctionNames.get(i))) {
-            return true;
-          }
-        }
+        return assertionFunctions.lookupByCallee(target) != null;
       }
       return false;
     }

test/com/google/javascript/jscomp/CheckConformanceTest.java

@@ -1418,16 +1418,41 @@ public void testCustomBanUnknownThis3() {
   }
 
   @Test
-  public void testCustomBanUnknownThis4() {
+  public void testCustomBanUnknownThis_allowsClosurePrimitiveAssert() {
     configuration =
         "requirement: {\n" +
         "  type: CUSTOM\n" +
         "  java_class: 'com.google.javascript.jscomp.ConformanceRules$BanUnknownThis'\n" +
         "  error_message: 'BanUnknownThis Message'\n" +
         "}";
 
-    testNoWarning(
-        "function f() {goog.asserts.assertInstanceof(this, Error);}");
+    String assertInstanceof =
+        lines(
+            "/** @const */ var asserts = {};",
+            "/**",
+            " * @param {?} value The value to check.",
+            " * @param {function(new: T, ...)} type A user-defined constructor.",
+            " * @return {T}",
+            " * @template T",
+            " * @closurePrimitive {asserts.matchesReturn}",
+            " */",
+            "asserts.assertInstanceof = function(value, type) {",
+            "  return value;",
+            "};");
+
+    testNoWarning(lines(assertInstanceof, "function f() {asserts.assertInstanceof(this, Error);}"));
+  }
+
+  @Test
+  public void testCustomBanUnknownThis_allowsGoogAssert() {
+    configuration =
+        "requirement: {\n"
+            + "  type: CUSTOM\n"
+            + "  java_class: 'com.google.javascript.jscomp.ConformanceRules$BanUnknownThis'\n"
+            + "  error_message: 'BanUnknownThis Message'\n"
+            + "}";
+
+    testNoWarning("function f() {goog.asserts.assertInstanceof(this, Error);}");
   }
 
   private static String config(String rule, String message, String... fields) {

test/com/google/javascript/jscomp/ClosureCodeRemovalTest.java

@@ -30,6 +30,12 @@ public final class ClosureCodeRemovalTest extends CompilerTestCase {
 
   private static final String EXTERNS = "var window;";
 
+  private static final String ASSERTIONS =
+      lines(
+          "const asserts = {};",
+          "/** @closurePrimitive {asserts.truthy} */",
+          "asserts.assert = function(...args) {};");
+
   public ClosureCodeRemovalTest() {
     super(EXTERNS);
   }
@@ -129,6 +135,30 @@ public void testAssertionRemoval4() {
     test("var x = goog.asserts.assert();", "var x = void 0;");
   }
 
+  @Test
+  public void testClosurePrimitiveAssertionRemoval1() {
+    enableTypeCheck();
+    test(ASSERTIONS + "var x = asserts.assert(y(), 'message');", ASSERTIONS + "var x = y();");
+  }
+
+  @Test
+  public void testClosurePrimitiveAssertionRemoval2() {
+    enableTypeCheck();
+    test(ASSERTIONS + "asserts.assert(y(), 'message');", ASSERTIONS);
+  }
+
+  @Test
+  public void testClosurePrimitiveAssertionRemoval3() {
+    enableTypeCheck();
+    test(ASSERTIONS + "asserts.assert();", ASSERTIONS);
+  }
+
+  @Test
+  public void testClosurePrimitiveAssertionRemoval4() {
+    enableTypeCheck();
+    test(ASSERTIONS + "var x = asserts.assert();", ASSERTIONS + "var x = void 0;");
+  }
+
   @Test
   public void testDoNotRemoveAbstractClass() {
     testSame(