Serialization: Use nested type lookup table for cross-module references also

Fixes a class of deserialization issues in the merge-modules
step.

The setup was the following:

- File A defines a typealias A whose underlying type is a nested
  type S of a type T, defined in a different module.

- File B defines an extension of T, and the extension member's
  type references A.

When deserializing A, we would proceed to deserialize the
underlying type, which references T.S. This would first deserialize
T and perform a name lookup to find S, which would deserialize all
members, including pulling in extensions. Deserialization of the
extension defined in file B would then fail, because the declaration
for A is not yet available.

We had a previous fix for these problems in the single-module case;
a per-file lookup table mapping mangled nested type names to
declarations, allowing a nested type to be deserialized without
pulling in all members and extensions of its parent type.

This patch generalizes the nested type lookup table allowing it to
be used to resolve cross-module references as well. Also, we were
only writing out the nested type table when serializing a partial
swiftmodule corresponding to a source file. Removing this check
allows the nested type table to be serialized for modules built
with WMO enabled as well, such as the standard library.

Fixes <rdar://problem/30976604> and
<https://bugs.swift.org/browse/SR-4208>.

Author: slavapestov

lib/Serialization/Deserialization.cpp

@@ -1338,55 +1338,55 @@ ModuleFile::resolveCrossReference(ModuleDecl *baseModule, uint32_t pathLen) {
     case XREF_TYPE_PATH_PIECE: {
       if (values.size() == 1) {
         ModuleDecl *module = values.front()->getModuleContext();
-        if (module == this->getAssociatedModule()) {
-          // Fast path for nested types in the same module.
-          IdentifierID IID;
-          bool onlyInNominal = false;
-          XRefTypePathPieceLayout::readRecord(scratch, IID, onlyInNominal);
-          Identifier memberName = getIdentifier(IID);
-          pathTrace.addValue(memberName);
-
-          llvm::PrettyStackTraceString message{
-              "If you're seeing a crash here, try passing "
-              "-Xfrontend -disable-serialization-nested-type-lookup-table"};
-
-          TypeDecl *nestedType = nullptr;
-          if (onlyInNominal) {
-            // Only look in the file containing the type itself.
-            const DeclContext *dc = values.front()->getDeclContext();
-            auto *serializedFile =
-                dyn_cast<SerializedASTFile>(dc->getModuleScopeContext());
-            if (serializedFile) {
-              nestedType =
-                  serializedFile->File.lookupNestedType(memberName,
-                                                        values.front());
-            }
-          } else {
-            // Fault in extensions, then ask every serialized AST in the module.
-            (void)cast<NominalTypeDecl>(values.front())->getExtensions();
-            for (FileUnit *file : module->getFiles()) {
-              if (file == getFile())
-                continue;
-              auto *serializedFile = dyn_cast<SerializedASTFile>(file);
-              if (!serializedFile)
-                continue;
-              nestedType =
-                  serializedFile->File.lookupNestedType(memberName,
-                                                        values.front());
-              if (nestedType)
-                break;
-            }
-          }
 
-          if (nestedType) {
-            values.clear();
-            values.push_back(nestedType);
-            ++NumNestedTypeShortcuts;
-            break;
+        // Fast path for nested types that avoids deserializing all
+        // members of the parent type.
+        IdentifierID IID;
+        bool onlyInNominal = false;
+        XRefTypePathPieceLayout::readRecord(scratch, IID, onlyInNominal);
+        Identifier memberName = getIdentifier(IID);
+        pathTrace.addValue(memberName);
+
+        llvm::PrettyStackTraceString message{
+          "If you're seeing a crash here, try passing "
+            "-Xfrontend -disable-serialization-nested-type-lookup-table"};
+
+        TypeDecl *nestedType = nullptr;
+        if (onlyInNominal) {
+          // Only look in the file containing the type itself.
+          const DeclContext *dc = values.front()->getDeclContext();
+          auto *serializedFile =
+            dyn_cast<SerializedASTFile>(dc->getModuleScopeContext());
+          if (serializedFile) {
+            nestedType =
+              serializedFile->File.lookupNestedType(memberName,
+                                                    values.front());
           }
+        } else {
+          // Fault in extensions, then ask every serialized AST in the module.
+          (void)cast<NominalTypeDecl>(values.front())->getExtensions();
+          for (FileUnit *file : module->getFiles()) {
+            if (file == getFile())
+              continue;
+            auto *serializedFile = dyn_cast<SerializedASTFile>(file);
+            if (!serializedFile)
+              continue;
+            nestedType =
+              serializedFile->File.lookupNestedType(memberName,
+                                                    values.front());
+            if (nestedType)
+              break;
+          }
+        }
 
-          pathTrace.removeLast();
+        if (nestedType) {
+          values.clear();
+          values.push_back(nestedType);
+          ++NumNestedTypeShortcuts;
+          break;
         }
+
+        pathTrace.removeLast();
       }
       LLVM_FALLTHROUGH;
     }

lib/Serialization/Serialization.cpp

@@ -4445,7 +4445,7 @@ void Serializer::writeAST(ModuleOrSourceFile DC,
     index_block::ObjCMethodTableLayout ObjCMethodTable(Out);
     writeObjCMethodTable(ObjCMethodTable, objcMethods);
 
-    if (DC.is<SourceFile *>() && enableNestedTypeLookupTable &&
+    if (enableNestedTypeLookupTable &&
         !nestedTypeDecls.empty()) {
       index_block::NestedTypeDeclsLayout NestedTypeDeclsTable(Out);
       writeNestedTypeDeclsTable(NestedTypeDeclsTable, nestedTypeDecls);

test/Serialization/Inputs/multi-module-nested-type-1.swift

@@ -0,0 +1,3 @@
+public struct X {
+  public typealias A = Int
+}

test/Serialization/Inputs/multi-module-nested-type-2.swift

@@ -0,0 +1,3 @@
+import other
+
+typealias B = X.A

test/Serialization/multi-module-nested-type-extension.swift

@@ -0,0 +1,19 @@
+// RUN: rm -rf %t && mkdir -p %t
+
+// Build the other module, which consists of a single source file.
+
+// RUN: %target-swift-frontend -emit-module -module-name other -o %t/multi-module-nested-type-1.swiftmodule -primary-file %S/Inputs/multi-module-nested-type-1.swift
+// RUN: %target-swift-frontend -emit-module -module-name other -o %t/other.swiftmodule %t/multi-module-nested-type-1.swiftmodule
+
+// Build this module, which consists of two source files.
+
+// RUN: %target-swift-frontend -emit-module -module-name me -o %t/multi-module-nested-type-2.swiftmodule -primary-file %S/Inputs/multi-module-nested-type-2.swift %s -I %t
+// RUN: %target-swift-frontend -emit-module -module-name me -o %t/multi-module-nested-type-3.swiftmodule %S/Inputs/multi-module-nested-type-2.swift -primary-file %s -I %t
+
+// RUN: %target-swift-frontend -emit-module -module-name me -o %t/me.swiftmodule %t/multi-module-nested-type-2.swiftmodule %t/multi-module-nested-type-3.swiftmodule -I %t
+
+import other
+
+extension X {
+  func takesB(_: B) {}
+}