AST: Refactor SelfBoundsFromWhereClause to support ProtocolDecls

Author: slavapestov

include/swift/AST/NameLookup.h

@@ -34,6 +34,7 @@ namespace swift {
   class Type;
   class TypeDecl;
   class ValueDecl;
+  struct SelfBounds;
 
 /// LookupResultEntry - One result of unqualified lookup.
 struct LookupResultEntry {
@@ -374,6 +375,12 @@ getDirectlyInheritedNominalTypeDecls(
                       llvm::PointerUnion<TypeDecl *, ExtensionDecl *> decl,
                       bool &anyObject);
 
+/// Retrieve the set of nominal type declarations that appear as the
+/// constraint type of any "Self" constraints in the where clause of the
+/// given protocol or protocol extension.
+SelfBounds getSelfBoundsFromWhereClause(
+    llvm::PointerUnion<TypeDecl *, ExtensionDecl *> decl);
+
 } // end namespace swift
 
 #endif

include/swift/AST/NameLookupRequests.h

@@ -182,26 +182,31 @@ class ExtendedNominalRequest :
   void noteCycleStep(DiagnosticEngine &diags) const;
 };
 
+struct SelfBounds {
+  llvm::TinyPtrVector<NominalTypeDecl *> decls;
+  bool anyObject = false;
+};
+
 /// Request the nominal types that occur as the right-hand side of "Self: Foo"
 /// constraints in the "where" clause of a protocol extension.
 class SelfBoundsFromWhereClauseRequest :
     public SimpleRequest<SelfBoundsFromWhereClauseRequest,
                          CacheKind::Uncached,
