[AST] NameLookup: Add a request to gather protocol requirements

Author: xedin

include/swift/AST/Decl.h

@@ -532,7 +532,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+1+8,
+  SWIFT_INLINE_BITFIELD_FULL(ProtocolDecl, NominalTypeDecl, 1+1+1+1+1+1+1+1+1+1+1+1+1+8,
     /// Whether the \c RequiresClass bit is valid.
     RequiresClassValid : 1,
 
@@ -570,6 +570,9 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
     /// Whether we have a lazy-loaded list of primary associated types.
     HasLazyPrimaryAssociatedTypes : 1,
 
+    /// Whether we've computed the protocol requirements list yet.
+    ProtocolRequirementsValid : 1,
+
     : NumPadBits,
 
     /// If this is a compiler-known protocol, this will be a KnownProtocolKind
@@ -4579,6 +4582,7 @@ class ProtocolDecl final : public NominalTypeDecl {
   ArrayRef<PrimaryAssociatedTypeName> PrimaryAssociatedTypeNames;
   ArrayRef<ProtocolDecl *> InheritedProtocols;
   ArrayRef<AssociatedTypeDecl *> AssociatedTypes;
+  ArrayRef<ValueDecl *> ProtocolRequirements;
 
   struct {
     /// The superclass decl and a bit to indicate whether the
@@ -4660,6 +4664,7 @@ class ProtocolDecl final : public NominalTypeDecl {
   friend class ExistentialRequiresAnyRequest;
   friend class InheritedProtocolsRequest;
   friend class PrimaryAssociatedTypesRequest;
+  friend class ProtocolRequirementsRequest;
 
 public:
   ProtocolDecl(DeclContext *DC, SourceLoc ProtocolLoc, SourceLoc NameLoc,
@@ -4703,6 +4708,10 @@ class ProtocolDecl final : public NominalTypeDecl {
   /// parametrized protocol type of the form SomeProtocol<Arg1, Arg2...>.
   ArrayRef<AssociatedTypeDecl *> getPrimaryAssociatedTypes() const;
 
+  /// Returns the list of all requirements (associated type and value)
+  /// associated with this protocol.
+  ArrayRef<ValueDecl *> getProtocolRequirements() const;
+
   /// Returns a protocol requirement with the given name, or nullptr if the
   /// name has multiple overloads, or no overloads at all.
   ValueDecl *getSingleRequirement(DeclName name) const;
@@ -4776,6 +4785,13 @@ class ProtocolDecl final : public NominalTypeDecl {
     Bits.ProtocolDecl.InheritedProtocolsValid = true;
   }
 
+  bool areProtocolRequirementsValid() const {
+    return Bits.ProtocolDecl.ProtocolRequirementsValid;
+  }
+  void setProtocolRequirementsValid() {
+    Bits.ProtocolDecl.ProtocolRequirementsValid = true;
+  }
+
 public:
   /// If this is known to be a compiler-known protocol, returns the kind.
   /// Otherwise returns None.

include/swift/AST/NameLookupRequests.h

@@ -196,6 +196,26 @@ class InheritedProtocolsRequest
                            ArrayRef<ProtocolDecl *> result) const;
 };
 
+class ProtocolRequirementsRequest
+    : public SimpleRequest<ProtocolRequirementsRequest,
+                           ArrayRef<ValueDecl *>(ProtocolDecl *),
+                           RequestFlags::SeparatelyCached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  ArrayRef<ValueDecl *> evaluate(Evaluator &, ProtocolDecl *) const;
+
+public:
+  // Caching
+  bool isCached() const { return true; }
+  Optional<ArrayRef<ValueDecl *>> getCachedResult() const;
+  void cacheResult(ArrayRef<ValueDecl *> value) const;
+};
+
 /// Requests whether or not this class has designated initializers that are
 /// not public or @usableFromInline.
 class HasMissingDesignatedInitializersRequest :

include/swift/AST/NameLookupTypeIDZone.def

@@ -47,6 +47,9 @@ SWIFT_REQUEST(NameLookup, InheritedDeclsReferencedRequest,
 SWIFT_REQUEST(NameLookup, InheritedProtocolsRequest,
               ArrayRef<ProtocolDecl *>(ProtocolDecl *), SeparatelyCached,
               NoLocationInfo)
+SWIFT_REQUEST(NameLookup, ProtocolRequirementsRequest,
+              ArrayRef<ValueDecl *>(ProtocolDecl *), SeparatelyCached,
+              NoLocationInfo)
 SWIFT_REQUEST(NameLookup, LookupConformanceInModuleRequest,
               ProtocolConformanceRef(LookupConformanceDescriptor),
               Uncached, NoLocationInfo)

lib/AST/Decl.cpp

@@ -5494,6 +5494,7 @@ ProtocolDecl::ProtocolDecl(DeclContext *DC, SourceLoc ProtocolLoc,
   Bits.ProtocolDecl.KnownProtocol = 0;
   Bits.ProtocolDecl.HasAssociatedTypes = 0;
   Bits.ProtocolDecl.HasLazyAssociatedTypes = 0;
+  Bits.ProtocolDecl.ProtocolRequirementsValid = false;
   setTrailingWhereClause(TrailingWhere);
 }
 
@@ -5533,7 +5534,7 @@ ProtocolDecl::getAssociatedTypeMembers() const {
     contextData->loader->loadAssociatedTypes(
         this, contextData->associatedTypesData, result);
   } else {
-    for (auto member : getABIMembers()) {
+    for (auto member : getProtocolRequirements()) {
       if (auto ATD = dyn_cast<AssociatedTypeDecl>(member)) {
         result.push_back(ATD);
       }
@@ -5544,6 +5545,12 @@ ProtocolDecl::getAssociatedTypeMembers() const {
   return AssociatedTypes;
 }
 
+ArrayRef<ValueDecl *> ProtocolDecl::getProtocolRequirements() const {
+  auto *mutableSelf = const_cast<ProtocolDecl *>(this);
+  return evaluateOrDefault(getASTContext().evaluator,
+                           ProtocolRequirementsRequest{mutableSelf}, {});
+}
+
 ValueDecl *ProtocolDecl::getSingleRequirement(DeclName name) const {
   auto results = const_cast<ProtocolDecl *>(this)->lookupDirect(name);
   ValueDecl *result = nullptr;

lib/AST/NameLookup.cpp

@@ -2723,6 +2723,20 @@ InheritedProtocolsRequest::evaluate(Evaluator &evaluator,
   return PD->getASTContext().AllocateCopy(result);
 }
 
+ArrayRef<ValueDecl *>
+ProtocolRequirementsRequest::evaluate(Evaluator &evaluator,
+                                      ProtocolDecl *PD) const {
+  SmallVector<ValueDecl *, 4> requirements;
+
+  for (auto *member : PD->getABIMembers()) {
+    auto *VD = dyn_cast<ValueDecl>(member);
+    if (VD && VD->isProtocolRequirement())
+      requirements.push_back(VD);
+  }
+
+  return PD->getASTContext().AllocateCopy(requirements);
+}
+
 NominalTypeDecl *
 ExtendedNominalRequest::evaluate(Evaluator &evaluator,
                                  ExtensionDecl *ext) const {

lib/AST/NameLookupRequests.cpp

@@ -111,6 +111,21 @@ void InheritedProtocolsRequest::writeDependencySink(
   }
 }
 
+Optional<ArrayRef<ValueDecl *>>
+ProtocolRequirementsRequest::getCachedResult() const {
+  auto proto = std::get<0>(getStorage());
+  if (!proto->areProtocolRequirementsValid())
+    return None;
+
+  return proto->ProtocolRequirements;
+}
+
+void ProtocolRequirementsRequest::cacheResult(ArrayRef<ValueDecl *> PDs) const {
+  auto proto = std::get<0>(getStorage());
+  proto->ProtocolRequirements = PDs;
+  proto->setProtocolRequirementsValid();
+}
+
 //----------------------------------------------------------------------------//
 // Missing designated initializers computation
 //----------------------------------------------------------------------------//

lib/Sema/TypeCheckProtocol.cpp

@@ -1817,10 +1817,7 @@ void MultiConformanceChecker::checkAllConformances() {
       return *checker;
     };
 
-    for (auto member : proto->getABIMembers()) {
-      auto req = dyn_cast<ValueDecl>(member);
-      if (!req || !req->isProtocolRequirement()) continue;
-
+    for (auto *req : proto->getProtocolRequirements()) {
       // If the requirement is unsatisfied, we might want to warn
       // about near misses; record it.
       if (isUnsatisfiedReq(getChecker(), conformance, req)) {
@@ -5124,19 +5121,11 @@ hasInvalidTypeInConformanceContext(const ValueDecl *requirement,
 }
 
 void ConformanceChecker::resolveValueWitnesses() {
-  for (auto member : Proto->getABIMembers()) {
-    auto requirement = dyn_cast<ValueDecl>(member);
-    if (!requirement)
-      continue;
-
+  for (auto *requirement : Proto->getProtocolRequirements()) {
     // Associated type requirements handled elsewhere.
     if (isa<TypeDecl>(requirement))
       continue;
 
-    // Type aliases don't have requirements themselves.
-    if (!requirement->isProtocolRequirement())
-      continue;
-
     /// Local function to finalize the witness.
     auto finalizeWitness = [&] {
       // Find the witness.
@@ -5467,7 +5456,7 @@ ConformanceChecker::getObjCRequirements(ObjCMethodKey key) {
 
   // Fill in the data structure if we haven't done so yet.
   if (!computedObjCMethodRequirements) {
-    for (auto requirement : proto->getABIMembers()) {
+    for (auto requirement : proto->getProtocolRequirements()) {
       auto funcRequirement = dyn_cast<AbstractFunctionDecl>(requirement);
       if (!funcRequirement)
         continue;
@@ -7129,15 +7118,11 @@ void TypeChecker::inferDefaultWitnesses(ProtocolDecl *proto) {
     return {defaultType, defaultedAssocType};
   };
 
-  for (auto *requirement : proto->getABIMembers()) {
+  for (auto *requirement : proto->getProtocolRequirements()) {
     if (requirement->isInvalid())
       continue;
 
-    auto *valueDecl = dyn_cast<ValueDecl>(requirement);
-    if (!valueDecl)
-      continue;
-
-    if (auto assocType = dyn_cast<AssociatedTypeDecl>(valueDecl)) {
+    if (auto assocType = dyn_cast<AssociatedTypeDecl>(requirement)) {
       if (assocType->getOverriddenDecls().empty()) {
         if (Type defaultType = findAssociatedTypeDefault(assocType).first)
           proto->setDefaultTypeWitness(assocType, defaultType);
@@ -7146,20 +7131,16 @@ void TypeChecker::inferDefaultWitnesses(ProtocolDecl *proto) {
       continue;
     }
 
-    if (isa<TypeDecl>(valueDecl))
-      continue;
-
-    if (!valueDecl->isProtocolRequirement())
-      continue;
+    assert(!isa<TypeDecl>(requirement));
 
-    ResolveWitnessResult result = checker.resolveWitnessViaLookup(valueDecl);
+    ResolveWitnessResult result = checker.resolveWitnessViaLookup(requirement);
 
     if (result == ResolveWitnessResult::Missing &&
         requirement->isSPI() &&
         !proto->isSPI()) {
       // SPI requirements need a default value, unless the protocol is SPI too.
-      valueDecl->diagnose(diag::spi_attribute_on_protocol_requirement,
-                          valueDecl->getName());
+      requirement->diagnose(diag::spi_attribute_on_protocol_requirement,
+                            requirement->getName());
     }
   }