[ClangImporter] Add @completionHandlerAsync for imported decls

Implicitly add the @completionHandlerAsync attribute for ObjCMethodDecl
that have a completion handler. Adds a link from the non-async to the
async function for use in diagnostics and refactorings.

Resolves rdar://74665226

Author: bnbarham

include/swift/AST/Attr.def

@@ -625,7 +625,7 @@ SIMPLE_DECL_ATTR(reasync, AtReasync,
   110)
 
 DECL_ATTR(completionHandlerAsync, CompletionHandlerAsync,
-  OnAbstractFunction | ConcurrencyOnly |
+  OnAbstractFunction | ConcurrencyOnly | LongAttribute |
   ABIStableToAdd | ABIStableToRemove |
   APIStableToAdd | APIStableToRemove,
   111)

include/swift/AST/Attr.h

@@ -2043,36 +2043,32 @@ class TransposeAttr final
 /// alternative, optionally providing a name (for cases when the alternative
 /// has a different name).
 class CompletionHandlerAsyncAttr final : public DeclAttribute {
-private:
-  /// DeclName of the async function in the attribute
+public:
+  /// Reference to the async alternative function. Only set for deserialized
+  /// attributes or inferred attributes from ObjectiveC code.
+  AbstractFunctionDecl *AsyncFunctionDecl;
+
+  /// DeclName of the async function in the attribute. Only set from actual
+  /// Swift code, deserialization/ObjectiveC imports will set the decl instead.
   const DeclNameRef AsyncFunctionName;
 
-public:
   /// Source location of the async function name in the attribute
   const SourceLoc AsyncFunctionNameLoc;
 
-  /// Get the name of the async function
-  ///
-  /// The name will come from the AsyncFunctionDecl if available, otherwise will
-  /// fall back on the user-provided name. If that is not defined, this function
-  /// will abort.
-  DeclNameRef getAsyncFunctionName() const;
-
   /// The index of the completion handler
   const size_t CompletionHandlerIndex;
 
   /// Source location of the completion handler index passed to the index
   const SourceLoc CompletionHandlerIndexLoc;
 
-  AbstractFunctionDecl *AsyncFunctionDecl = nullptr;
-
   CompletionHandlerAsyncAttr(DeclNameRef asyncFunctionName,
                              SourceLoc asyncFunctionNameLoc,
                              size_t completionHandlerIndex,
                              SourceLoc completionHandlerIndexLoc,
                              SourceLoc atLoc, SourceRange range)
       : DeclAttribute(DAK_CompletionHandlerAsync, atLoc, range,
                       /*implicit*/ false),
+        AsyncFunctionDecl(nullptr),
         AsyncFunctionName(asyncFunctionName),
         AsyncFunctionNameLoc(asyncFunctionNameLoc),
         CompletionHandlerIndex(completionHandlerIndex),
@@ -2084,11 +2080,9 @@ class CompletionHandlerAsyncAttr final : public DeclAttribute {
                              SourceLoc atLoc, SourceRange range)
       : DeclAttribute(DAK_CompletionHandlerAsync, atLoc, range,
                       /*implicit*/ false),
+        AsyncFunctionDecl(&asyncFunctionDecl) ,
         CompletionHandlerIndex(completionHandlerIndex),
-        CompletionHandlerIndexLoc(completionHandlerIndexLoc),
-        AsyncFunctionDecl(&asyncFunctionDecl) {}
-
-
+        CompletionHandlerIndexLoc(completionHandlerIndexLoc) {}
 
   static bool classof(const DeclAttribute *DA) {
     return DA->getKind() == DAK_CompletionHandlerAsync;

include/swift/AST/Decl.h

@@ -6097,6 +6097,8 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// constructor.
   bool hasDynamicSelfResult() const;
 
+  AbstractFunctionDecl *getAsyncAlternative() const;
+
   using DeclContext::operator new;
   using Decl::getASTContext;
 };

include/swift/AST/TypeCheckRequests.h

@@ -2954,23 +2954,20 @@ class ConditionalRequirementsRequest
   bool isCached() const { return true; }
 };
 
-class TypeCheckCompletionHandlerAsyncAttrRequest
-    : public SimpleRequest<TypeCheckCompletionHandlerAsyncAttrRequest,
-                           bool(AbstractFunctionDecl *,
-                                CompletionHandlerAsyncAttr *),
+class AsyncAlternativeRequest
+    : public SimpleRequest<AsyncAlternativeRequest,
+                           AbstractFunctionDecl *(AbstractFunctionDecl *),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
 
 private:
   friend SimpleRequest;
 
