Serialization: Round-trip primary associated type list

slavapestov · slavapestov · commit 6982b90c6e7c · 2022-04-16T22:36:29.000-04:00
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -527,7 +527,7 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
     IsComputingSemanticMembers : 1
   );
 
-  SWIFT_INLINE_BITFIELD_FULL(ProtocolDecl, NominalTypeDecl, 1+1+1+1+1+1+1+1+1+1+1+8,
+  SWIFT_INLINE_BITFIELD_FULL(ProtocolDecl, NominalTypeDecl, 1+1+1+1+1+1+1+1+1+1+1+1+8,
     /// Whether the \c RequiresClass bit is valid.
     RequiresClassValid : 1,
 
@@ -562,6 +562,9 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
     /// Whether we have a lazy-loaded list of associated types.
     HasLazyAssociatedTypes : 1,
 
+    /// Whether we have a lazy-loaded list of primary associated types.
+    HasLazyPrimaryAssociatedTypes : 1,
+
     : NumPadBits,
 
     /// If this is a compiler-known protocol, this will be a KnownProtocolKind
@@ -4450,6 +4453,10 @@ class ProtocolDecl final : public NominalTypeDecl {
     return Bits.ProtocolDecl.HasLazyRequirementSignature;
   }
 
+  bool hasLazyPrimaryAssociatedTypes() const {
+    return Bits.ProtocolDecl.HasLazyPrimaryAssociatedTypes;
+  }
+
   friend class SuperclassDeclRequest;
   friend class SuperclassTypeRequest;
   friend class StructuralRequirementsRequest;
@@ -4462,6 +4469,7 @@ class ProtocolDecl final : public NominalTypeDecl {
   friend class ExistentialConformsToSelfRequest;
   friend class ExistentialRequiresAnyRequest;
   friend class InheritedProtocolsRequest;
+  friend class PrimaryAssociatedTypesRequest;
   
 public:
   ProtocolDecl(DeclContext *DC, SourceLoc ProtocolLoc, SourceLoc NameLoc,
@@ -4686,6 +4694,9 @@ class ProtocolDecl final : public NominalTypeDecl {
   void setLazyAssociatedTypeMembers(LazyMemberLoader *lazyLoader,
                                     uint64_t associatedTypesData);
 
+  void setLazyPrimaryAssociatedTypeMembers(LazyMemberLoader *lazyLoader,
+                                           uint64_t associatedTypesData);
+
 public:
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Decl *D) {
diff --git a/include/swift/AST/LazyResolver.h b/include/swift/AST/LazyResolver.h
@@ -64,6 +64,8 @@ class LazyProtocolData : public LazyIterableDeclContextData {
   uint64_t requirementSignatureData = 0;
   /// The context data used for loading the list of associated types.
   uint64_t associatedTypesData = 0;
+  /// The context data used for loading the list of primary associated types.
+  uint64_t primaryAssociatedTypesData = 0;
 };
 
 /// A class that can lazily load members from a serialized format.
@@ -107,6 +109,11 @@ class alignas(void*) LazyMemberLoader {
   loadAssociatedTypes(const ProtocolDecl *proto, uint64_t contextData,
                       SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) = 0;
 
+  /// Loads the primary associated types of a protocol.
+  virtual void
+  loadPrimaryAssociatedTypes(const ProtocolDecl *proto, uint64_t contextData,
+                             SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) = 0;
+
   /// Returns the replaced decl for a given @_dynamicReplacement(for:) attribute.
   virtual ValueDecl *
   loadDynamicallyReplacedFunctionDecl(const DynamicReplacementAttr *DRA,
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -5585,8 +5585,8 @@ ProtocolDecl::setLazyRequirementSignature(LazyMemberLoader *lazyLoader,
 }
 
 void
-ProtocolDecl::setLazyAssociatedTypeMembers(LazyMemberLoader *lazyLoader,
-                                           uint64_t associatedTypesData) {
+ProtocolDecl::setLazyAssociatedTypeMembers(
+    LazyMemberLoader *lazyLoader, uint64_t associatedTypesData) {
   assert(!Bits.ProtocolDecl.HasAssociatedTypes);
   assert(!Bits.ProtocolDecl.HasLazyAssociatedTypes);
 
@@ -5596,6 +5596,17 @@ ProtocolDecl::setLazyAssociatedTypeMembers(LazyMemberLoader *lazyLoader,
   Bits.ProtocolDecl.HasLazyAssociatedTypes = true;
 }
 
+void
+ProtocolDecl::setLazyPrimaryAssociatedTypeMembers(
+    LazyMemberLoader *lazyLoader, uint64_t associatedTypesData) {
+  assert(!Bits.ProtocolDecl.HasLazyPrimaryAssociatedTypes);
+
+  auto contextData = static_cast<LazyProtocolData *>(
+      getASTContext().getOrCreateLazyContextData(this, lazyLoader));
+  contextData->primaryAssociatedTypesData = associatedTypesData;
+  Bits.ProtocolDecl.HasLazyPrimaryAssociatedTypes = true;
+}
+
 void ProtocolDecl::computeKnownProtocolKind() const {
   auto module = getModuleContext();
   if (module != module->getASTContext().getStdlibModule() &&
diff --git a/lib/ClangImporter/ImporterImpl.h b/lib/ClangImporter/ImporterImpl.h
@@ -1590,8 +1590,15 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
     llvm_unreachable("unimplemented for ClangImporter");
   }
 
-  void loadAssociatedTypes(const ProtocolDecl *decl, uint64_t contextData,
-                           SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) override {
+  void loadAssociatedTypes(
+      const ProtocolDecl *decl, uint64_t contextData,
+      SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) override {
+    llvm_unreachable("unimplemented for ClangImporter");
+  }
+
+  void loadPrimaryAssociatedTypes(
+      const ProtocolDecl *decl, uint64_t contextData,
+      SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) override {
     llvm_unreachable("unimplemented for ClangImporter");
   }
 
diff --git a/lib/Sema/TypeCheckDecl.cpp b/lib/Sema/TypeCheckDecl.cpp
@@ -703,6 +703,18 @@ ArrayRef<AssociatedTypeDecl *>
 PrimaryAssociatedTypesRequest::evaluate(Evaluator &evaluator,
                                         ProtocolDecl *decl) const {
   SmallVector<AssociatedTypeDecl *, 2> assocTypes;
+
+  if (decl->hasLazyPrimaryAssociatedTypes()) {
+    auto &ctx = decl->getASTContext();
+    auto contextData = static_cast<LazyProtocolData *>(
+        ctx.getOrCreateLazyContextData(decl, nullptr));
+
+    contextData->loader->loadPrimaryAssociatedTypes(
+        decl, contextData->primaryAssociatedTypesData, assocTypes);
+
+    return decl->getASTContext().AllocateCopy(assocTypes);
+  }
+
   llvm::SmallDenseSet<Identifier, 2> assocTypeNames;
 
   for (auto pair : decl->getPrimaryAssociatedTypeNames()) {
diff --git a/lib/Serialization/DeclTypeRecordNodes.def b/lib/Serialization/DeclTypeRecordNodes.def
@@ -128,6 +128,7 @@ DECL(PATTERN_BINDING)
 DECL(PROTOCOL)
 TRAILING_INFO(DEFAULT_WITNESS_TABLE)
 TRAILING_INFO(ASSOCIATED_TYPE)
+TRAILING_INFO(PRIMARY_ASSOCIATED_TYPE)
 DECL(PREFIX_OPERATOR)
 DECL(POSTFIX_OPERATOR)
 DECL(INFIX_OPERATOR)
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp
@@ -991,42 +991,37 @@ void ModuleFile::readAssociatedTypes(
   }
 }
 
-/// Advances past any records that might be part of a protocol requirement
-/// signature, which consists of generic requirements together with protocol
-/// typealias records.
-static llvm::Error skipRequirementSignature(llvm::BitstreamCursor &Cursor) {
+void ModuleFile::readPrimaryAssociatedTypes(
+                   SmallVectorImpl<AssociatedTypeDecl *> &assocTypes,
+                   llvm::BitstreamCursor &Cursor) {
   using namespace decls_block;
 
   BCOffsetRAII lastRecordOffset(Cursor);
+  SmallVector<uint64_t, 8> scratch;
+  StringRef blobData;
 
   while (true) {
-    Expected<llvm::BitstreamEntry> maybeEntry =
-        Cursor.advance(AF_DontPopBlockAtEnd);
-    if (!maybeEntry)
-      return maybeEntry.takeError();
-    llvm::BitstreamEntry entry = maybeEntry.get();
+    lastRecordOffset.reset();
+
+    llvm::BitstreamEntry entry =
+        fatalIfUnexpected(Cursor.advance(AF_DontPopBlockAtEnd));
     if (entry.Kind != llvm::BitstreamEntry::Record)
       break;
 
-    Expected<unsigned> maybeRecordID = Cursor.skipRecord(entry.ID);
-    if (!maybeRecordID)
-      return maybeRecordID.takeError();
-    switch (maybeRecordID.get()) {
-    case REQUIREMENT_SIGNATURE:
+    scratch.clear();
+    unsigned recordID = fatalIfUnexpected(
+        Cursor.readRecord(entry.ID, scratch, &blobData));
+    if (recordID != PRIMARY_ASSOCIATED_TYPE)
       break;
 
-    default:
-      // This record is not a generic requirement.
-      return llvm::Error::success();
-    }
+    DeclID declID;
+    PrimaryAssociatedTypeLayout::readRecord(scratch, declID);
 
-    lastRecordOffset.reset();
+    assocTypes.push_back(cast<AssociatedTypeDecl>(getDecl(declID)));
   }
-  return llvm::Error::success();
 }
 
-/// Advances past any lazy associated type member records.
-static llvm::Error skipAssociatedTypeMembers(llvm::BitstreamCursor &Cursor) {
+static llvm::Error skipRecords(llvm::BitstreamCursor &Cursor, unsigned kind) {
   using namespace decls_block;
 
   BCOffsetRAII lastRecordOffset(Cursor);
@@ -1043,20 +1038,39 @@ static llvm::Error skipAssociatedTypeMembers(llvm::BitstreamCursor &Cursor) {
     Expected<unsigned> maybeRecordID = Cursor.skipRecord(entry.ID);
     if (!maybeRecordID)
       return maybeRecordID.takeError();
-    switch (maybeRecordID.get()) {
-    case ASSOCIATED_TYPE:
-      break;
-
-    default:
-      // This record is not an associated type.
+    if (maybeRecordID.get() != kind)
       return llvm::Error::success();
-    }
 
     lastRecordOffset.reset();
   }
+
   return llvm::Error::success();
 }
 
+/// Advances past any records that might be part of a protocol requirement
+/// signature, which consists of generic requirements together with protocol
+/// typealias records.
+static llvm::Error skipRequirementSignature(llvm::BitstreamCursor &Cursor) {
+  using namespace decls_block;
+
+  return skipRecords(Cursor, REQUIREMENT_SIGNATURE);
+}
+
+/// Advances past any lazy associated type member records.
+static llvm::Error skipAssociatedTypeMembers(llvm::BitstreamCursor &Cursor) {
+  using namespace decls_block;
+
+  return skipRecords(Cursor, ASSOCIATED_TYPE);
+}
+
+/// Advances past any lazy primary associated type member records.
+static llvm::Error skipPrimaryAssociatedTypeMembers(
+    llvm::BitstreamCursor &Cursor) {
+  using namespace decls_block;
+
+  return skipRecords(Cursor, PRIMARY_ASSOCIATED_TYPE);
+}
+
 GenericSignature ModuleFile::getGenericSignature(
     serialization::GenericSignatureID ID) {
   auto signature = getGenericSignatureChecked(ID);
@@ -3749,16 +3763,21 @@ class DeclDeserializer {
       proto->setImplicit();
     proto->setIsObjC(isObjC);
 
-    proto->setLazyRequirementSignature(&MF,
-                                       MF.DeclTypeCursor.GetCurrentBitNo());
+    proto->setLazyRequirementSignature(
+        &MF, MF.DeclTypeCursor.GetCurrentBitNo());
     if (llvm::Error Err = skipRequirementSignature(MF.DeclTypeCursor))
       MF.fatal(std::move(Err));
 
-    proto->setLazyAssociatedTypeMembers(&MF,
-                                        MF.DeclTypeCursor.GetCurrentBitNo());
+    proto->setLazyAssociatedTypeMembers(
+        &MF, MF.DeclTypeCursor.GetCurrentBitNo());
     if (llvm::Error Err = skipAssociatedTypeMembers(MF.DeclTypeCursor))
       MF.fatal(std::move(Err));
 
+    proto->setLazyPrimaryAssociatedTypeMembers(
+        &MF, MF.DeclTypeCursor.GetCurrentBitNo());
+    if (llvm::Error Err = skipPrimaryAssociatedTypeMembers(MF.DeclTypeCursor))
+      MF.fatal(std::move(Err));
+
     proto->setMemberLoader(&MF, MF.DeclTypeCursor.GetCurrentBitNo());
 
     return proto;
@@ -6987,6 +7006,14 @@ void ModuleFile::loadAssociatedTypes(const ProtocolDecl *decl,
   readAssociatedTypes(assocTypes, DeclTypeCursor);
 }
 
+void ModuleFile::loadPrimaryAssociatedTypes(const ProtocolDecl *decl,
+                                            uint64_t contextData,
+                           SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) {
+  BCOffsetRAII restoreOffset(DeclTypeCursor);
+  fatalIfNotSuccess(DeclTypeCursor.JumpToBit(contextData));
+  readPrimaryAssociatedTypes(assocTypes, DeclTypeCursor);
+}
+
 static Optional<ForeignErrorConvention::Kind>
 decodeRawStableForeignErrorConventionKind(uint8_t kind) {
   switch (kind) {
diff --git a/lib/Serialization/ModuleFile.h b/lib/Serialization/ModuleFile.h
@@ -414,8 +414,14 @@ class ModuleFile
                                 llvm::BitstreamCursor &Cursor);
 
   /// Read a list of associated type declarations in a protocol.
-  void readAssociatedTypes(SmallVectorImpl<AssociatedTypeDecl *> &assocTypes,
-                           llvm::BitstreamCursor &Cursor);
+  void readAssociatedTypes(
+      SmallVectorImpl<AssociatedTypeDecl *> &assocTypes,
+      llvm::BitstreamCursor &Cursor);
+
+  /// Read a list of primary associated type declarations in a protocol.
+  void readPrimaryAssociatedTypes(
+      SmallVectorImpl<AssociatedTypeDecl *> &assocTypes,
+      llvm::BitstreamCursor &Cursor);
 
   /// Populates the protocol's default witness table.
   ///
@@ -747,8 +753,14 @@ class ModuleFile
                            SmallVectorImpl<ProtocolTypeAlias> &typeAliases) override;
 
   void
-  loadAssociatedTypes(const ProtocolDecl *proto, uint64_t contextData,
-                      SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) override;
+  loadAssociatedTypes(
+      const ProtocolDecl *proto, uint64_t contextData,
+      SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) override;
+
+  void
+  loadPrimaryAssociatedTypes(
+      const ProtocolDecl *proto, uint64_t contextData,
+      SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) override;
 
   Optional<StringRef> getGroupNameById(unsigned Id) const;
   Optional<StringRef> getSourceFileNameById(unsigned Id) const;
diff --git a/lib/Serialization/ModuleFormat.h b/lib/Serialization/ModuleFormat.h
@@ -56,7 +56,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// 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.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 684; // Cond. available underlying types for opaque
+const uint16_t SWIFTMODULE_VERSION_MINOR = 685; // Primary associated types
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -1657,6 +1657,11 @@ namespace decls_block {
     DeclIDField                  // associated type decl
   >;
 
+  using PrimaryAssociatedTypeLayout = BCRecordLayout<
+    PRIMARY_ASSOCIATED_TYPE,
+    DeclIDField                  // associated type decl
+  >;
+
   /// Specifies the private discriminator string for a private declaration. This
   /// identifies the declaration's original source file in some opaque way.
   using PrivateDiscriminatorLayout = BCRecordLayout<
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
@@ -1396,7 +1396,8 @@ void Serializer::serializeGenericRequirements(
   }
 }
 
-void Serializer::writeAssociatedTypes(ArrayRef<AssociatedTypeDecl *> assocTypes) {
+void Serializer::writeAssociatedTypes(
+    ArrayRef<AssociatedTypeDecl *> assocTypes) {
   using namespace decls_block;
 
   auto assocTypeAbbrCode = DeclTypeAbbrCodes[AssociatedTypeLayout::Code];
@@ -1408,6 +1409,19 @@ void Serializer::writeAssociatedTypes(ArrayRef<AssociatedTypeDecl *> assocTypes)
   }
 }
 
+void Serializer::writePrimaryAssociatedTypes(
+    ArrayRef<AssociatedTypeDecl *> assocTypes) {
+  using namespace decls_block;
+
+  auto assocTypeAbbrCode = DeclTypeAbbrCodes[PrimaryAssociatedTypeLayout::Code];
+
+  for (auto *assocType : assocTypes) {
+    PrimaryAssociatedTypeLayout::emitRecord(
+        Out, ScratchRecord, assocTypeAbbrCode,
+        addDeclRef(assocType));
+  }
+}
+
 void Serializer::writeRequirementSignature(
     const RequirementSignature &requirementSig) {
   using namespace decls_block;
@@ -3706,6 +3720,7 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
     writeGenericParams(proto->getGenericParams());
     S.writeRequirementSignature(proto->getRequirementSignature());
     S.writeAssociatedTypes(proto->getAssociatedTypeMembers());
+    S.writePrimaryAssociatedTypes(proto->getPrimaryAssociatedTypes());
     writeMembers(id, proto->getAllMembers(), true);
     writeDefaultWitnessTable(proto);
   }
@@ -5044,6 +5059,7 @@ void Serializer::writeAllDeclsAndTypes() {
   registerDeclTypeAbbr<PatternBindingLayout>();
   registerDeclTypeAbbr<ProtocolLayout>();
   registerDeclTypeAbbr<AssociatedTypeLayout>();
+  registerDeclTypeAbbr<PrimaryAssociatedTypeLayout>();
   registerDeclTypeAbbr<DefaultWitnessTableLayout>();
   registerDeclTypeAbbr<PrefixOperatorLayout>();
   registerDeclTypeAbbr<PostfixOperatorLayout>();
diff --git a/lib/Serialization/Serialization.h b/lib/Serialization/Serialization.h
diff --git a/test/Serialization/Inputs/parameterized_protocol_other.swift b/test/Serialization/Inputs/parameterized_protocol_other.swift
diff --git a/test/Serialization/parameterized_protocol.swift b/test/Serialization/parameterized_protocol.swift
diff --git a/test/type/parameterized_protocol.swift b/test/type/parameterized_protocol.swift
