[Serialization] Add a "nested types" lookup table for partial modules

include/swift/Frontend/FrontendOptions.h

@@ -194,6 +194,11 @@ class FrontendOptions {
   /// until the end of all files.
   bool DelayedFunctionBodyParsing = false;
 
+  /// If true, serialization encodes an extra lookup table for use in module-
+  /// merging when emitting partial modules (the per-file modules in a non-WMO
+  /// build).
+  bool EnableSerializationNestedTypeLookupTable = true;
+
   /// Indicates whether or not an import statement can pick up a Swift source
   /// file (as opposed to a module file).
   bool EnableSourceImport = false;

include/swift/Option/FrontendOptions.td

@@ -119,6 +119,10 @@ def autolink_library : Separate<["-"], "autolink-library">,
 def disable_typo_correction : Flag<["-"], "disable-typo-correction">,
   HelpText<"Disable typo correction">;
 
+def disable_serialization_nested_type_lookup_table :
+  Flag<["-"], "disable-serialization-nested-type-lookup-table">,
+  HelpText<"Force module merging to use regular lookups to find nested types">;
+
 } // end let Flags = [FrontendOption, NoDriverOption]
 
 def debug_crash_Group : OptionGroup<"<automatic crashing options>">;

include/swift/Serialization/ModuleFile.h

