AST: Introduce ProtocolDecl::getAllInheritedProtocols()

Author: slavapestov

include/swift/AST/Decl.h

@@ -5172,6 +5172,10 @@ class ProtocolDecl final : public NominalTypeDecl {
   /// Retrieve the set of protocols inherited from this protocol.
   ArrayRef<ProtocolDecl *> getInheritedProtocols() const;
 
+  /// Retrieve the transitive closure of the inherited protocols, not including
+  /// this protocol itself.
+  ArrayRef<ProtocolDecl *> getAllInheritedProtocols() const;
+
   /// Determine whether this protocol has a superclass.
   bool hasSuperclass() const { return (bool)getSuperclassDecl(); }
 

include/swift/AST/NameLookupRequests.h

@@ -199,6 +199,25 @@ class InheritedProtocolsRequest
                            ArrayRef<ProtocolDecl *> result) const;
 };
 
+class AllInheritedProtocolsRequest
+    : public SimpleRequest<
+          AllInheritedProtocolsRequest, ArrayRef<ProtocolDecl *>(ProtocolDecl *),
+          RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  ArrayRef<ProtocolDecl *>
+  evaluate(Evaluator &evaluator, ProtocolDecl *PD) const;
+
+public:
+  // Caching
+  bool isCached() const { return true; }
+};
+
 class ProtocolRequirementsRequest
     : public SimpleRequest<ProtocolRequirementsRequest,
                            ArrayRef<ValueDecl *>(ProtocolDecl *),

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, AllInheritedProtocolsRequest,
+              ArrayRef<ProtocolDecl *>(ProtocolDecl *), Cached,
+              NoLocationInfo)
 SWIFT_REQUEST(NameLookup, ProtocolRequirementsRequest,
               ArrayRef<ValueDecl *>(ProtocolDecl *), SeparatelyCached,
               NoLocationInfo)

lib/AST/Decl.cpp

@@ -6470,6 +6470,13 @@ ArrayRef<ProtocolDecl *> ProtocolDecl::getInheritedProtocols() const {
                            {});
 }
 
+ArrayRef<ProtocolDecl *> ProtocolDecl::getAllInheritedProtocols() const {
+  auto *mutThis = const_cast<ProtocolDecl *>(this);
+  return evaluateOrDefault(getASTContext().evaluator,
+                           AllInheritedProtocolsRequest{mutThis},
+                           {});
+}
+
 ArrayRef<AssociatedTypeDecl *>
 ProtocolDecl::getAssociatedTypeMembers() const {
   if (Bits.ProtocolDecl.HasAssociatedTypes)

lib/AST/NameLookup.cpp

@@ -3323,6 +3323,20 @@ InheritedProtocolsRequest::evaluate(Evaluator &evaluator,
   return ctx.AllocateCopy(inherited.getArrayRef());
 }
 
+ArrayRef<ProtocolDecl *>
+AllInheritedProtocolsRequest::evaluate(Evaluator &evaluator,
+                                       ProtocolDecl *PD) const {
+  llvm::SmallSetVector<ProtocolDecl *, 2> result;
+
+  PD->walkInheritedProtocols([&](ProtocolDecl *inherited) {
+    if (inherited != PD)
+      result.insert(inherited);
+    return TypeWalker::Action::Continue;
+  });
+
+  return PD->getASTContext().AllocateCopy(result.getArrayRef());
+}
+
 ArrayRef<ValueDecl *>
 ProtocolRequirementsRequest::evaluate(Evaluator &evaluator,
                                       ProtocolDecl *PD) const {