[Serialization] Improve extensions of nested types with the same name (swiftlang#7397)

Previously looking up an extension would result in all extensions for
types with the same name (nested or not) being deserialized; this
could even bring in base types that had not been deserialized yet. Add
in a string to distinguish an extension's base type; in the top-level
case this is just a module name, but for nested types it's a full
mangled name.

This is a little heavier than I'd like it to be, since it means we
mangle names and then throw them away, and since it means there's a
whole bunch of extra string data in the module just for uniquely
identifying a declaration. But it's correct, and does less work than
before, and fixes a circularity issue with a nested type A.B.A that
apparently used to work.

https://bugs.swift.org/browse/SR-3915

Author: jrose-apple

include/swift/Serialization/ModuleFile.h

@@ -312,6 +312,10 @@ class ModuleFile : public LazyMemberLoader {
   using SerializedDeclTable =
       llvm::OnDiskIterableChainedHashTable<DeclTableInfo>;
 
+  class ExtensionTableInfo;
+  using SerializedExtensionTable =
+      llvm::OnDiskIterableChainedHashTable<ExtensionTableInfo>;
+
   class LocalDeclTableInfo;
   using SerializedLocalDeclTable =
       llvm::OnDiskIterableChainedHashTable<LocalDeclTableInfo>;
@@ -323,9 +327,9 @@ class ModuleFile : public LazyMemberLoader {
   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<SerializedExtensionTable> ExtensionDecls;
   std::unique_ptr<SerializedLocalDeclTable> LocalTypeDecls;
   std::unique_ptr<SerializedNestedTypeDeclsTable> NestedTypeDecls;
 
@@ -446,6 +450,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::ExtensionTableLayout format.
+  std::unique_ptr<SerializedExtensionTable>
+  readExtensionTable(ArrayRef<uint64_t> fields, StringRef blobData);
+
   /// Read an on-disk local decl hash table stored in
   /// index_block::NestedTypeDeclsLayout format.
   std::unique_ptr<SerializedNestedTypeDeclsTable>

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 = 314; // Last change: change synthetic substitution serialization
+const uint16_t VERSION_MINOR = 315; // Last change: uniquely identify extensions
 
 using DeclID = PointerEmbeddedInt<unsigned, 31>;
 using DeclIDField = BCFixed<31>;
@@ -1476,6 +1476,12 @@ namespace index_block {
     BCBlob       // actual names
   >;
 
+  using ExtensionTableLayout = BCRecordLayout<
+    EXTENSIONS, // record ID
+    BCVBR<16>,  // table offset within the blob (see below)
+    BCBlob  // map from identifier strings to decl kinds / decl IDs
+  >;
+
   using ObjCMethodTableLayout = BCRecordLayout<
     OBJC_METHODS,  // record ID
     BCVBR<16>,     // table offset within the blob (see below)

lib/Serialization/Deserialization.cpp

@@ -23,12 +23,14 @@
 #include "swift/Parse/Parser.h"
 #include "swift/Serialization/BCReadingExtras.h"
 #include "swift/Serialization/SerializedModuleLoader.h"
+#include "swift/Basic/Defer.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "Serialization"
 
+STATISTIC(NumDeclsLoaded, "# of decls deserialized");
 STATISTIC(NumMemberListsLoaded,
           "# of nominals/extensions whose members were loaded");
 STATISTIC(NumNestedTypeShortcuts,
@@ -2163,6 +2165,7 @@ Decl *ModuleFile::getDecl(DeclID DID, Optional<DeclContext *> ForcedContext) {
   if (declOrOffset.isComplete())
     return declOrOffset;
 
+  ++NumDeclsLoaded;
   BCOffsetRAII restoreOffset(DeclTypeCursor);
   DeclTypeCursor.JumpToBit(declOrOffset);
   auto entry = DeclTypeCursor.advance();

lib/Serialization/ModuleFile.cpp

@@ -14,6 +14,7 @@
 #include "swift/Serialization/ModuleFormat.h"
 #include "swift/Subsystems.h"
 #include "swift/AST/AST.h"
+#include "swift/AST/ASTMangler.h"
 #include "swift/AST/ModuleLoader.h"
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/USRGeneration.h"
@@ -339,6 +340,66 @@ class ModuleFile::DeclTableInfo {
   }
 };
 
+/// Used to deserialize entries in the on-disk decl hash table.
+class ModuleFile::ExtensionTableInfo {
+  ModuleFile &File;
+public:
+  using internal_key_type = StringRef;
+  using external_key_type = Identifier;
+  using data_type = SmallVector<std::pair<StringRef, DeclID>, 8>;
+  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);
+  }
+
+  data_type ReadData(internal_key_type key, const uint8_t *data,
+                     unsigned length) {
+    data_type result;
+    const uint8_t *limit = data + length;
+    while (data < limit) {
+      DeclID offset = endian::readNext<uint32_t, little, unaligned>(data);
+
+      int32_t nameIDOrLength =
+          endian::readNext<int32_t, little, unaligned>(data);
+      StringRef moduleNameOrMangledBase;
+      if (nameIDOrLength < 0) {
+        const ModuleDecl *module = File.getModule(-nameIDOrLength);
+        moduleNameOrMangledBase = module->getName().str();
+      } else {
+        moduleNameOrMangledBase =
+            StringRef(reinterpret_cast<const char *>(data), nameIDOrLength);
+        data += nameIDOrLength;
+      }
+
+      result.push_back({ moduleNameOrMangledBase, offset });
+    }
+
+    return result;
+  }
+
+  explicit ExtensionTableInfo(ModuleFile &file) : File(file) {}
+};
+
 /// Used to deserialize entries in the on-disk decl hash table.
 class ModuleFile::LocalDeclTableInfo {
 public:
@@ -432,6 +493,17 @@ ModuleFile::readDeclTable(ArrayRef<uint64_t> fields, StringRef blobData) {
                                                 base + sizeof(uint32_t), base));
 }
 
+std::unique_ptr<ModuleFile::SerializedExtensionTable>
+ModuleFile::readExtensionTable(ArrayRef<uint64_t> fields, StringRef blobData) {
+  uint32_t tableOffset;
+  index_block::DeclListLayout::readRecord(fields, tableOffset);
+  auto base = reinterpret_cast<const uint8_t *>(blobData.data());
+
+  using OwnedTable = std::unique_ptr<SerializedExtensionTable>;
+  return OwnedTable(SerializedExtensionTable::Create(base + tableOffset,
+    base + sizeof(uint32_t), base, ExtensionTableInfo(*this)));
+}
+
 std::unique_ptr<ModuleFile::SerializedLocalDeclTable>
 ModuleFile::readLocalDeclTable(ArrayRef<uint64_t> fields, StringRef blobData) {
   uint32_t tableOffset;
@@ -569,7 +641,7 @@ bool ModuleFile::readIndexBlock(llvm::BitstreamCursor &cursor) {
         PrecedenceGroupDecls = readDeclTable(scratch, blobData);
         break;
       case index_block::EXTENSIONS:
-        ExtensionDecls = readDeclTable(scratch, blobData);
+        ExtensionDecls = readExtensionTable(scratch, blobData);
         break;
       case index_block::CLASS_MEMBERS:
         ClassMembersByName = readDeclTable(scratch, blobData);
@@ -1437,9 +1509,25 @@ void ModuleFile::loadExtensions(NominalTypeDecl *nominal) {
   if (iter == ExtensionDecls->end())
     return;
 
-  for (auto item : *iter) {
-    if (item.first == getKindForTable(nominal))
-      (void)getDecl(item.second);
+  if (nominal->hasAccessibility() &&
+      nominal->getEffectiveAccess() < Accessibility::Internal) {
+    if (nominal->getModuleScopeContext() != getFile())
+      return;
+  }
+
+  if (nominal->getParent()->isModuleScopeContext()) {
+    Identifier moduleName = nominal->getParentModule()->getName();
+    for (auto item : *iter) {
+      if (item.first == moduleName.str())
+        (void)getDecl(item.second);
+    }
+  } else {
+    std::string mangledName =
+        NewMangling::ASTMangler().mangleNominalType(nominal);
+    for (auto item : *iter) {
+      if (item.first == mangledName)
+        (void)getDecl(item.second);
+    }
   }
 }
 

lib/Serialization/Serialization.cpp

@@ -13,13 +13,13 @@
 #include "Serialization.h"
 #include "SILFormat.h"
 #include "swift/AST/AST.h"
+#include "swift/AST/ASTMangler.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/DiagnosticsCommon.h"
 #include "swift/AST/ForeignErrorConvention.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/Initializer.h"
 #include "swift/AST/LinkLibrary.h"
-#include "swift/AST/Mangle.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/ASTMangler.h"
 #include "swift/AST/RawComment.h"
@@ -112,6 +112,74 @@ namespace {
     }
   };
 
+  class ExtensionTableInfo {
+    serialization::Serializer &Serializer;
+    llvm::SmallDenseMap<const NominalTypeDecl *,std::string,4> MangledNameCache;
+
+  public:
+    explicit ExtensionTableInfo(serialization::Serializer &serializer)
+        : Serializer(serializer) {}
+
+    using key_type = Identifier;
+    using key_type_ref = key_type;
+    using data_type = Serializer::ExtensionTableData;
+    using data_type_ref = const data_type &;
+    using hash_value_type = uint32_t;
+    using offset_type = unsigned;
+
+    hash_value_type ComputeHash(key_type_ref key) {
+      assert(!key.empty());
+      return llvm::HashString(key.str());
+    }
+
+    int32_t getNameDataForBase(const NominalTypeDecl *nominal,
+                               StringRef *dataToWrite = nullptr) {
+      if (nominal->getDeclContext()->isModuleScopeContext())
+        return -Serializer.addModuleRef(nominal->getParentModule());
+
+      auto &mangledName = MangledNameCache[nominal];
+      if (mangledName.empty())
+        mangledName = NewMangling::ASTMangler().mangleNominalType(nominal);
+
+      assert(llvm::isUInt<31>(mangledName.size()));
+      if (dataToWrite)
+        *dataToWrite = mangledName;
+      return mangledName.size();
+    }
+
+    std::pair<unsigned, unsigned> EmitKeyDataLength(raw_ostream &out,
+                                                    key_type_ref key,
+                                                    data_type_ref data) {
+      uint32_t keyLength = key.str().size();
+      uint32_t dataLength = (sizeof(uint32_t) * 2) * data.size();
+      for (auto dataPair : data) {
+        int32_t nameData = getNameDataForBase(dataPair.first);
+        if (nameData > 0)
+          dataLength += nameData;
+      }
+      endian::Writer<little> writer(out);
+      writer.write<uint16_t>(keyLength);
+      writer.write<uint16_t>(dataLength);
+      return { keyLength, dataLength };
+    }
+
+    void EmitKey(raw_ostream &out, key_type_ref key, unsigned len) {
+      out << key.str();
+    }
+
+    void EmitData(raw_ostream &out, key_type_ref key, data_type_ref data,
+                  unsigned len) {
+      static_assert(declIDFitsIn32Bits(), "DeclID too large");
+      endian::Writer<little> writer(out);
+      for (auto entry : data) {
+        StringRef dataToWrite;
+        writer.write<uint32_t>(entry.second);
+        writer.write<int32_t>(getNameDataForBase(entry.first, &dataToWrite));
+        out << dataToWrite;
+      }
+    }
+  };
+
   class LocalDeclTableInfo {
   public:
     using key_type = std::string;
@@ -3592,6 +3660,32 @@ static void writeDeclTable(const index_block::DeclListLayout &DeclList,
   DeclList.emit(scratch, kind, tableOffset, hashTableBlob);
 }
 
+static void
+writeExtensionTable(const index_block::ExtensionTableLayout &ExtensionTable,
+                    const Serializer::ExtensionTable &table,
+                    Serializer &serializer) {
+  if (table.empty())
+    return;
+
+  SmallVector<uint64_t, 8> scratch;
+  llvm::SmallString<4096> hashTableBlob;
+  uint32_t tableOffset;
+  {
+    llvm::OnDiskChainedHashTableGenerator<ExtensionTableInfo> generator;
+    ExtensionTableInfo info{serializer};
+    for (auto &entry : table) {
+      generator.insert(entry.first, entry.second, info);
+    }
+
+    llvm::raw_svector_ostream blobStream(hashTableBlob);
+    // Make sure that no bucket is at offset 0
+    endian::Writer<little>(blobStream).write<uint32_t>(0);
+    tableOffset = generator.Emit(blobStream, info);
+  }
+
+  ExtensionTable.emit(scratch, tableOffset, hashTableBlob);
+}
+
 static void writeLocalDeclTable(const index_block::DeclListLayout &DeclList,
                                 index_block::RecordKind kind,
                                 LocalTypeHashTableGenerator &generator) {
@@ -4158,11 +4252,12 @@ static void collectInterestingNestedDeclarations(
 
 void Serializer::writeAST(ModuleOrSourceFile DC,
                           bool enableNestedTypeLookupTable) {
-  DeclTable topLevelDecls, extensionDecls, operatorDecls, operatorMethodDecls;
+  DeclTable topLevelDecls, operatorDecls, operatorMethodDecls;
   DeclTable precedenceGroupDecls;
   ObjCMethodTable objcMethods;
   NestedTypeDeclsTable nestedTypeDecls;
   LocalTypeHashTableGenerator localTypeGenerator;
+  ExtensionTable extensionDecls;
   bool hasLocalTypes = false;
 
   Optional<DeclID> entryPointClassID;
@@ -4196,7 +4291,7 @@ void Serializer::writeAST(ModuleOrSourceFile DC,
         Type extendedTy = ED->getExtendedType();
         const NominalTypeDecl *extendedNominal = extendedTy->getAnyNominal();
         extensionDecls[extendedNominal->getName()]
-          .push_back({ getKindForTable(extendedNominal), addDeclRef(D) });
+          .push_back({ extendedNominal, addDeclRef(D) });
       } else if (auto OD = dyn_cast<OperatorDecl>(D)) {
         operatorDecls[OD->getName()]
           .push_back({ getStableFixity(OD->getKind()), addDeclRef(D) });
@@ -4264,13 +4359,17 @@ void Serializer::writeAST(ModuleOrSourceFile DC,
     writeDeclTable(DeclList, index_block::TOP_LEVEL_DECLS, topLevelDecls);
     writeDeclTable(DeclList, index_block::OPERATORS, operatorDecls);
     writeDeclTable(DeclList, index_block::PRECEDENCE_GROUPS, precedenceGroupDecls);
-    writeDeclTable(DeclList, index_block::EXTENSIONS, extensionDecls);
     writeDeclTable(DeclList, index_block::CLASS_MEMBERS, ClassMembersByName);
     writeDeclTable(DeclList, index_block::OPERATOR_METHODS, operatorMethodDecls);
     if (hasLocalTypes)
       writeLocalDeclTable(DeclList, index_block::LOCAL_TYPE_DECLS,
                           localTypeGenerator);
 
+    if (!extensionDecls.empty()) {
+      index_block::ExtensionTableLayout ExtensionTable(Out);
+      writeExtensionTable(ExtensionTable, extensionDecls, *this);
+    }
+
     index_block::ObjCMethodTableLayout ObjCMethodTable(Out);
     writeObjCMethodTable(ObjCMethodTable, objcMethods);
 

lib/Serialization/Serialization.h

@@ -128,12 +128,16 @@ class Serializer {
 
   // In-memory representation of what will eventually be an on-disk
   // hash table of all defined Objective-C methods.
-  using ObjCMethodTable = llvm::DenseMap<ObjCSelector, ObjCMethodTableData>;
+  using ObjCMethodTable = llvm::MapVector<ObjCSelector, ObjCMethodTableData>;
 
   using NestedTypeDeclsData = SmallVector<std::pair<DeclID, DeclID>, 4>;
   // In-memory representation of what will eventually be an on-disk
   // hash table of all defined Objective-C methods.
-  using NestedTypeDeclsTable = llvm::DenseMap<Identifier, NestedTypeDeclsData>;
+  using NestedTypeDeclsTable = llvm::MapVector<Identifier, NestedTypeDeclsData>;
+
+  using ExtensionTableData =
+      SmallVector<std::pair<const NominalTypeDecl *, DeclID>, 4>;
+  using ExtensionTable = llvm::MapVector<Identifier, ExtensionTableData>;
 
 private:
   /// A map from identifiers to methods and properties with the given name.