@@ -316,13 +316,18 @@ class ModuleFile : public LazyMemberLoader {
   using SerializedLocalDeclTable =
       llvm::OnDiskIterableChainedHashTable<LocalDeclTableInfo>;
 
+  class NestedTypeDeclsTableInfo;
+  using SerializedNestedTypeDeclsTable =
+      llvm::OnDiskIterableChainedHashTable<NestedTypeDeclsTableInfo>;
+
   std::unique_ptr<SerializedDeclTable> TopLevelDecls;
   std::unique_ptr<SerializedDeclTable> OperatorDecls;
   std::unique_ptr<SerializedDeclTable> PrecedenceGroupDecls;
   std::unique_ptr<SerializedDeclTable> ExtensionDecls;
   std::unique_ptr<SerializedDeclTable> ClassMembersByName;
   std::unique_ptr<SerializedDeclTable> OperatorMethodDecls;
   std::unique_ptr<SerializedLocalDeclTable> LocalTypeDecls;
+  std::unique_ptr<SerializedNestedTypeDeclsTable> NestedTypeDecls;
 
   class ObjCMethodTableInfo;
   using SerializedObjCMethodTable =
@@ -441,6 +446,11 @@ class ModuleFile : public LazyMemberLoader {
   std::unique_ptr<ModuleFile::SerializedObjCMethodTable>
   readObjCMethodTable(ArrayRef<uint64_t> fields, StringRef blobData);
 
+  /// Read an on-disk local decl hash table stored in
+  /// index_block::NestedTypeDeclsLayout format.
+  std::unique_ptr<SerializedNestedTypeDeclsTable>
+  readNestedTypeDeclsTable(ArrayRef<uint64_t> fields, StringRef blobData);
+
   /// Reads the index block, which contains global tables.
   ///
   /// Returns false if there was an error.
@@ -580,6 +590,10 @@ class ModuleFile : public LazyMemberLoader {
   /// Searches the module's local type decls for the given mangled name.
   TypeDecl *lookupLocalType(StringRef MangledName);
 
+  /// Searches the module's nested type decls table for the given member of
+  /// the given type.
+  TypeDecl *lookupNestedType(Identifier name, const ValueDecl *parent);
+
   /// Searches the module's operators for one with the given name and fixity.
   ///
   /// If none is found, returns null.

include/swift/Serialization/ModuleFormat.h

@@ -54,7 +54,7 @@ const uint16_t VERSION_MAJOR = 0;
 /// in source control, you should also update the comment to briefly
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
-const uint16_t VERSION_MINOR = 307; // Last change: layout requirements
+const uint16_t VERSION_MINOR = 308; // Last change: nested type table
 
 using DeclID = PointerEmbeddedInt<unsigned, 31>;
 using DeclIDField = BCFixed<31>;
@@ -1463,7 +1463,8 @@ namespace index_block {
     SIL_LAYOUT_OFFSETS,
 
     PRECEDENCE_GROUPS,
-
+    NESTED_TYPE_DECLS,
+
     LastRecordKind = PRECEDENCE_GROUPS,
   };
 
@@ -1494,6 +1495,12 @@ namespace index_block {
     BCBlob         // map from Objective-C selectors to methods with that selector
   >;
 
+  using NestedTypeDeclsLayout = BCRecordLayout<
+    NESTED_TYPE_DECLS, // record ID
+    BCVBR<16>,  // table offset within the blob (see below)
+    BCBlob  // map from identifier strings to decl kinds / decl IDs
+  >;
+
   using EntryPointLayout = BCRecordLayout<
     ENTRY_POINT,
     DeclIDField  // the ID of the main class; 0 if there was a main source file

include/swift/Serialization/SerializationOptions.h

@@ -36,6 +36,7 @@ namespace swift {
     ArrayRef<std::string> ExtraClangOptions;
 
     bool AutolinkForceLoad = false;
+    bool EnableNestedTypeLookupTable = false;
     bool SerializeAllSIL = false;
     bool SerializeOptionsForDebugging = false;
     bool IsSIB = false;

include/swift/Serialization/SerializedModuleLoader.h

@@ -104,6 +104,7 @@ class SerializedModuleLoader : public ModuleLoader {
 class SerializedASTFile final : public LoadedFile {
   friend class SerializedModuleLoader;
   friend class SerializedSILLoader;
+  friend class ModuleFile;
 
   ModuleFile &File;
   bool IsSIB;

lib/Frontend/CompilerInvocation.cpp

@@ -767,6 +767,8 @@ static bool ParseFrontendArgs(FrontendOptions &Opts, ArgList &Args,
   Opts.EnableSourceImport |= Args.hasArg(OPT_enable_source_import);
   Opts.ImportUnderlyingModule |= Args.hasArg(OPT_import_underlying_module);
   Opts.SILSerializeAll |= Args.hasArg(OPT_sil_serialize_all);
+  Opts.EnableSerializationNestedTypeLookupTable &=
+      !Args.hasArg(OPT_disable_serialization_nested_type_lookup_table);
 
   if (const Arg *A = Args.getLastArgNoClaim(OPT_import_objc_header)) {
     Opts.ImplicitObjCHeaderPath = A->getValue();

lib/FrontendTool/FrontendTool.cpp

@@ -666,6 +666,8 @@ static bool performCompile(std::unique_ptr<CompilerInstance> &Instance,
       serializationOpts.ModuleLinkName = opts.ModuleLinkName;
       serializationOpts.ExtraClangOptions =
           Invocation.getClangImporterOptions().ExtraArgs;
+      serializationOpts.EnableNestedTypeLookupTable =
+          opts.EnableSerializationNestedTypeLookupTable;
       if (!IRGenOpts.ForceLoadSymbolName.empty())
         serializationOpts.AutolinkForceLoad = true;
 

lib/Serialization/Deserialization.cpp

@@ -23,8 +23,17 @@
 #include "swift/ClangImporter/ClangImporter.h"
 #include "swift/Parse/Parser.h"
 #include "swift/Serialization/BCReadingExtras.h"
+#include "swift/Serialization/SerializedModuleLoader.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/Support/raw_ostream.h"
 
+#define DEBUG_TYPE "Serialization"
+
+STATISTIC(NumMemberListsLoaded,
+          "# of nominals/extensions whose members were loaded");
+STATISTIC(NumNestedTypeShortcuts,
+          "# of same-module nested types resolved without lookup");
+
 using namespace swift;
 using namespace swift::serialization;
 
@@ -226,6 +235,10 @@ namespace {
       path.push_back({ PathPiece::Kind::Unknown, kind });
     }
 
+    void removeLast() {
+      path.pop_back();
+    }
+
     void print(raw_ostream &os) const override {
       os << "Cross-reference to module '" << baseM.getName() << "'\n";
       for (auto &piece : path) {
@@ -1455,7 +1468,61 @@ Decl *ModuleFile::resolveCrossReference(ModuleDecl *M, uint32_t pathLen) {
     unsigned recordID = DeclTypeCursor.readRecord(entry.ID, scratch,
                                                   &blobData);
     switch (recordID) {
-    case XREF_TYPE_PATH_PIECE:
+    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;
+          }
+
+          pathTrace.removeLast();
+        }
+      }
+      SWIFT_FALLTHROUGH;
+    }
     case XREF_VALUE_PATH_PIECE:
     case XREF_INITIALIZER_PATH_PIECE: {
       TypeID TID = 0;
@@ -4254,6 +4321,7 @@ Type ModuleFile::getType(TypeID TID) {
 
 void ModuleFile::loadAllMembers(Decl *D, uint64_t contextData) {
   PrettyStackTraceDecl trace("loading members for", D);
+  ++NumMemberListsLoaded;
 
   BCOffsetRAII restoreOffset(DeclTypeCursor);
   DeclTypeCursor.JumpToBit(contextData);

lib/Serialization/ModuleFile.cpp

@@ -377,6 +377,50 @@ class ModuleFile::LocalDeclTableInfo {
   }
 };
 
+class ModuleFile::NestedTypeDeclsTableInfo {
+public:
+  using internal_key_type = StringRef;
+  using external_key_type = Identifier;
+  using data_type = SmallVector<std::pair<DeclID, DeclID>, 4>;
+  using hash_value_type = uint32_t;
+  using offset_type = unsigned;
+
+  internal_key_type GetInternalKey(external_key_type ID) {
+    return ID.str();
+  }
+
+  hash_value_type ComputeHash(internal_key_type key) {
+    return llvm::HashString(key);
+  }
+
+  static bool EqualKey(internal_key_type lhs, internal_key_type rhs) {
+    return lhs == rhs;
+  }
+
+  static std::pair<unsigned, unsigned> ReadKeyDataLength(const uint8_t *&data) {
+    unsigned keyLength = endian::readNext<uint16_t, little, unaligned>(data);
+    unsigned dataLength = endian::readNext<uint16_t, little, unaligned>(data);
+    return { keyLength, dataLength };
+  }
+
+  static internal_key_type ReadKey(const uint8_t *data, unsigned length) {
+    return StringRef(reinterpret_cast<const char *>(data), length);
+  }
+
+  static data_type ReadData(internal_key_type key, const uint8_t *data,
+                            unsigned length) {
+    data_type result;
+    while (length > 0) {
+      DeclID parentID = endian::readNext<uint32_t, little, unaligned>(data);
+      DeclID childID = endian::readNext<uint32_t, little, unaligned>(data);
+      result.push_back({ parentID, childID });
+      length -= sizeof(uint32_t) * 2;
+    }
+
+    return result;
+  }
+};
+
 std::unique_ptr<ModuleFile::SerializedDeclTable>
 ModuleFile::readDeclTable(ArrayRef<uint64_t> fields, StringRef blobData) {
   uint32_t tableOffset;
@@ -399,6 +443,18 @@ ModuleFile::readLocalDeclTable(ArrayRef<uint64_t> fields, StringRef blobData) {
     base + sizeof(uint32_t), base));
 }
 
+std::unique_ptr<ModuleFile::SerializedNestedTypeDeclsTable>
+ModuleFile::readNestedTypeDeclsTable(ArrayRef<uint64_t> fields,
+                                     StringRef blobData) {
+  uint32_t tableOffset;
+  index_block::NestedTypeDeclsLayout::readRecord(fields, tableOffset);
+  auto base = reinterpret_cast<const uint8_t *>(blobData.data());
+
+  using OwnedTable = std::unique_ptr<SerializedNestedTypeDeclsTable>;
+  return OwnedTable(SerializedNestedTypeDeclsTable::Create(base + tableOffset,
+      base + sizeof(uint32_t), base));
+}
+
 /// Used to deserialize entries in the on-disk Objective-C method table.
 class ModuleFile::ObjCMethodTableInfo {
 public:
@@ -531,6 +587,9 @@ bool ModuleFile::readIndexBlock(llvm::BitstreamCursor &cursor) {
       case index_block::LOCAL_TYPE_DECLS:
         LocalTypeDecls = readLocalDeclTable(scratch, blobData);
         break;
+      case index_block::NESTED_TYPE_DECLS:
+        NestedTypeDecls = readNestedTypeDeclsTable(scratch, blobData);
+        break;
       case index_block::LOCAL_DECL_CONTEXT_OFFSETS:
         assert(blobData.empty());
         LocalDeclContexts.assign(scratch.begin(), scratch.end());
@@ -1194,6 +1253,31 @@ TypeDecl *ModuleFile::lookupLocalType(StringRef MangledName) {
   return cast<TypeDecl>(getDecl((*iter).first));
 }
 
+TypeDecl *ModuleFile::lookupNestedType(Identifier name,
+                                       const ValueDecl *parent) {
+  PrettyModuleFileDeserialization stackEntry(*this);
+
+  if (!NestedTypeDecls)
+    return nullptr;
+
+  auto iter = NestedTypeDecls->find(name);
+  if (iter == NestedTypeDecls->end())
+    return nullptr;
+
+  auto data = *iter;
+  for (std::pair<DeclID, DeclID> entry : data) {
+    assert(entry.first);
+    auto declOrOffset = Decls[entry.first - 1];
+    if (!declOrOffset.isComplete())
+      continue;
+    if (declOrOffset != parent)
+      continue;
+    return cast<TypeDecl>(getDecl(entry.second));
+  }
+
+  return nullptr;
+}
+
 OperatorDecl *ModuleFile::lookupOperator(Identifier name, DeclKind fixity) {
   PrettyModuleFileDeserialization stackEntry(*this);