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

There's a class of errors in Serialization called "circularity
issues", where declaration A in file A.swift depends on declaration B
in file B.swift, and B also depends on A. In some cases we can manage
to type-check each of these files individually due to the laziness of
'validateDecl', but then fail to merge the "partial modules" generated
from A.swift and B.swift to form a single swiftmodule for the library
(because deserialization is a little less lazy for some things). A
common case of this is when at least one of the declarations is
nested, in which case a lookup to find that declaration needs to load
all the members of the parent type. This gets even worse when the
nested type is defined in an extension.

This commit sidesteps that issue specifically for nested types by
creating a top-level, per-file table of nested types in the "partial
modules". When a type is in the same module, we can then look it up
/without/ importing all other members of the parent type.

The long-term solution is to allow accessing any members of a type
without having to load them all, something we should support not just
for module-merging while building a single target but when reading
from imported modules as well. This should improve both compile time
and memory usage, though I'm not sure to what extent. (Unfortunately,
too many things still depend on the whole members list being loaded.)

Because this is a new code path, I put in a switch to turn it off:
frontend flag -disable-serialization-nested-type-lookup-table

https://bugs.swift.org/browse/SR-3707 (and possibly others)

Author: jrose-apple

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);