-  bool evaluate(Evaluator &evaluator,
-                AbstractFunctionDecl *attachedFucntionDecl,
-                CompletionHandlerAsyncAttr *attr) const;
+  AbstractFunctionDecl *evaluate(
+      Evaluator &evaluator, AbstractFunctionDecl *attachedFunctionDecl) const;
 
 public:
-  // Caching
   bool isCached() const { return true; }
 };
 

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -325,6 +325,6 @@ SWIFT_REQUEST(TypeChecker, SynthesizeMainFunctionRequest,
 SWIFT_REQUEST(TypeChecker, GetImplicitSendableRequest,
               NormalProtocolConformance *(NominalTypeDecl *),
               Cached, NoLocationInfo)
-SWIFT_REQUEST(TypeChecker, TypeCheckCompletionHandlerAsyncAttrRequest,
-              bool (AbstractFunctionDecl *, CompletionHandlerAsyncAttr *),
+SWIFT_REQUEST(TypeChecker, AsyncAlternativeRequest,
+              AbstractFunctionDecl *(AbstractFunctionDecl *),
               Cached, NoLocationInfo)

lib/AST/Attr.cpp

@@ -1081,6 +1081,20 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
 #include "swift/AST/Attr.def"
     llvm_unreachable("handled above");
 
+  case DAK_CompletionHandlerAsync: {
+    auto *attr = cast<CompletionHandlerAsyncAttr>(this);
+    Printer.printAttrName("@completionHandlerAsync");
+    Printer << "(\"";
+    if (attr->AsyncFunctionDecl) {
+      Printer << attr->AsyncFunctionDecl->getName();
+    } else {
+      Printer << attr->AsyncFunctionName;
+    }
+    Printer << "\", completionHandleIndex: " <<
+        attr->CompletionHandlerIndex << ')';
+    break;
+  }
+
   default:
     assert(DeclAttribute::isDeclModifier(getKind()) &&
            "handled above");
@@ -2003,11 +2017,3 @@ void swift::simple_display(llvm::raw_ostream &out, const DeclAttribute *attr) {
   if (attr)
     attr->print(out);
 }
-
-DeclNameRef CompletionHandlerAsyncAttr::getAsyncFunctionName() const {
-  if (AsyncFunctionDecl)
-    return DeclNameRef(AsyncFunctionDecl->getName());
-  if (AsyncFunctionName)
-    return AsyncFunctionName;
-  llvm_unreachable("completionHandlerAsync attr missing async function name");
-}

lib/AST/Decl.cpp

@@ -6843,6 +6843,12 @@ bool AbstractFunctionDecl::hasDynamicSelfResult() const {
   return isa<ConstructorDecl>(this);
 }
 
+AbstractFunctionDecl *AbstractFunctionDecl::getAsyncAlternative() const {
+  auto mutableFunc = const_cast<AbstractFunctionDecl *>(this);
+  return evaluateOrDefault(getASTContext().evaluator,
+                           AsyncAlternativeRequest{mutableFunc}, nullptr);
+}
+
 bool AbstractFunctionDecl::argumentNameIsAPIByDefault() const {
   // Initializers have argument labels.
   if (isa<ConstructorDecl>(this))

lib/ClangImporter/ImportDecl.cpp

@@ -7961,6 +7961,34 @@ static bool suppressOverriddenMethods(ClangImporter::Implementation &importer,
   return false;
 }
 
+void addCompletionHandlerAttribute(Decl *asyncImport,
+                                   ArrayRef<Decl *> members,
+                                   ASTContext &SwiftContext) {
+  auto *ayncFunc = dyn_cast_or_null<AbstractFunctionDecl>(asyncImport);
+  if (!ayncFunc)
+    return;
+
+  auto errorConvention = ayncFunc->getForeignErrorConvention();
+  auto asyncConvention = ayncFunc->getForeignAsyncConvention();
+  if (!asyncConvention)
+    return;
+
+  unsigned completionIndex = asyncConvention->completionHandlerParamIndex();
+  if (errorConvention &&
+      completionIndex >= errorConvention->getErrorParameterIndex()) {
+    completionIndex--;
+  }
+
+  for (auto *member : members) {
+    if (member != asyncImport) {
+      member->getAttrs().add(
+          new (SwiftContext) CompletionHandlerAsyncAttr(
+              cast<AbstractFunctionDecl>(*asyncImport), completionIndex,
+              SourceLoc(), SourceLoc(), SourceRange()));
+    }
+  }
+}
+
 /// Given a set of methods with the same selector, each taken from a
 /// different protocol in the protocol hierarchy of a class into which
 /// we want to introduce mirror imports, import only the methods which
@@ -8003,6 +8031,8 @@ void SwiftDeclConverter::importNonOverriddenMirroredMethods(DeclContext *dc,
     auto proto = entries[i].second;
     if (auto imported =
             Impl.importMirroredDecl(objcMethod, dc, getVersion(), proto)) {
+      size_t start = members.size();
+
       members.push_back(imported);
 
       for (auto alternate : Impl.getAlternateDecls(imported)) {
@@ -8014,8 +8044,13 @@ void SwiftDeclConverter::importNonOverriddenMirroredMethods(DeclContext *dc,
         auto asyncVersion = getVersion().withConcurrency(true);
         if (auto asyncImport = Impl.importMirroredDecl(
                 objcMethod, dc, asyncVersion, proto)) {
-          if (asyncImport != imported)
+          if (asyncImport != imported) {
+            addCompletionHandlerAttribute(
+                asyncImport,
+                llvm::makeArrayRef(members).drop_front(start),
+                Impl.SwiftContext);
             members.push_back(asyncImport);
+          }
         }
       }
     }
@@ -9661,7 +9696,10 @@ void ClangImporter::Implementation::loadAllMembersOfObjcContainer(
 
 void ClangImporter::Implementation::insertMembersAndAlternates(
     const clang::NamedDecl *nd, SmallVectorImpl<Decl *> &members) {
+
+  size_t start = members.size();
   llvm::SmallPtrSet<Decl *, 4> knownAlternateMembers;
+  Decl *asyncImport = nullptr;
   forEachDistinctName(
       nd, [&](ImportedName name, ImportNameVersion nameVersion) -> bool {
     auto member = importDecl(nd, nameVersion);
@@ -9682,8 +9720,18 @@ void ClangImporter::Implementation::insertMembersAndAlternates(
       return true;
 
     members.push_back(member);
+    if (nameVersion.supportsConcurrency()) {
+      assert(!asyncImport &&
+             "Should only have a single version with concurrency enabled");
+      asyncImport = member;
+    }
+
     return true;
   });
+
+  addCompletionHandlerAttribute(asyncImport,
+                                llvm::makeArrayRef(members).drop_front(start),
+                                SwiftContext);
 }
 
 void ClangImporter::Implementation::importInheritedConstructors(

lib/Sema/MiscDiagnostics.cpp

@@ -4655,28 +4655,18 @@ class CompletionHandlerUsageChecker final : public ASTWalker {
     if (expr->getType().isNull())
       return {false, expr}; // Something failed to typecheck, bail out
 
-    if (ClosureExpr *closure = dyn_cast<ClosureExpr>(expr))
+    if (auto *closure = dyn_cast<ClosureExpr>(expr))
       return {closure->isBodyAsync(), closure};
 
-    if (ApplyExpr *call = dyn_cast<ApplyExpr>(expr)) {
-      if (DeclRefExpr *fnDeclExpr = dyn_cast<DeclRefExpr>(call->getFn())) {
-        ValueDecl *fnDecl = fnDeclExpr->getDecl();
-        CompletionHandlerAsyncAttr *asyncAltAttr =
-            fnDecl->getAttrs().getAttribute<CompletionHandlerAsyncAttr>();
-        if (asyncAltAttr) {
-          // Ensure that the attribute typechecks,
-          // this also resolves the async function decl.
-          if (!evaluateOrDefault(
-                  ctx.evaluator,
-                  TypeCheckCompletionHandlerAsyncAttrRequest{
-                      cast<AbstractFunctionDecl>(fnDecl), asyncAltAttr},
-                  false)) {
+    if (auto *call = dyn_cast<ApplyExpr>(expr)) {
+      if (auto *fn = dyn_cast<DeclRefExpr>(call->getFn())) {
+        if (auto *afd = dyn_cast<AbstractFunctionDecl>(fn->getDecl())) {
+          auto *asyncFunc = afd->getAsyncAlternative();
+          if (!asyncFunc)
             return {false, call};
-          }
           ctx.Diags.diagnose(call->getLoc(), diag::warn_use_async_alternative);
-          ctx.Diags.diagnose(asyncAltAttr->AsyncFunctionDecl->getLoc(),
-                             diag::decl_declared_here,
-                             asyncAltAttr->AsyncFunctionDecl->getName());
+          ctx.Diags.diagnose(asyncFunc->getLoc(), diag::decl_declared_here,
+                             asyncFunc->getName());
         }
       }
     }

lib/Sema/TypeCheckAttr.cpp

@@ -5653,17 +5653,20 @@ void TypeChecker::checkClosureAttributes(ClosureExpr *closure) {
 
 void AttributeChecker::visitCompletionHandlerAsyncAttr(
     CompletionHandlerAsyncAttr *attr) {
-  AbstractFunctionDecl *AFD = dyn_cast<AbstractFunctionDecl>(D);
-  if (!AFD)
-    return;
-
-  evaluateOrDefault(Ctx.evaluator,
-                    TypeCheckCompletionHandlerAsyncAttrRequest{AFD, attr}, {});
+  if (AbstractFunctionDecl *AFD = dyn_cast<AbstractFunctionDecl>(D))
+    AFD->getAsyncAlternative();
 }
 
-bool TypeCheckCompletionHandlerAsyncAttrRequest::evaluate(
-    Evaluator &evaluator, AbstractFunctionDecl *attachedFunctionDecl,
-    CompletionHandlerAsyncAttr *attr) const {
+AbstractFunctionDecl *AsyncAlternativeRequest::evaluate(
+    Evaluator &evaluator, AbstractFunctionDecl *attachedFunctionDecl) const {
+  auto attr = attachedFunctionDecl->getAttrs()
+    .getAttribute<CompletionHandlerAsyncAttr>();
+  if (!attr)
+    return nullptr;
+
+  if (attr->AsyncFunctionDecl)
+    return attr->AsyncFunctionDecl;
+
   auto &Diags = attachedFunctionDecl->getASTContext().Diags;
   // Check phases:
   //  1. Attached function shouldn't be async and should have enough args
@@ -5681,7 +5684,7 @@ bool TypeCheckCompletionHandlerAsyncAttrRequest::evaluate(
                    diag::attr_completion_handler_async_handler_not_func, attr);
     Diags.diagnose(attachedFunctionDecl->getAsyncLoc(),
                    diag::note_attr_function_declared_async);
-    return false;
+    return nullptr;
   }
 
   const ParameterList *attachedFunctionParams =
@@ -5690,15 +5693,15 @@ bool TypeCheckCompletionHandlerAsyncAttrRequest::evaluate(
   if (attachedFunctionParams->size() == 0) {
     Diags.diagnose(attr->getLocation(),
                    diag::attr_completion_handler_async_handler_not_func, attr);
-    return false;
+    return nullptr;
   }
   size_t completionHandlerIndex = attr->CompletionHandlerIndexLoc.isValid()
                                       ? attr->CompletionHandlerIndex
                                       : attachedFunctionParams->size() - 1;
   if (attachedFunctionParams->size() < completionHandlerIndex) {
     Diags.diagnose(attr->CompletionHandlerIndexLoc,
                    diag::attr_completion_handler_async_handler_out_of_range);
-    return false;
+    return nullptr;
   }
 
   // Phase 2: Typecheck the completion handler
@@ -5718,7 +5721,7 @@ bool TypeCheckCompletionHandlerAsyncAttrRequest::evaluate(
               completionHandlerParamDecl->getType())
           .highlight(
               completionHandlerParamDecl->getTypeRepr()->getSourceRange());
-      return false;
+      return nullptr;
     }
 
     auto handlerTypeRepr =
@@ -5759,7 +5762,7 @@ bool TypeCheckCompletionHandlerAsyncAttrRequest::evaluate(
             handlerTypeAttrs->getLoc(TAK_autoclosure),
             diag::note_attr_completion_handler_async_handler_attr_req, false,
             "autoclosure");
-      return false;
+      return nullptr;
     }
   }
 
@@ -5793,7 +5796,7 @@ bool TypeCheckCompletionHandlerAsyncAttrRequest::evaluate(
     //  should return all of those types in a tuple
 
     SmallVector<ValueDecl *, 2> allCandidates;
-    lookupReplacedDecl(attr->getAsyncFunctionName(), attr, attachedFunctionDecl,
+    lookupReplacedDecl(attr->AsyncFunctionName, attr, attachedFunctionDecl,
                        allCandidates);
     SmallVector<AbstractFunctionDecl *, 2> candidates;
     candidates.reserve(allCandidates.size());
@@ -5810,21 +5813,20 @@ bool TypeCheckCompletionHandlerAsyncAttrRequest::evaluate(
     if (candidates.empty()) {
       Diags.diagnose(attr->AsyncFunctionNameLoc,
                      diag::attr_completion_handler_async_no_suitable_function,
-                     attr->getAsyncFunctionName());
-      return false;
+                     attr->AsyncFunctionName);
+      return nullptr;
     } else if (candidates.size() > 1) {
       Diags.diagnose(attr->AsyncFunctionNameLoc,
                      diag::attr_completion_handler_async_ambiguous_function,
-                     attr, attr->getAsyncFunctionName());
+                     attr, attr->AsyncFunctionName);
 
       for (AbstractFunctionDecl *candidate : candidates) {
         Diags.diagnose(candidate->getLoc(), diag::decl_declared_here,
                        candidate->getName());
       }
-      return false;
+      return nullptr;
     }
-    attr->AsyncFunctionDecl = candidates.front();
-  }
 
-  return true;
+    return candidates.front();
+  }
 }