AST: Requestify lookup of protocol referenced by ImplementsAttr

Direct lookup relied in primary file checking to have filled in the
protocol type stored in the ImplementsAttr. This was already wrong
with multi-file test cases in non-WMO mode, and crashed in the
ASTPrinter if printing a declaration in a non-primary file.

I don't have a standalone test case that is independent of my
upcoming ASTPrinter changes, but this is a nice cleanup regardless.

Author: slavapestov

Diff for: include/swift/AST/Attr.h

@@ -1551,25 +1551,28 @@ class SpecializeAttr final
 /// The @_implements attribute, which treats a decl as the implementation for
 /// some named protocol requirement (but otherwise not-visible by that name).
 class ImplementsAttr : public DeclAttribute {
-  TypeExpr *ProtocolType;
+  TypeRepr *TyR;
   DeclName MemberName;
   DeclNameLoc MemberNameLoc;
 
-public:
   ImplementsAttr(SourceLoc atLoc, SourceRange Range,
-                 TypeExpr *ProtocolType,
+                 TypeRepr *TyR,
                  DeclName MemberName,
                  DeclNameLoc MemberNameLoc);
 
+public:
   static ImplementsAttr *create(ASTContext &Ctx, SourceLoc atLoc,
                                 SourceRange Range,
-                                TypeExpr *ProtocolType,
+                                TypeRepr *TyR,
                                 DeclName MemberName,
                                 DeclNameLoc MemberNameLoc);
 
-  void setProtocolType(Type ty);
-  Type getProtocolType() const;
-  TypeRepr *getProtocolTypeRepr() const;
+  static ImplementsAttr *create(DeclContext *DC,
+                                ProtocolDecl *Proto,
+                                DeclName MemberName);
+
+  ProtocolDecl *getProtocol(DeclContext *dc) const;
+  TypeRepr *getProtocolTypeRepr() const { return TyR; }
 
   DeclName getMemberName() const { return MemberName; }
   DeclNameLoc getMemberNameLoc() const { return MemberNameLoc; }

Diff for: include/swift/AST/NameLookupRequests.h

@@ -912,6 +912,25 @@ class PotentialMacroExpansionsInContextRequest
   bool isCached() const { return true; }
 };
 
