AST: Cache result of ProtocolDecl::getAssociatedTypeMembers()

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.

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,17 +4808,23 @@ 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 *>();
+
+  SmallVector<AssociatedTypeDecl *, 2> result;
 
   // Find the associated type declarations.
   for (auto member : getMembers()) {
@@ -4825,7 +4833,8 @@ ProtocolDecl::getAssociatedTypeMembers() const {
     }
   }
 
-  return result;
+  self->AssociatedTypes = getASTContext().AllocateCopy(result);
+  return AssociatedTypes;
 }
 
 ValueDecl *ProtocolDecl::getSingleRequirement(DeclName name) const {

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