[clang] Rebase swiftasynccall's musttail support onto the [[clang::musttail]] logic (PR #86011)

[John McCall via cfe-commits]

https://github.com/rjmccall created https://github.com/llvm/llvm-project/pull/86011

The old logic expects the call to be the last thing we emitted, and since it kicks in before we emit cleanups, and since `swiftasynccall` functions always return void, that's likely to be true.  "Likely" isn't very reassuring when we're talking about slapping attributes on random calls, though.  And indeed, while I can't find any way to break the logic directly in current main, our previous (ongoing?) experiments with shortening argument temporary lifetimes definitely broke it wide open.  So while this commit is prophylactic for now, it's clearly the right thing to do, and it can cherry-picked to other branches to fix problems.

>From 22ee234b13d5adec71627e9df8912b551bbb0b8e Mon Sep 17 00:00:00 2001
From: John McCall <rjmccall at apple.com>
Date: Wed, 20 Mar 2024 17:04:28 -0400
Subject: [PATCH] Rebase swiftasynccall's musttail support onto the
 [[clang::musttail]] logic

The old logic expects the call to be the last thing we emitted, and since
it kicks in before we emit cleanups, and since swiftasynccall functions
always return void, that's likely to be true.  "Likely" isn't very
reassuring when we're talking about slapping attributes on random calls,
though.  And indeed, while I can't find any way to break the logic directly
in current main, our previous (ongoing?) experiments with shortening
argument temporary lifetimes definitely broke it wide open.  So while
this commit is prophylactic for now, it's clearly the right thing to do,
and it can cherry-picked to other branches to fix problems.
---
 clang/lib/CodeGen/CGStmt.cpp               | 33 ++++++++++++----------
 clang/test/CodeGen/swift-async-call-conv.c | 16 +++++++++++
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/clang/lib/CodeGen/CGStmt.cpp b/clang/lib/CodeGen/CGStmt.cpp
index 8898e3f22a7df6..cb5a004e4f4a6c 100644
--- a/clang/lib/CodeGen/CGStmt.cpp
+++ b/clang/lib/CodeGen/CGStmt.cpp
@@ -1341,10 +1341,8 @@ struct SaveRetExprRAII {
 };
 } // namespace
 
-/// If we have 'return f(...);', where both caller and callee are SwiftAsync,
-/// codegen it as 'tail call ...; ret void;'.
-static void makeTailCallIfSwiftAsync(const CallExpr *CE, CGBuilderTy &Builder,
-                                     const CGFunctionInfo *CurFnInfo) {
+/// Determine if the given call uses the swiftasync calling convention.
+static bool isSwiftAsyncCallee(const CallExpr *CE) {
   auto calleeQualType = CE->getCallee()->getType();
   const FunctionType *calleeType = nullptr;
   if (calleeQualType->isFunctionPointerType() ||
@@ -1359,18 +1357,12 @@ static void makeTailCallIfSwiftAsync(const CallExpr *CE, CGBuilderTy &Builder,
       // getMethodDecl() doesn't handle member pointers at the moment.
       calleeType = methodDecl->getType()->castAs<FunctionType>();
     } else {
-      return;
+      return false;
     }
   } else {
-    return;
-  }
-  if (calleeType->getCallConv() == CallingConv::CC_SwiftAsync &&
-      (CurFnInfo->getASTCallingConvention() == CallingConv::CC_SwiftAsync)) {
-    auto CI = cast<llvm::CallInst>(&Builder.GetInsertBlock()->back());
-    CI->setTailCallKind(llvm::CallInst::TCK_MustTail);
-    Builder.CreateRetVoid();
-    Builder.ClearInsertionPoint();
+    return false;
   }
+  return calleeType->getCallConv() == CallingConv::CC_SwiftAsync;
 }
 
 /// EmitReturnStmt - Note that due to GCC extensions, this can have an operand
@@ -1410,6 +1402,19 @@ void CodeGenFunction::EmitReturnStmt(const ReturnStmt &S) {
   RunCleanupsScope cleanupScope(*this);
   if (const auto *EWC = dyn_cast_or_null<ExprWithCleanups>(RV))
     RV = EWC->getSubExpr();
+
+  // If we're in a swiftasynccall function, and the return expression is a
+  // call to a swiftasynccall function, mark the call as the musttail call.
+  std::optional<llvm::SaveAndRestore<const CallExpr *>> SaveMustTail;
+  if (RV && CurFnInfo &&
+      CurFnInfo->getASTCallingConvention() == CallingConv::CC_SwiftAsync) {
+    if (auto CE = dyn_cast<CallExpr>(RV)) {
+      if (isSwiftAsyncCallee(CE)) {
+        SaveMustTail.emplace(MustTailCall, CE);
+      }
+    }
+  }
+
   // FIXME: Clean this up by using an LValue for ReturnTemp,
   // EmitStoreThroughLValue, and EmitAnyExpr.
   // Check if the NRVO candidate was not globalized in OpenMP mode.
@@ -1432,8 +1437,6 @@ void CodeGenFunction::EmitReturnStmt(const ReturnStmt &S) {
     // for side effects.
     if (RV) {
       EmitAnyExpr(RV);
-      if (auto *CE = dyn_cast<CallExpr>(RV))
-        makeTailCallIfSwiftAsync(CE, Builder, CurFnInfo);
     }
   } else if (!RV) {
     // Do nothing (return value is left uninitialized)
diff --git a/clang/test/CodeGen/swift-async-call-conv.c b/clang/test/CodeGen/swift-async-call-conv.c
index ce32c22fe80984..39511698bbae9d 100644
--- a/clang/test/CodeGen/swift-async-call-conv.c
+++ b/clang/test/CodeGen/swift-async-call-conv.c
@@ -182,3 +182,19 @@ SWIFTASYNCCALL void async_struct_field_and_methods(int i, S &sref, S *sptr) {
 // CPPONLY-LABEL: define{{.*}} swifttailcc void @{{.*}}async_nonleaf_method2
 // CPPONLY: musttail call swifttailcc void @{{.*}}async_leaf_method
 #endif
+
+// Passing this as an argument requires a coerce-and-expand operation,
+// which requires a temporary.  Make sure that cleaning up that temporary
+// doesn't mess around with the musttail handling.
+struct coerce_and_expand {
+  char a,b,c,d;
+};
+struct coerce_and_expand return_coerced(void);
+SWIFTASYNCCALL void take_coerced_async(struct coerce_and_expand);
+
+// CHECK-LABEL: swifttailcc void @{{.*}}test_coerced
+SWIFTASYNCCALL void test_coerced() {
+  // CHECK:      musttail call swifttailcc void @{{.*}}take_coerced_async
+  // CHECK-NEXT: ret void
+  return take_coerced_async(return_coerced());
+}