[Serialization] Improve extensions of nested types with the same name (#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

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 @@ public:
   }
 };
 
+/// 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);
+    }
   }
 }