+/// Resolves the protocol referenced by an @_implements attribute.
+class ImplementsAttrProtocolRequest
+    : public SimpleRequest<ImplementsAttrProtocolRequest,
+                           ProtocolDecl *(const ImplementsAttr *, DeclContext *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  ProtocolDecl *evaluate(Evaluator &evaluator, const ImplementsAttr *attr,
+                         DeclContext *dc) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 #define SWIFT_TYPEID_ZONE NameLookup
 #define SWIFT_TYPEID_HEADER "swift/AST/NameLookupTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

Diff for: include/swift/AST/NameLookupTypeIDZone.def

@@ -109,3 +109,5 @@ SWIFT_REQUEST(NameLookup, HasDynamicCallableAttributeRequest,
               bool(NominalTypeDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(NameLookup, PotentialMacroExpansionsInContextRequest,
               PotentialMacroExpansions(TypeOrExtension), Cached, NoLocationInfo)
+SWIFT_REQUEST(NameLookup, ImplementsAttrProtocolRequest,
+              ProtocolDecl *(const ImplementsAttr *, DeclContext *), Cached, NoLocationInfo)

Diff for: lib/AST/Attr.cpp

@@ -1239,7 +1239,10 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     Printer.printAttrName("@_implements");
     Printer << "(";
     auto *attr = cast<ImplementsAttr>(this);
-    attr->getProtocolType().print(Printer, Options);
+    if (auto *proto = attr->getProtocol(D->getDeclContext()))
+      proto->getDeclaredInterfaceType()->print(Printer, Options);
+    else
+      attr->getProtocolTypeRepr()->print(Printer, Options);
     Printer << ", " << attr->getMemberName() << ")";
     break;
   }
@@ -2360,37 +2363,41 @@ TransposeAttr *TransposeAttr::create(ASTContext &context, bool implicit,
 }
 
 ImplementsAttr::ImplementsAttr(SourceLoc atLoc, SourceRange range,
-                               TypeExpr *ProtocolType,
+                               TypeRepr *TyR,
                                DeclName MemberName,
                                DeclNameLoc MemberNameLoc)
     : DeclAttribute(DAK_Implements, atLoc, range, /*Implicit=*/false),
-      ProtocolType(ProtocolType),
+      TyR(TyR),
       MemberName(MemberName),
       MemberNameLoc(MemberNameLoc) {
 }
 
-
 ImplementsAttr *ImplementsAttr::create(ASTContext &Ctx, SourceLoc atLoc,
                                        SourceRange range,
-                                       TypeExpr *ProtocolType,
+                                       TypeRepr *TyR,
                                        DeclName MemberName,
                                        DeclNameLoc MemberNameLoc) {
   void *mem = Ctx.Allocate(sizeof(ImplementsAttr), alignof(ImplementsAttr));
-  return new (mem) ImplementsAttr(atLoc, range, ProtocolType,
+  return new (mem) ImplementsAttr(atLoc, range, TyR,
                                   MemberName, MemberNameLoc);
 }
 
-void ImplementsAttr::setProtocolType(Type ty) {
-  assert(ty);
-  ProtocolType->setType(MetatypeType::get(ty));
-}
-
-Type ImplementsAttr::getProtocolType() const {
-  return ProtocolType->getInstanceType();
+ImplementsAttr *ImplementsAttr::create(DeclContext *DC,
+                                       ProtocolDecl *Proto,
+                                       DeclName MemberName) {
+  auto &ctx = DC->getASTContext();
+  void *mem = ctx.Allocate(sizeof(ImplementsAttr), alignof(ImplementsAttr));
+  auto *attr = new (mem) ImplementsAttr(
+      SourceLoc(), SourceRange(), nullptr,
+      MemberName, DeclNameLoc());
+  ctx.evaluator.cacheOutput(ImplementsAttrProtocolRequest{attr, DC},
+                            std::move(Proto));
+  return attr;
 }
 
-TypeRepr *ImplementsAttr::getProtocolTypeRepr() const {
-  return ProtocolType->getTypeRepr();
+ProtocolDecl *ImplementsAttr::getProtocol(DeclContext *dc) const {
+  return evaluateOrDefault(dc->getASTContext().evaluator,
+        ImplementsAttrProtocolRequest{this, dc}, nullptr);
 }
 
 CustomAttr::CustomAttr(SourceLoc atLoc, SourceRange range, TypeExpr *type,

Diff for: lib/AST/NameLookup.cpp

@@ -3644,6 +3644,28 @@ bool TypeBase::hasDynamicCallableAttribute() {
   });
 }
 
+ProtocolDecl *ImplementsAttrProtocolRequest::evaluate(
+    Evaluator &evaluator, const ImplementsAttr *attr, DeclContext *dc) const {
+
+  auto typeRepr = attr->getProtocolTypeRepr();
+
+  ASTContext &ctx = dc->getASTContext();
+  DirectlyReferencedTypeDecls referenced =
+    directReferencesForTypeRepr(evaluator, ctx, typeRepr, dc);
+
+  // Resolve those type declarations to nominal type declarations.
+  SmallVector<ModuleDecl *, 2> modulesFound;
+  bool anyObject = false;
+  auto nominalTypes
+    = resolveTypeDeclsToNominal(evaluator, ctx, referenced, modulesFound,
+                                anyObject);
+
+  if (nominalTypes.empty())
+    return nullptr;
+
+  return dyn_cast<ProtocolDecl>(nominalTypes.front());
+}
+
 void FindLocalVal::checkPattern(const Pattern *Pat, DeclVisibilityKind Reason) {
   Pat->forEachVariable([&](VarDecl *VD) { checkValueDecl(VD, Reason); });
 }

Diff for: lib/Parse/ParseDecl.cpp

@@ -1218,10 +1218,9 @@ Parser::parseImplementsAttribute(SourceLoc AtLoc, SourceLoc Loc) {
   }
 
   // FIXME(ModQual): Reject module qualification on MemberName.
-  auto *TE = new (Context) TypeExpr(ProtocolType.get());
   return ParserResult<ImplementsAttr>(
     ImplementsAttr::create(Context, AtLoc, SourceRange(Loc, rParenLoc),
-                           TE, MemberName.getFullName(),
+                           ProtocolType.get(), MemberName.getFullName(),
                            MemberNameLoc));
 }
 

Diff for: lib/Sema/DerivedConformanceComparable.cpp

@@ -259,16 +259,12 @@ deriveComparable_lt(
   // Add the @_implements(Comparable, < (_:_:)) attribute
   if (generatedIdentifier != C.Id_LessThanOperator) {
     auto comparable = C.getProtocol(KnownProtocolKind::Comparable);
-    auto comparableType = comparable->getDeclaredInterfaceType();
-    auto comparableTypeExpr = TypeExpr::createImplicit(comparableType, C);
     SmallVector<Identifier, 2> argumentLabels = { Identifier(), Identifier() };
     auto comparableDeclName = DeclName(C, DeclBaseName(C.Id_LessThanOperator),
                                    argumentLabels);
-    comparableDecl->getAttrs().add(new (C) ImplementsAttr(SourceLoc(),
-                                                          SourceRange(),
-                                                          comparableTypeExpr,
-                                                          comparableDeclName,
-                                                          DeclNameLoc()));
+    comparableDecl->getAttrs().add(ImplementsAttr::create(parentDC,
+                                                          comparable,
+                                                          comparableDeclName));
   }
 
   if (!C.getLessThanIntDecl()) {

Diff for: lib/Sema/DerivedConformanceEquatableHashable.cpp

@@ -417,16 +417,12 @@ deriveEquatable_eq(
   // Add the @_implements(Equatable, ==(_:_:)) attribute
   if (generatedIdentifier != C.Id_EqualsOperator) {
     auto equatableProto = C.getProtocol(KnownProtocolKind::Equatable);
-    auto equatableTy = equatableProto->getDeclaredInterfaceType();
-    auto equatableTyExpr = TypeExpr::createImplicit(equatableTy, C);
     SmallVector<Identifier, 2> argumentLabels = { Identifier(), Identifier() };
     auto equalsDeclName = DeclName(C, DeclBaseName(C.Id_EqualsOperator),
                                    argumentLabels);
-    eqDecl->getAttrs().add(new (C) ImplementsAttr(SourceLoc(),
-                                                  SourceRange(),
-                                                  equatableTyExpr,
-                                                  equalsDeclName,
-                                                  DeclNameLoc()));
+    eqDecl->getAttrs().add(ImplementsAttr::create(parentDC,
+                                                  equatableProto,
+                                                  equalsDeclName));
   }
 
   if (!C.getEqualIntDecl()) {

Diff for: lib/Sema/TypeCheckAttr.cpp

@@ -3572,58 +3572,45 @@ void AttributeChecker::visitTypeEraserAttr(TypeEraserAttr *attr) {
 void AttributeChecker::visitImplementsAttr(ImplementsAttr *attr) {
   DeclContext *DC = D->getDeclContext();
 
-  Type T = attr->getProtocolType();
-  if (!T && attr->getProtocolTypeRepr()) {
-    auto context = TypeResolverContext::GenericRequirement;
-    T = TypeResolution::resolveContextualType(attr->getProtocolTypeRepr(), DC,
-                                              TypeResolutionOptions(context),
-                                              /*unboundTyOpener*/ nullptr,
-                                              /*placeholderHandler*/ nullptr,
-                                              /*packElementOpener*/ nullptr);
-  }
-
-  // Definite error-types were already diagnosed in resolveType.
-  if (T->hasError())
+  ProtocolDecl *PD = attr->getProtocol(DC);
+
+  if (!PD) {
+    diagnose(attr->getLocation(), diag::implements_attr_non_protocol_type)
+      .highlight(attr->getProtocolTypeRepr()->getSourceRange());
     return;
-  attr->setProtocolType(T);
+  }
 
-  // Check that we got a ProtocolType.
-  if (auto PT = T->getAs<ProtocolType>()) {
-    ProtocolDecl *PD = PT->getDecl();
+  // Check that the ProtocolType has the specified member.
+  LookupResult R =
+      TypeChecker::lookupMember(PD->getDeclContext(),
+                                PD->getDeclaredInterfaceType(),
+                                DeclNameRef(attr->getMemberName()));
+  if (!R) {
+    diagnose(attr->getLocation(),
+             diag::implements_attr_protocol_lacks_member,
+             PD->getName(), attr->getMemberName())
+      .highlight(attr->getMemberNameLoc().getSourceRange());
+    return;
+  }
 
-    // Check that the ProtocolType has the specified member.
-    LookupResult R =
-        TypeChecker::lookupMember(PD->getDeclContext(), PT,
-                                  DeclNameRef(attr->getMemberName()));
-    if (!R) {
+  // Check that the decl we're decorating is a member of a type that actually
+  // conforms to the specified protocol.
+  NominalTypeDecl *NTD = DC->getSelfNominalTypeDecl();
+  if (auto *OtherPD = dyn_cast<ProtocolDecl>(NTD)) {
+    if (!OtherPD->inheritsFrom(PD)) {
       diagnose(attr->getLocation(),
-               diag::implements_attr_protocol_lacks_member,
-               PD->getName(), attr->getMemberName())
-        .highlight(attr->getMemberNameLoc().getSourceRange());
-    }
-
-    // Check that the decl we're decorating is a member of a type that actually
-    // conforms to the specified protocol.
-    NominalTypeDecl *NTD = DC->getSelfNominalTypeDecl();
-    if (auto *OtherPD = dyn_cast<ProtocolDecl>(NTD)) {
-      if (!OtherPD->inheritsFrom(PD)) {
-        diagnose(attr->getLocation(),
-                 diag::implements_attr_protocol_not_conformed_to,
-                 NTD->getName(), PD->getName())
-          .highlight(attr->getProtocolTypeRepr()->getSourceRange());
-      }
-    } else {
-      SmallVector<ProtocolConformance *, 2> conformances;
-      if (!NTD->lookupConformance(PD, conformances)) {
-        diagnose(attr->getLocation(),
-                 diag::implements_attr_protocol_not_conformed_to,
-                 NTD->getName(), PD->getName())
-          .highlight(attr->getProtocolTypeRepr()->getSourceRange());
-      }
+               diag::implements_attr_protocol_not_conformed_to,
+               NTD->getName(), PD->getName())
+        .highlight(attr->getProtocolTypeRepr()->getSourceRange());
     }
   } else {
-    diagnose(attr->getLocation(), diag::implements_attr_non_protocol_type)
-      .highlight(attr->getProtocolTypeRepr()->getSourceRange());
+    SmallVector<ProtocolConformance *, 2> conformances;
+    if (!NTD->lookupConformance(PD, conformances)) {
+      diagnose(attr->getLocation(),
+               diag::implements_attr_protocol_not_conformed_to,
+               NTD->getName(), PD->getName())
+        .highlight(attr->getProtocolTypeRepr()->getSourceRange());
+    }
   }
 }
 

Diff for: lib/Sema/TypeCheckProtocol.cpp

@@ -1237,11 +1237,9 @@ witnessHasImplementsAttrForExactRequirement(ValueDecl *witness,
   assert(requirement->isProtocolRequirement());
   auto *PD = cast<ProtocolDecl>(requirement->getDeclContext());
   if (auto A = witness->getAttrs().getAttribute<ImplementsAttr>()) {
-    if (Type T = A->getProtocolType()) {
-      if (auto ProtoTy = T->getAs<ProtocolType>()) {
-        if (ProtoTy->getDecl() == PD) {
-          return A->getMemberName() == requirement->getName();
-        }
+    if (auto *OtherPD = A->getProtocol(witness->getDeclContext())) {
+      if (OtherPD == PD) {
+        return A->getMemberName() == requirement->getName();
       }
     }
   }