[Coroutines] Mark 'coroutine_handle<>::address' as always-inline

Close #65054

The direct issue is still the call to coroutine_handle<>::address()
after await_suspend(). Without optimizations, the current logic will put
the temporary result of await_suspend() to the coroutine frame since the
middle end feel the temporary is escaped from
coroutine_handle<>::address. To fix this fundamentally, we should wrap
the whole logic about await-suspend into a standalone function. See
#64945

And as a short-term workaround, we probably can mark
coroutine_handle<>::address() as always-inline so that the temporary
result may not be thought to be escaped then it won't be put on the
coroutine frame. Although it looks dirty, it is probably do-able since
the compiler are allowed to do special tricks to standard library
components.

Author: ChuanqiXu9

clang/lib/Sema/SemaCoroutine.cpp

@@ -344,6 +344,28 @@ static Expr *maybeTailCall(Sema &S, QualType RetType, Expr *E,
 
   Expr *JustAddress = AddressExpr.get();
 
+  // FIXME: Without optimizations, the temporary result from `await_suspend()`
+  // may be put on the coroutine frame since the coroutine frame constructor
+  // will think the temporary variable will escape from the
+  // `coroutine_handle<>::address()` call. This is problematic since the
+  // coroutine should be considered to be suspended after it enters
+  // `await_suspend` so it shouldn't access/update the coroutine frame after
+  // that.
+  //
+  // See https://github.com/llvm/llvm-project/issues/65054 for the report.
+  //
+  // The long term solution may wrap the whole logic about `await-suspend`
+  // into a standalone function. This is similar to the proposed solution
+  // in tryMarkAwaitSuspendNoInline. See the comments there for details.
+  //
+  // The short term solution here is to mark `coroutine_handle<>::address()`
+  // function as always-inline so that the coroutine frame constructor won't
+  // think the temporary result is escaped incorrectly.
+  if (auto *FD = cast<CallExpr>(JustAddress)->getDirectCallee())
+    if (!FD->hasAttr<AlwaysInlineAttr>() && !FD->hasAttr<NoInlineAttr>())
+      FD->addAttr(AlwaysInlineAttr::CreateImplicit(S.getASTContext(),
+                                                   FD->getLocation()));
+
   // Check that the type of AddressExpr is void*
   if (!JustAddress->getType().getTypePtr()->isVoidPointerType())
     S.Diag(cast<CallExpr>(JustAddress)->getCalleeDecl()->getLocation(),

clang/test/CodeGenCoroutines/pr65054.cpp

@@ -0,0 +1,60 @@
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -std=c++20 \
+// RUN:      -O0 -disable-llvm-passes -emit-llvm %s -o - \
+// RUN:      | FileCheck %s --check-prefix=FRONTEND
+
+// The output of O0 is highly redundant and hard to test. Also it is not good
+// limit the output of O0. So we test the optimized output from O0. The idea
+// is the optimizations shouldn't change the semantics of the program.
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -std=c++20 \
+// RUN:      -O0 -emit-llvm %s -o - -disable-O0-optnone \
+// RUN:      | opt -passes='sroa,mem2reg,simplifycfg' -S | FileCheck %s --check-prefix=CHECK-O0
+
+#include "Inputs/coroutine.h"
+
+// A simple awaiter type with an await_suspend method that can't be
+// inlined.
+struct Awaiter {
+  const int& x;
+
+  bool await_ready() { return false; }
+  std::coroutine_handle<> await_suspend(const std::coroutine_handle<> h);
+  void await_resume() {}
+};
+
+struct MyTask {
+  // A lazy promise with an await_transform method that supports awaiting
+  // integer references using the Awaiter struct above.
+  struct promise_type {
+    MyTask get_return_object() {
+      return MyTask{
+          std::coroutine_handle<promise_type>::from_promise(*this),
+      };
+    }
+
+    std::suspend_always initial_suspend() { return {}; }
+    std::suspend_always final_suspend() noexcept { return {}; }
+    void unhandled_exception();
+
+    auto await_transform(const int& x) { return Awaiter{x}; }
+  };
+
+  std::coroutine_handle<> h;
+};
+
+// A global array of integers.
+int g_array[32];
+
+// A coroutine that awaits each integer in the global array.
+MyTask FooBar() {
+  for (const int& x : g_array) {
+    co_await x;
+  }
+}
+
+// FRONTEND: define{{.*}}@_ZNKSt16coroutine_handleIvE7addressEv{{.*}}#[[address_attr:[0-9]+]]
+// FRONTEND: attributes #[[address_attr]] = {{.*}}alwaysinline
+
+// CHECK-O0: define{{.*}}@_Z6FooBarv.resume
+// CHECK-O0: call{{.*}}@_ZN7Awaiter13await_suspendESt16coroutine_handleIvE
+// CHECK-O0-NOT: store
+// CHECK-O0: ret void