Make native module checks able to resolve Closure namespaces in types

This CL mimics the current behavior of type resolution when ClosureRewriteModule runs pre-typechecking. Essentially: if a prefix of a type name matches a non-legacy goog.module, and no longer prefix match any legacy Closure namespaces, then assume that prefix must correspond to the module's namespace.

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

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

@@ -587,6 +587,10 @@ private void gatherAllProvides(Node root) {
         // once leaving the module.
         providedNamesFromCall.put(name.getFirstProvideCall(), name);
       }
+
+      if (metadata != null && metadata.isGoogProvide()) {
+        typeRegistry.registerLegacyClosureModule(name.getNamespace());
+      }
     }
   }
 
@@ -845,12 +849,13 @@ void build() {
     }
 
     private void finishDeclaringGoogModule() {
-      if (module == null || !module.metadata().isLegacyGoogModule()) {
-        // TODO(b/134523248): consider supporting referring to non-legacy modules in types.
+      if (module == null || !module.metadata().isGoogModule()) {
         return;
       }
       TypedVar exportsVar = checkNotNull(currentScope.getSlot("exports"));
+
       if (module.metadata().isLegacyGoogModule()) {
+        typeRegistry.registerLegacyClosureModule(module.closureNamespace());
         QualifiedName moduleNamespace = QualifiedName.of(module.closureNamespace());
         currentScope
             .getGlobalScope()
@@ -870,14 +875,16 @@ private void finishDeclaringGoogModule() {
                     moduleNamespace.getComponent(), exportsVar.getType(), exportsVar.getNameNode());
           }
         }
-        // All goog.modules are accessible by their namespace as types, but not as values.
         declareAliasTypeIfRvalueIsAliasable(
             module.closureNamespace(),
             exportsVar.getNameNode(), // Pretend that 'exports = '... is the lvalue node.
             QualifiedName.of("exports"),
             exportsVar.getType(),
             currentScope,
             currentScope.getGlobalScope());
+      } else {
+        typeRegistry.registerClosureModule(
+            module.closureNamespace(), exportsVar.getNameNode(), exportsVar.getType());
       }
     }
 

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

@@ -47,6 +47,7 @@
 import static com.google.javascript.rhino.jstype.JSTypeNative.UNKNOWN_TYPE;
 import static com.google.javascript.rhino.jstype.JSTypeNative.VOID_TYPE;
 
+import com.google.auto.value.AutoValue;
 import com.google.common.annotations.GwtIncompatible;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -153,6 +154,9 @@ public class JSTypeRegistry implements Serializable {
 
   private final Table<Node, String, JSType> scopedNameTable = HashBasedTable.create();
 
+  // Only needed for type resolution at the moment
+  private final transient Map<String, ModuleSlot> moduleToSlotMap = new HashMap<>();
+
   // NOTE: This would normally be "static final" but that causes unit test failures
   // when serializing and deserializing compiler state for multistage builds.
   private final Node nameTableGlobalRoot = new Node(Token.ROOT);
@@ -2349,4 +2353,43 @@ private FunctionType nativeRecord(String name, TemplateType... templateKeys) {
     type.setImplicitMatch(true);
     return type;
   }
+
+  /** Ensures that a type annotation pointing to a Closure modules is correctly resolved. */
+  public void registerClosureModule(String moduleName, Node definitionNode, JSType type) {
+    moduleToSlotMap.put(moduleName, ModuleSlot.create(/* isLegacy= */ false, definitionNode, type));
+  }
+
+  /**
+   * Ensures that a type annotation pointing to a Closure modules is correctly resolved.
+   *
+   * <p>Currently this is useful because module rewriting will prevent type resolution given a
+   */
+  public void registerLegacyClosureModule(String moduleName) {
+    moduleToSlotMap.put(moduleName, ModuleSlot.create(/* isLegacy= */ true, null, null));
+  }
+
+  /**
+   * Returns the associated slot, if any, for the given module namespace.
+   *
+   * <p>Returns null if the given name is not a Closure namespace from a goog.provide or goog.module
+   */
+  ModuleSlot getModuleSlot(String moduleName) {
+    return moduleToSlotMap.get(moduleName);
+  }
+
+  /** Stores information about a Closure namespace. */
+  @AutoValue
+  abstract static class ModuleSlot {
+    abstract boolean isLegacyModule();
+
+    @Nullable
+    abstract Node definitionNode();
+
+    @Nullable
+    abstract JSType type();
+
+    static ModuleSlot create(boolean isLegacy, Node definitionNode, JSType type) {
+      return new AutoValue_JSTypeRegistry_ModuleSlot(isLegacy, definitionNode, type);
+    }
+  }
 }

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

