Merge pull request #38713 from slavapestov/lazy-assoc-type-list-deserialization

More efficient getAssociatedTypeMembers() on serialized ProtocolDecls

Author: slavapestov

include/swift/AST/Decl.h

@@ -511,7 +511,7 @@ class alignas(1 << DeclAlignInBits) Decl {
     IsComputingSemanticMembers : 1
   );
 
-  SWIFT_INLINE_BITFIELD_FULL(ProtocolDecl, NominalTypeDecl, 1+1+1+1+1+1+1+1+1+8+16,
+  SWIFT_INLINE_BITFIELD_FULL(ProtocolDecl, NominalTypeDecl, 1+1+1+1+1+1+1+1+1+1+1+8+16,
     /// Whether the \c RequiresClass bit is valid.
     RequiresClassValid : 1,
 
@@ -540,6 +540,12 @@ class alignas(1 << DeclAlignInBits) Decl {
     /// Whether we have a lazy-loaded requirement signature.
     HasLazyRequirementSignature : 1,
 
+    /// Whether we have computed the list of associated types.
+    HasAssociatedTypes : 1,
+
+    /// Whether we have a lazy-loaded list of associated types.
+    HasLazyAssociatedTypes : 1,
+
     : NumPadBits,
 
     /// If this is a compiler-known protocol, this will be a KnownProtocolKind
@@ -4094,6 +4100,7 @@ class ProtocolDecl final : public NominalTypeDecl {
   SourceLoc ProtocolLoc;
 
   ArrayRef<ProtocolDecl *> InheritedProtocols;
+  ArrayRef<AssociatedTypeDecl *> AssociatedTypes;
 
   struct {
     /// The superclass decl and a bit to indicate whether the
@@ -4192,7 +4199,7 @@ class ProtocolDecl final : public NominalTypeDecl {
   /// Retrieve the set of AssociatedTypeDecl members of this protocol; this
   /// saves loading the set of members in cases where there's no possibility of
   /// a protocol having nested types (ObjC protocols).
-  llvm::TinyPtrVector<AssociatedTypeDecl *> getAssociatedTypeMembers() const;
+  ArrayRef<AssociatedTypeDecl *> getAssociatedTypeMembers() const;
 
   /// Returns a protocol requirement with the given name, or nullptr if the
   /// name has multiple overloads, or no overloads at all.
@@ -4376,6 +4383,9 @@ class ProtocolDecl final : public NominalTypeDecl {
   void setLazyRequirementSignature(LazyMemberLoader *lazyLoader,
                                    uint64_t requirementSignatureData);
 
+  void setLazyAssociatedTypeMembers(LazyMemberLoader *lazyLoader,
+                                    uint64_t associatedTypesData);
+
 private:
   ArrayRef<Requirement> getCachedRequirementSignature() const;
 

include/swift/AST/LazyResolver.h

@@ -60,9 +60,10 @@ class LazyIterableDeclContextData : public LazyContextData {
 /// Context data for protocols.
 class LazyProtocolData : public LazyIterableDeclContextData {
 public:
-  /// The context data used for loading all of the members of the iterable
-  /// context.
+  /// The context data used for loading a requirement signature.
   uint64_t requirementSignatureData = 0;
+  /// The context data used for loading the list of associated types.
+  uint64_t associatedTypesData = 0;
 };
 
 /// A class that can lazily load members from a serialized format.
@@ -100,6 +101,11 @@ class alignas(void*) LazyMemberLoader {
   loadRequirementSignature(const ProtocolDecl *proto, uint64_t contextData,
                            SmallVectorImpl<Requirement> &requirements) = 0;
 
+  /// Loads the associated types of a protocol.
+  virtual void
+  loadAssociatedTypes(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,

lib/AST/Decl.cpp

@@ -4787,7 +4787,9 @@ ProtocolDecl::ProtocolDecl(DeclContext *DC, SourceLoc ProtocolLoc,
   Bits.ProtocolDecl.NumRequirementsInSignature = 0;
   Bits.ProtocolDecl.HasMissingRequirements = false;
   Bits.ProtocolDecl.KnownProtocol = 0;
-    setTrailingWhereClause(TrailingWhere);
+  Bits.ProtocolDecl.HasAssociatedTypes = 0;
+  Bits.ProtocolDecl.HasLazyAssociatedTypes = 0;
+  setTrailingWhereClause(TrailingWhere);
 }
 
 bool ProtocolDecl::isMarkerProtocol() const {
@@ -4806,26 +4808,40 @@ ArrayRef<ProtocolDecl *> ProtocolDecl::getInheritedProtocols() const {
                            {});
 }
 
-llvm::TinyPtrVector<AssociatedTypeDecl *>
+ArrayRef<AssociatedTypeDecl *>
 ProtocolDecl::getAssociatedTypeMembers() const {
-  llvm::TinyPtrVector<AssociatedTypeDecl *> result;
+  if (Bits.ProtocolDecl.HasAssociatedTypes)
+    return AssociatedTypes;
+
+  auto *self = const_cast<ProtocolDecl *>(this);
+  self->Bits.ProtocolDecl.HasAssociatedTypes = 1;
 
   // Clang-imported protocols never have associated types.
   if (hasClangNode())
-    return result;
+    return ArrayRef<AssociatedTypeDecl *>();
 
   // Deserialized @objc protocols never have associated types.
-  if (!getParentSourceFile() && isObjC())
-    return result;
+  if (getParentSourceFile() == nullptr && isObjC())
+    return ArrayRef<AssociatedTypeDecl *>();
 
-  // Find the associated type declarations.
-  for (auto member : getMembers()) {
-    if (auto ATD = dyn_cast<AssociatedTypeDecl>(member)) {
-      result.push_back(ATD);
+  SmallVector<AssociatedTypeDecl *, 2> result;
+  if (Bits.ProtocolDecl.HasLazyAssociatedTypes) {
+    auto &ctx = getASTContext();
+    auto contextData = static_cast<LazyProtocolData *>(
+        ctx.getOrCreateLazyContextData(this, nullptr));
+
+    contextData->loader->loadAssociatedTypes(
+        this, contextData->associatedTypesData, result);
+  } else {
+    for (auto member : getMembers()) {
+      if (auto ATD = dyn_cast<AssociatedTypeDecl>(member)) {
+        result.push_back(ATD);
+      }
     }
   }
 
-  return result;
+  self->AssociatedTypes = getASTContext().AllocateCopy(result);
+  return AssociatedTypes;
 }
 
 ValueDecl *ProtocolDecl::getSingleRequirement(DeclName name) const {
@@ -5190,6 +5206,18 @@ ProtocolDecl::setLazyRequirementSignature(LazyMemberLoader *lazyLoader,
     ++Stats->getFrontendCounters().NumLazyRequirementSignatures;
 }
 
+void
+ProtocolDecl::setLazyAssociatedTypeMembers(LazyMemberLoader *lazyLoader,
+                                           uint64_t associatedTypesData) {
+  assert(!Bits.ProtocolDecl.HasAssociatedTypes);
+  assert(!Bits.ProtocolDecl.HasLazyAssociatedTypes);
+
+  auto contextData = static_cast<LazyProtocolData *>(
+      getASTContext().getOrCreateLazyContextData(this, lazyLoader));
+  contextData->associatedTypesData = associatedTypesData;
+  Bits.ProtocolDecl.HasLazyAssociatedTypes = true;
+}
+
 ArrayRef<Requirement> ProtocolDecl::getCachedRequirementSignature() const {
   assert(RequirementSignature &&
          "getting requirement signature before computing it");

lib/AST/RequirementMachine/ProtocolGraph.h

@@ -35,7 +35,7 @@ struct ProtocolInfo {
   llvm::TinyPtrVector<const ProtocolDecl *> AllInherited;
 
   /// Associated types defined in the protocol itself.
-  llvm::TinyPtrVector<AssociatedTypeDecl *> AssociatedTypes;
+  ArrayRef<AssociatedTypeDecl *> AssociatedTypes;
 
   /// Associated types from all inherited protocols, not including duplicates or
   /// those defined in the protocol itself. Computed by
@@ -64,7 +64,7 @@ struct ProtocolInfo {
   }
 
   ProtocolInfo(ArrayRef<ProtocolDecl *> inherited,
-               llvm::TinyPtrVector<AssociatedTypeDecl *> &&types,
+               ArrayRef<AssociatedTypeDecl *> &&types,
                ArrayRef<Requirement> reqs)
     : Inherited(inherited),
       AssociatedTypes(types),

lib/ClangImporter/ImporterImpl.h

@@ -1423,6 +1423,11 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
     llvm_unreachable("unimplemented for ClangImporter");
   }
 
+  void loadAssociatedTypes(const ProtocolDecl *decl, uint64_t contextData,
+                           SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) override {
+    llvm_unreachable("unimplemented for ClangImporter");
+  }
+
   template <typename DeclTy, typename ...Targs>
   DeclTy *createDeclWithClangNode(ClangNode ClangN, AccessLevel access,
                                   Targs &&... Args) {

lib/Serialization/DeclTypeRecordNodes.def

@@ -123,6 +123,7 @@ DECL(OPAQUE_TYPE)
 DECL(PATTERN_BINDING)
 DECL(PROTOCOL)
 TRAILING_INFO(DEFAULT_WITNESS_TABLE)
+TRAILING_INFO(ASSOCIATED_TYPE)
 DECL(PREFIX_OPERATOR)
 DECL(POSTFIX_OPERATOR)
 DECL(INFIX_OPERATOR)

lib/Serialization/Deserialization.cpp

@@ -886,6 +886,36 @@ llvm::Error ModuleFile::readGenericRequirementsChecked(
   return llvm::Error::success();
 }
 
+void ModuleFile::readAssociatedTypes(
+                   SmallVectorImpl<AssociatedTypeDecl *> &assocTypes,
+                   llvm::BitstreamCursor &Cursor) {
+  using namespace decls_block;
+
+  BCOffsetRAII lastRecordOffset(Cursor);
+  SmallVector<uint64_t, 8> scratch;
+  StringRef blobData;
+
+  while (true) {
+    lastRecordOffset.reset();
+
+    llvm::BitstreamEntry entry =
+        fatalIfUnexpected(Cursor.advance(AF_DontPopBlockAtEnd));
+    if (entry.Kind != llvm::BitstreamEntry::Record)
+      break;
+
+    scratch.clear();
+    unsigned recordID = fatalIfUnexpected(
+        Cursor.readRecord(entry.ID, scratch, &blobData));
+    if (recordID != ASSOCIATED_TYPE)
+      break;
+
+    DeclID declID;
+    AssociatedTypeLayout::readRecord(scratch, declID);
+
+    assocTypes.push_back(cast<AssociatedTypeDecl>(getDecl(declID)));
+  }
+}
+
 /// Advances past any records that might be part of a requirement signature.
 static llvm::Error skipGenericRequirements(llvm::BitstreamCursor &Cursor) {
   using namespace decls_block;
@@ -919,6 +949,38 @@ static llvm::Error skipGenericRequirements(llvm::BitstreamCursor &Cursor) {
   return llvm::Error::success();
 }
 
+/// Advances past any lazy associated type member records.
+static llvm::Error skipAssociatedTypeMembers(llvm::BitstreamCursor &Cursor) {
+  using namespace decls_block;
+
+  BCOffsetRAII lastRecordOffset(Cursor);
+
+  while (true) {
+    Expected<llvm::BitstreamEntry> maybeEntry =
+        Cursor.advance(AF_DontPopBlockAtEnd);
+    if (!maybeEntry)
+      return maybeEntry.takeError();
+    llvm::BitstreamEntry entry = maybeEntry.get();
+    if (entry.Kind != llvm::BitstreamEntry::Record)
+      break;
+
+    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.
+      return llvm::Error::success();
+    }
+
+    lastRecordOffset.reset();
+  }
+  return llvm::Error::success();
+}
+
 GenericSignature ModuleFile::getGenericSignature(
     serialization::GenericSignatureID ID) {
   auto signature = getGenericSignatureChecked(ID);
@@ -3552,6 +3614,11 @@ class DeclDeserializer {
     if (llvm::Error Err = skipGenericRequirements(MF.DeclTypeCursor))
       MF.fatal(std::move(Err));
 
+    proto->setLazyAssociatedTypeMembers(&MF,
+                                        MF.DeclTypeCursor.GetCurrentBitNo());
+    if (llvm::Error Err = skipAssociatedTypeMembers(MF.DeclTypeCursor))
+      MF.fatal(std::move(Err));
+
     proto->setMemberLoader(&MF, MF.DeclTypeCursor.GetCurrentBitNo());
 
     return proto;
@@ -6558,6 +6625,14 @@ void ModuleFile::loadRequirementSignature(const ProtocolDecl *decl,
   readGenericRequirements(reqs, DeclTypeCursor);
 }
 
+void ModuleFile::loadAssociatedTypes(const ProtocolDecl *decl,
+                                     uint64_t contextData,
+                           SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) {
+  BCOffsetRAII restoreOffset(DeclTypeCursor);
+  fatalIfNotSuccess(DeclTypeCursor.JumpToBit(contextData));
+  readAssociatedTypes(assocTypes, DeclTypeCursor);
+}
+
 static Optional<ForeignErrorConvention::Kind>
 decodeRawStableForeignErrorConventionKind(uint8_t kind) {
   switch (kind) {

lib/Serialization/ModuleFile.h

@@ -389,6 +389,10 @@ class ModuleFile
   readGenericRequirementsChecked(SmallVectorImpl<Requirement> &requirements,
                                  llvm::BitstreamCursor &Cursor);
 
+  /// Read a list of associated type declarations in a protocol.
+  void readAssociatedTypes(SmallVectorImpl<AssociatedTypeDecl *> &assocTypes,
+                           llvm::BitstreamCursor &Cursor);
+
   /// Populates the protocol's default witness table.
   ///
   /// Returns true if there is an error.
@@ -709,6 +713,10 @@ class ModuleFile
   loadRequirementSignature(const ProtocolDecl *proto, uint64_t contextData,
                            SmallVectorImpl<Requirement> &requirements) override;
 
+  void
+  loadAssociatedTypes(const ProtocolDecl *proto, uint64_t contextData,
+                      SmallVectorImpl<AssociatedTypeDecl *> &assocTypes) override;
+
   Optional<StringRef> getGroupNameById(unsigned Id) const;
   Optional<StringRef> getSourceFileNameById(unsigned Id) const;
   Optional<StringRef> getGroupNameForDecl(const Decl *D) const;

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 = 620; // use @available(renamed:) for async alternative warning
+const uint16_t SWIFTMODULE_VERSION_MINOR = 621; // protocol associated type list
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -1612,6 +1612,11 @@ namespace decls_block {
     BCVBR<8>                     // alignment
   >;
 
+  using AssociatedTypeLayout = BCRecordLayout<
+    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<

lib/Serialization/Serialization.cpp

@@ -1346,6 +1346,19 @@ void Serializer::writeGenericRequirements(ArrayRef<Requirement> requirements,
   }
 }
 
+void Serializer::writeAssociatedTypes(ArrayRef<AssociatedTypeDecl *> assocTypes,
+                                      const std::array<unsigned, 256> &abbrCodes) {
+  using namespace decls_block;
+
+  auto assocTypeAbbrCode = abbrCodes[AssociatedTypeLayout::Code];
+
+  for (auto *assocType : assocTypes) {
+    AssociatedTypeLayout::emitRecord(
+        Out, ScratchRecord, assocTypeAbbrCode,
+        addDeclRef(assocType));
+  }
+}
+
 void Serializer::writeASTBlockEntity(GenericSignature sig) {
   using namespace decls_block;
 
@@ -3536,6 +3549,8 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
     writeGenericParams(proto->getGenericParams());
     S.writeGenericRequirements(
       proto->getRequirementSignature(), S.DeclTypeAbbrCodes);
+    S.writeAssociatedTypes(
+      proto->getAssociatedTypeMembers(), S.DeclTypeAbbrCodes);
     writeMembers(id, proto->getAllMembers(), true);
     writeDefaultWitnessTable(proto);
   }
@@ -4761,6 +4776,7 @@ void Serializer::writeAllDeclsAndTypes() {
   registerDeclTypeAbbr<OpaqueTypeLayout>();
   registerDeclTypeAbbr<PatternBindingLayout>();
   registerDeclTypeAbbr<ProtocolLayout>();
+  registerDeclTypeAbbr<AssociatedTypeLayout>();
   registerDeclTypeAbbr<DefaultWitnessTableLayout>();
   registerDeclTypeAbbr<PrefixOperatorLayout>();
   registerDeclTypeAbbr<PostfixOperatorLayout>();