-                         llvm::TinyPtrVector<NominalTypeDecl *>,
-                         ExtensionDecl *> {
+                         SelfBounds,
+                         llvm::PointerUnion<TypeDecl *, ExtensionDecl *>> {
 public:
   using SimpleRequest::SimpleRequest;
 
 private:
   friend SimpleRequest;
 
   // Evaluation.
-  llvm::TinyPtrVector<NominalTypeDecl *> evaluate(Evaluator &evaluator,
-                                                  ExtensionDecl *ext) const;
+  SelfBounds evaluate(Evaluator &evaluator,
+                      llvm::PointerUnion<TypeDecl *, ExtensionDecl *>) const;
 
 public:
   // Cycle handling
-  llvm::TinyPtrVector<NominalTypeDecl *> breakCycle() const { return { }; }
+  SelfBounds breakCycle() const { return { }; }
   void diagnoseCycle(DiagnosticEngine &diags) const;
   void noteCycleStep(DiagnosticEngine &diags) const;
 };

lib/AST/Decl.cpp

@@ -3642,7 +3642,7 @@ ProtocolDecl::getInheritedProtocols() const {
   SmallPtrSet<const ProtocolDecl *, 4> known;
   known.insert(this);
   bool anyObject = false;
-  for (const auto &found :
+  for (const auto found :
            getDirectlyInheritedNominalTypeDecls(
              const_cast<ProtocolDecl *>(this), anyObject)) {
     if (auto proto = dyn_cast<ProtocolDecl>(found.second)) {
@@ -3758,12 +3758,14 @@ bool ProtocolDecl::requiresClassSlow() {
     getDirectlyInheritedNominalTypeDecls(this, anyObject);
 
   // Quick check: do we inherit AnyObject?
-  if (anyObject)
-    return Bits.ProtocolDecl.RequiresClass = true;
+  if (anyObject) {
+    Bits.ProtocolDecl.RequiresClass = true;
+    return true;
+  }
 
   // Look through all of the inherited nominals for a superclass or a
   // class-bound protocol.
-  for (const auto &found : allInheritedNominals) {
+  for (const auto found : allInheritedNominals) {
     // Superclass bound.
     if (isa<ClassDecl>(found.second))
       return Bits.ProtocolDecl.RequiresClass = true;

lib/AST/NameLookup.cpp

@@ -556,18 +556,26 @@ resolveTypeDeclsToNominal(Evaluator &evaluator,
                           SmallVectorImpl<ModuleDecl *> &modulesFound,
                           bool &anyObject);
 
-TinyPtrVector<NominalTypeDecl *>
-SelfBoundsFromWhereClauseRequest::evaluate(Evaluator &evaluator,
-                                           ExtensionDecl *ext) const {
-  auto proto = ext->getExtendedProtocolDecl();
-  assert(proto && "Not a protocol extension?");
-
-  ASTContext &ctx = proto->getASTContext();
-  TinyPtrVector<NominalTypeDecl *> result;
-  if (!ext->getGenericParams())
+SelfBounds
+SelfBoundsFromWhereClauseRequest::evaluate(
+    Evaluator &evaluator,
+    llvm::PointerUnion<TypeDecl *, ExtensionDecl *> decl) const {
+  auto *typeDecl = decl.dyn_cast<TypeDecl *>();
+  auto *protoDecl = dyn_cast_or_null<ProtocolDecl>(typeDecl);
+  auto *extDecl = decl.dyn_cast<ExtensionDecl *>();
+
+  DeclContext *dc = protoDecl ? (DeclContext *)protoDecl : (DeclContext *)extDecl;
+  auto requirements = protoDecl ? protoDecl->getTrailingWhereClause()
+                                : extDecl->getTrailingWhereClause();
+
+  ASTContext &ctx = dc->getASTContext();
+
+  SelfBounds result;
+
+  if (requirements == nullptr)
     return result;
 
-  for (const auto &req : ext->getGenericParams()->getTrailingRequirements()) {
+  for (const auto &req : requirements->getRequirements()) {
     // We only care about type constraints.
     if (req.getKind() != RequirementReprKind::TypeConstraint)
       continue;
@@ -578,29 +586,41 @@ SelfBoundsFromWhereClauseRequest::evaluate(Evaluator &evaluator,
       if (auto identTypeRepr = dyn_cast<SimpleIdentTypeRepr>(typeRepr))
         isSelfLHS = (identTypeRepr->getIdentifier() == ctx.Id_Self);
     } else if (Type type = req.getSubject()) {
-      isSelfLHS = type->isEqual(proto->getSelfInterfaceType());
+      isSelfLHS = type->isEqual(dc->getSelfInterfaceType());
     }
     if (!isSelfLHS)
       continue;
 
     // Resolve the right-hand side.
     DirectlyReferencedTypeDecls rhsDecls;
     if (auto typeRepr = req.getConstraintRepr()) {
-      rhsDecls = directReferencesForTypeRepr(evaluator, ctx, typeRepr, ext);
+      rhsDecls = directReferencesForTypeRepr(evaluator, ctx, typeRepr, dc);
     } else if (Type type = req.getConstraint()) {
       rhsDecls = directReferencesForType(type);
     }
 
     SmallVector<ModuleDecl *, 2> modulesFound;
-    bool anyObject = false;
     auto rhsNominals = resolveTypeDeclsToNominal(evaluator, ctx, rhsDecls,
-                                                 modulesFound, anyObject);
-    result.insert(result.end(), rhsNominals.begin(), rhsNominals.end());
+                                                 modulesFound,
+                                                 result.anyObject);
+    result.decls.insert(result.decls.end(),
+                        rhsNominals.begin(),
+                        rhsNominals.end());
   }
 
   return result;
 }
 
+SelfBounds swift::getSelfBoundsFromWhereClause(
+    llvm::PointerUnion<TypeDecl *, ExtensionDecl *> decl) {
+  auto *typeDecl = decl.dyn_cast<TypeDecl *>();
+  auto *extDecl = decl.dyn_cast<ExtensionDecl *>();
+  auto &ctx = typeDecl ? typeDecl->getASTContext()
+                       : extDecl->getASTContext();
+  return evaluateOrDefault(ctx.evaluator,
+                           SelfBoundsFromWhereClauseRequest{decl}, {});
+}
+
 static void
 populateLookupDeclsFromContext(DeclContext *dc,
                                SmallVectorImpl<NominalTypeDecl *> &lookupDecls) {
@@ -614,9 +634,8 @@ populateLookupDeclsFromContext(DeclContext *dc,
   // constraints that can affect name lookup.
   if (dc->getExtendedProtocolDecl()) {
     auto ext = cast<ExtensionDecl>(dc);
-    auto bounds = evaluateOrDefault(dc->getASTContext().evaluator,
-                                    SelfBoundsFromWhereClauseRequest{ext}, {});
-    for (auto bound : bounds)
+    auto bounds = getSelfBoundsFromWhereClause(ext);
+    for (auto bound : bounds.decls)
       lookupDecls.push_back(bound);
   }
 }

lib/AST/NameLookupRequests.cpp

@@ -121,15 +121,21 @@ void ExtendedNominalRequest::noteCycleStep(DiagnosticEngine &diags) const {
 void SelfBoundsFromWhereClauseRequest::diagnoseCycle(
                                               DiagnosticEngine &diags) const {
   // FIXME: Improve this diagnostic.
-  auto ext = std::get<0>(getStorage());
-  diags.diagnose(ext, diag::circular_reference);
+  auto subject = std::get<0>(getStorage());
+  Decl *decl = subject.dyn_cast<TypeDecl *>();
+  if (decl == nullptr)
+    decl = subject.get<ExtensionDecl *>();
+  diags.diagnose(decl, diag::circular_reference);
 }
 
 void SelfBoundsFromWhereClauseRequest::noteCycleStep(
                                               DiagnosticEngine &diags) const {
-  auto ext = std::get<0>(getStorage());
   // FIXME: Customize this further.
-  diags.diagnose(ext, diag::circular_reference_through);
+  auto subject = std::get<0>(getStorage());
+  Decl *decl = subject.dyn_cast<TypeDecl *>();
+  if (decl == nullptr)
+    decl = subject.get<ExtensionDecl *>();
+  diags.diagnose(decl, diag::circular_reference_through);
 }
 
 void TypeDeclsFromWhereClauseRequest::diagnoseCycle(