@@ -43,11 +43,13 @@
 import static com.google.common.base.Preconditions.checkState;
 
 import com.google.common.base.Predicate;
+import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
 import com.google.javascript.rhino.ErrorReporter;
 import com.google.javascript.rhino.Node;
 import com.google.javascript.rhino.StaticScope;
 import com.google.javascript.rhino.StaticSlot;
+import com.google.javascript.rhino.jstype.JSTypeRegistry.ModuleSlot;
 import java.util.ArrayList;
 import java.util.List;
 import javax.annotation.Nullable;
@@ -244,9 +246,14 @@ JSType resolveInternal(ErrorReporter reporter) {
 
     // TODO(user): Investigate whether it is really necessary to keep two
     // different mechanisms for resolving named types, and if so, which order
-    // makes more sense. Now, resolution via registry is first in order to
-    // avoid triggering the warnings built into the resolution via properties.
-    boolean resolved = resolveViaRegistry(reporter);
+    // makes more sense.
+    // The `resolveViaClosureNamespace` mechanism can probably be deleted (or reworked) once the
+    // compiler supports type annotations via path. The `resolveViaProperties` and
+    // `resolveViaRegistry` are, unfortunately, both needed now with no migration plan.
+    boolean resolved = resolveViaClosureNamespace(reporter);
+    if (!resolved) {
+      resolved = resolveViaRegistry(reporter);
+    }
     if (!resolved) {
       resolveViaProperties(reporter);
     }
@@ -298,46 +305,65 @@ private boolean resolveViaRegistry(ErrorReporter reporter) {
    * as properties. The scope must have been fully parsed and a symbol table constructed.
    */
   private void resolveViaProperties(ErrorReporter reporter) {
-    String[] componentNames = reference.split("\\.", -1);
-    if (componentNames[0].length() == 0) {
+    List<String> componentNames = Splitter.on('.').splitToList(reference);
+    if (componentNames.get(0).isEmpty()) {
       handleUnresolvedType(reporter, /* ignoreForwardReferencedTypes= */ true);
       return;
     }
 
     StaticTypedSlot slot =
         checkNotNull(resolutionScope, "resolutionScope")
-            .getSlot(checkNotNull(componentNames, "componentNames")[0]);
+            .getSlot(checkNotNull(componentNames, "componentNames").get(0));
     if (slot == null) {
       handleUnresolvedType(reporter, /* ignoreForwardReferencedTypes= */ true);
       return;
     }
+    Node definitionNode = slot.getDeclaration() != null ? slot.getDeclaration().getNode() : null;
+    resolveViaPropertyGivenSlot(
+        slot.getType(), definitionNode, componentNames, reporter, /* componentIndex= */ 1);
+  }
+
+  /**
+   * Resolve a type using a given StaticTypedSlot and list of properties on that type.
+   *
+   * @param slotType the JSType of teh slot, possibly null
+   * @param definitionNode If known, the Node representing the type definition.
+   * @param componentIndex the index into {@code componentNames} at which to start resolving
+   */
+  private void resolveViaPropertyGivenSlot(
+      JSType slotType,
+      Node definitionNode,
+      List<String> componentNames,
+      ErrorReporter reporter,
+      int componentIndex) {
+    if (resolveTypeFromNodeIfTypedef(definitionNode, reporter)) {
+      return;
+    }
 
     // If the first component has a type of 'Unknown', then any type
     // names using it should be regarded as silently 'Unknown' rather than be
     // noisy about it.
-    JSType slotType = slot.getType();
     if (slotType == null || slotType.isAllType() || slotType.isNoType()) {
       handleUnresolvedType(reporter, /* ignoreForwardReferencedTypes= */ true);
       return;
     }
 
     // resolving component by component
-    for (int i = 1; i < componentNames.length; i++) {
+    for (int i = componentIndex; i < componentNames.size(); i++) {
+      String component = componentNames.get(i);
       ObjectType parentObj = ObjectType.cast(slotType);
-      if (parentObj == null || componentNames[i].length() == 0) {
+      if (parentObj == null || component.length() == 0) {
         handleUnresolvedType(reporter, /* ignoreForwardReferencedTypes= */ true);
         return;
       }
-      if (i == componentNames.length - 1) {
+      if (i == componentNames.size() - 1) {
         // Look for a typedefTypeProp on the definition node of the last component.
-        Node def = parentObj.getPropertyDefSite(componentNames[i]);
-        JSType typedefType = def != null ? def.getTypedefTypeProp() : null;
-        if (typedefType != null) {
-          setReferencedAndResolvedType(typedefType, reporter);
+        Node def = parentObj.getPropertyDefSite(component);
+        if (resolveTypeFromNodeIfTypedef(def, reporter)) {
           return;
         }
       }
-      slotType = parentObj.getPropertyType(componentNames[i]);
+      slotType = parentObj.getPropertyType(component);
     }
 
     // Translate "constructor" types to "instance" types.
@@ -361,6 +387,72 @@ private void resolveViaProperties(ErrorReporter reporter) {
     }
   }
 
+  /**
+   * Resolves a named type by checking for the longest prefix that matches some Closure namespace,
+   * if any, then attempting to resolve via properties based on the type of the `exports` object in
+   * that namespace.
+   */
+  private boolean resolveViaClosureNamespace(ErrorReporter reporter) {
+    List<String> componentNames = Splitter.on('.').splitToList(reference);
+    if (componentNames.get(0).isEmpty()) {
+      return false;
+    }
+
+    StaticTypedSlot slot = resolutionScope.getSlot(componentNames.get(0));
+    // Skip types whose root component is defined in a local scope (not a global scope). Those will
+    // follow the normal resolution scheme. (For legacy compatibility reasons we don't check for
+    // global names that are the same as the module root).
+    if (slot != null && slot.getScope() != null && slot.getScope().getParentScope() != null) {
+      return false;
+    }
+
+    // Find the `exports` type of the longest prefix match of this namespace, if any. Then resolve
+    // it via property.
+    String prefix = reference;
+
+    for (int remainingComponentIndex = componentNames.size();
+        remainingComponentIndex > 0;
+        remainingComponentIndex--) {
+      ModuleSlot module = registry.getModuleSlot(prefix);
+      if (module == null) {
+        int lastDot = prefix.lastIndexOf(".");
+        if (lastDot >= 0) {
+          prefix = prefix.substring(0, lastDot);
+        }
+        continue;
+      }
+
+      if (module.isLegacyModule()) {
+        // Try to resolve this name via registry or properties.
+        return false;
+      } else {
+        // Always stop resolution here whether successful or not, instead of continuing with
+        // resolution via registry or via properties, to match legacy behavior.
+        resolveViaPropertyGivenSlot(
+            module.type(),
+            module.definitionNode(),
+            componentNames,
+            reporter,
+            remainingComponentIndex);
+        return true;
+      }
+    }
+    return false; // Keep trying to resolve this name.
+  }
+
+  /** Checks the given Node for a typedef annotation, resolving to that type if existent. */
+  private boolean resolveTypeFromNodeIfTypedef(Node node, ErrorReporter reporter) {
+    if (node == null) {
+      return false;
+    }
+    JSType typedefType = node.getTypedefTypeProp();
+    if (typedefType == null) {
+      return false;
+    }
+    setReferencedAndResolvedType(typedefType, reporter);
+    return true;
+  }
+
   private void setReferencedAndResolvedType(
       JSType type, ErrorReporter reporter) {
     if (nonNull) {

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

@@ -87,6 +87,6 @@ public JSDocInfo getJSDocInfo() {
 
   @Override
   public StaticTypedScope getScope() {
-    throw new UnsupportedOperationException();
+    return null;
   }
 }

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

@@ -4781,21 +4781,60 @@ public void testGoogModuleRequireAndRequireType_invalidDestructuringImport() {
 
   @Test
   public void testReferenceGoogModulesByType() {
-    // TODO(b/133501660): consider supporting referring to non-legacy modules in types, so that the
-    // following code would not emit an error. This is consistent with non-native module checks.
-    testWarning(
+    testSame(
         srcs(
             lines(
                 "goog.module('mod.A');", //
                 "class A {}",
                 "exports = A;"),
+            "var /** !mod.A */ a;"));
+
+    assertThat(globalScope.getVar("a")).hasJSTypeThat().toStringIsEqualTo("A");
+  }
+
+  @Test
+  public void testCannotReferenceGoogDeclareModuleIdByType() {
+    testWarning(
+        srcs(
+            lines(
+                "goog.declareModuleId('mod');", //
+                "export class A {}"),
             "var /** !mod.A */ a;"),
         warning(RhinoErrorReporter.UNRECOGNIZED_TYPE_ERROR));
   }
 
   @Test
   public void testReferenceGoogModuleNestedClassByType() {
-    // TODO(b/133501660): consider supporting referring to non-legacy modules in types.
+    testSame(
+        srcs(
+            lines(
+                "goog.module('mod.A');", //
+                "class A {}",
+                "A.B = class {};",
+                "A.B.C = class {};",
+                "exports = A;"),
+            "var /** !mod.A.B.C */ a;"));
+
+    assertThat(globalScope.getVar("a")).hasJSTypeThat().toStringIsEqualTo("A.B.C");
+  }
+
+  @Test
+  public void testReferenceGoogModuleNestedClassByType_exportDefaultTypedef() {
+    // Verify that the unusal resolution of @typedefs via Node annotation works correctly.
+    testSame(
+        srcs(
+            lines(
+                "goog.module('mod.A');", //
+                "/** @typedef {string} */",
+                "let MyTypedef;",
+                "exports = MyTypedef;"),
+            "var /** !mod.A */ a;"));
+
+    assertThat(globalScope.getVar("a")).hasJSTypeThat().isString();
+  }
+
+  @Test
+  public void testReferenceGoogModuleByType_shadowingLocalCausingResolutionFailure() {
     testWarning(
         srcs(
             lines(
@@ -4804,7 +4843,27 @@ public void testReferenceGoogModuleNestedClassByType() {
                 "A.B = class {};",
                 "A.B.C = class {};",
                 "exports = A;"),
-            "var /** !mod.A.B.C */ a;"),
+            lines(
+                "goog.module('mod.B');", //
+                "const mod = {};",
+                "var /** !mod.A.B.C */ a;")), // Resolves relative to the local `mod`.
+        warning(RhinoErrorReporter.UNRECOGNIZED_TYPE_ERROR));
+  }
+
+  @Test
+  public void testReferenceGoogModuleByType_googProvideExtensionCausesResolutionFailure() {
+    processClosurePrimitives = true;
+    testWarning(
+        srcs(
+            lines(
+                "goog.module('mod.A');", //
+                "class A {}",
+                "A.B = class {};",
+                "A.B.C = class {};",
+                "exports = A;"),
+            lines(
+                "goog.provide('mod.A.B');", //
+                "var /** !mod.A.B.C */ a;")), // Assumes this refers to the goog.provide
         warning(RhinoErrorReporter.UNRECOGNIZED_TYPE_ERROR));
   }