[Coroutines] Mark parameter allocas with coro.outside.frame metadata (#127653)

Parameters to a coroutine get copied (moved) to coroutine-local
instances which code inside the coroutine then uses.

The original parameters should not be part of the frame. Normally
CoroSplit figures that out by itself, but for [[clang::trivial_abi]]
parameters which, get destructed at the end of the ramp function, it
does not (see bug), causing use-after-free's if the frame is destroyed
before the end of the ramp (as happens if it doesn't suspend).

Since Clang knows these should never be part of the frame, use metadata
to make it so.

Fixes #127499

Author: zmodem

clang/lib/CodeGen/CGCoroutine.cpp

@@ -855,6 +855,20 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
     // Create parameter copies. We do it before creating a promise, since an
     // evolution of coroutine TS may allow promise constructor to observe
     // parameter copies.
+    for (const ParmVarDecl *Parm : FnArgs) {
+      // If the original param is in an alloca, exclude it from the coroutine
+      // frame. The parameter copy will be part of the frame, but the original
+      // parameter memory should remain on the stack. This is necessary to
+      // ensure that parameters destroyed in callees, as with `trivial_abi` or
+      // in the MSVC C++ ABI, are appropriately destroyed after setting up the
+      // coroutine.
+      Address ParmAddr = GetAddrOfLocalVar(Parm);
+      if (auto *ParmAlloca =
+              dyn_cast<llvm::AllocaInst>(ParmAddr.getBasePointer())) {
+        ParmAlloca->setMetadata(llvm::LLVMContext::MD_coro_outside_frame,
+                                llvm::MDNode::get(CGM.getLLVMContext(), {}));
+      }
+    }
     for (auto *PM : S.getParamMoves()) {
       EmitStmt(PM);
       ParamReplacer.addCopy(cast<DeclStmt>(PM));

clang/test/CodeGenCoroutines/coro-params.cpp

@@ -3,6 +3,7 @@
 // Vefifies that parameter copies are used in the body of the coroutine
 // Verifies that parameter copies are used to construct the promise type, if that type has a matching constructor
 // RUN: %clang_cc1 -std=c++20 -triple=x86_64-unknown-linux-gnu -emit-llvm -o - %s -disable-llvm-passes -fexceptions | FileCheck %s
+// RUN: %clang_cc1 -std=c++20 -triple=x86_64-pc-win32          -emit-llvm -o - %s -disable-llvm-passes -fexceptions | FileCheck %s --check-prefix=MSABI
 
 namespace std {
 template <typename... T> struct coroutine_traits;
@@ -59,46 +60,65 @@ struct MoveAndCopy {
   ~MoveAndCopy();
 };
 
-void consume(int,int,int) noexcept;
+struct [[clang::trivial_abi]] TrivialABI {
+  int val;
+  TrivialABI(TrivialABI&&) noexcept;
+  ~TrivialABI();
+};
+
+void consume(int,int,int,int) noexcept;
 
 // TODO: Add support for CopyOnly params
-// CHECK: define{{.*}} void @_Z1fi8MoveOnly11MoveAndCopy(i32 noundef %val, ptr noundef %[[MoParam:.+]], ptr noundef %[[McParam:.+]]) #0 personality ptr @__gxx_personality_v0
-void f(int val, MoveOnly moParam, MoveAndCopy mcParam) {
+// CHECK: define{{.*}} void @_Z1fi8MoveOnly11MoveAndCopy10TrivialABI(i32 noundef %val, ptr noundef %[[MoParam:.+]], ptr noundef %[[McParam:.+]], i32 %[[TrivialParam:.+]]) #0 personality ptr @__gxx_personality_v0
+void f(int val, MoveOnly moParam, MoveAndCopy mcParam, TrivialABI trivialParam) {
+  // CHECK: %[[TrivialAlloca:.+]] = alloca %struct.TrivialABI,
+  // CHECK-SAME: !coro.outside.frame
   // CHECK: %[[MoCopy:.+]] = alloca %struct.MoveOnly,
   // CHECK: %[[McCopy:.+]] = alloca %struct.MoveAndCopy,
+  // CHECK: %[[TrivialCopy:.+]] = alloca %struct.TrivialABI,
   // CHECK: store i32 %val, ptr %[[ValAddr:.+]]
 
   // CHECK: call ptr @llvm.coro.begin(
   // CHECK: call void @_ZN8MoveOnlyC1EOS_(ptr {{[^,]*}} %[[MoCopy]], ptr noundef nonnull align 4 dereferenceable(4) %[[MoParam]])
   // CHECK-NEXT: call void @llvm.lifetime.start.p0(
   // CHECK-NEXT: call void @_ZN11MoveAndCopyC1EOS_(ptr {{[^,]*}} %[[McCopy]], ptr noundef nonnull align 4 dereferenceable(4) %[[McParam]]) #
   // CHECK-NEXT: call void @llvm.lifetime.start.p0(
-  // CHECK-NEXT: invoke void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopyEE12promise_typeC1Ev(
+  // CHECK-NEXT: call void @_ZN10TrivialABIC1EOS_(ptr {{[^,]*}} %[[TrivialCopy]], ptr {{[^,]*}} %[[TrivialAlloca]])
+  // CHECK-NEXT: call void @llvm.lifetime.start.p0(
+  // CHECK-NEXT: invoke void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopy10TrivialABIEE12promise_typeC1Ev(
 
   // CHECK: call void @_ZN14suspend_always12await_resumeEv(
   // CHECK: %[[IntParam:.+]] = load i32, ptr %{{.*}}
   // CHECK: %[[MoGep:.+]] = getelementptr inbounds nuw %struct.MoveOnly, ptr %[[MoCopy]], i32 0, i32 0
   // CHECK: %[[MoVal:.+]] = load i32, ptr %[[MoGep]]
-  // CHECK: %[[McGep:.+]] =  getelementptr inbounds nuw %struct.MoveAndCopy, ptr %[[McCopy]], i32 0, i32 0
+  // CHECK: %[[McGep:.+]] = getelementptr inbounds nuw %struct.MoveAndCopy, ptr %[[McCopy]], i32 0, i32 0
   // CHECK: %[[McVal:.+]] = load i32, ptr %[[McGep]]
-  // CHECK: call void @_Z7consumeiii(i32 noundef %[[IntParam]], i32 noundef %[[MoVal]], i32 noundef %[[McVal]])
+  // CHECK: %[[TrivialGep:.+]] = getelementptr inbounds nuw %struct.TrivialABI, ptr %[[TrivialCopy]], i32 0, i32 0
+  // CHECK: %[[TrivialVal:.+]] = load i32, ptr %[[TrivialGep]]
+  // CHECK: call void @_Z7consumeiiii(i32 noundef %[[IntParam]], i32 noundef %[[MoVal]], i32 noundef %[[McVal]], i32 noundef %[[TrivialVal]])
 
-  consume(val, moParam.val, mcParam.val);
+  consume(val, moParam.val, mcParam.val, trivialParam.val);
   co_return;
 
   // Skip to final suspend:
-  // CHECK: call void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopyEE12promise_type13final_suspendEv(
+  // CHECK: call void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopy10TrivialABIEE12promise_type13final_suspendEv(
   // CHECK: call void @_ZN14suspend_always12await_resumeEv(
 
   // Destroy promise, then parameter copies:
-  // CHECK: call void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopyEE12promise_typeD1Ev(ptr {{[^,]*}} %__promise)
+  // CHECK: call void @_ZNSt16coroutine_traitsIJvi8MoveOnly11MoveAndCopy10TrivialABIEE12promise_typeD1Ev(ptr {{[^,]*}} %__promise)
+  // CHECK-NEXT: call void @llvm.lifetime.end.p0(
+  // CHECK-NEXT: call void @_ZN10TrivialABID1Ev(ptr {{[^,]*}} %[[TrivialCopy]])
   // CHECK-NEXT: call void @llvm.lifetime.end.p0(
   // CHECK-NEXT: call void @_ZN11MoveAndCopyD1Ev(ptr {{[^,]*}} %[[McCopy]])
   // CHECK-NEXT: call void @llvm.lifetime.end.p0(
   // CHECK-NEXT: call void @_ZN8MoveOnlyD1Ev(ptr {{[^,]*}} %[[MoCopy]]
   // CHECK-NEXT: call void @llvm.lifetime.end.p0(
   // CHECK-NEXT: call void @llvm.lifetime.end.p0(
   // CHECK-NEXT: call ptr @llvm.coro.free(
+
+  // The original trivial_abi parameter is destroyed when returning from the ramp.
+  // CHECK: call i1 @llvm.coro.end
+  // CHECK: call void @_ZN10TrivialABID1Ev(ptr {{[^,]*}} %[[TrivialAlloca]])
 }
 
 // CHECK-LABEL: void @_Z16dependent_paramsI1A1BEvT_T0_S3_(ptr noundef %x, ptr noundef %0, ptr noundef %y)
@@ -190,3 +210,38 @@ method some_class::good_coroutine_calls_custom_constructor(float) {
   // CHECK: invoke void @_ZNSt16coroutine_traitsIJ6methodR10some_classfEE12promise_typeC1ES2_f(ptr {{[^,]*}} %__promise, ptr noundef nonnull align 1 dereferenceable(1) %{{.+}}, float
   co_return;
 }
+
+
+struct MSParm {
+  int val;
+  ~MSParm();
+};
+
+void consume(int) noexcept;
+
+// Similarly to the [[clang::trivial_abi]] parameters, with the MSVC ABI
+// parameters are also destroyed by the callee, and on x86-64 such parameters
+// may get passed in registers. In that case it's again important that the
+// parameter's local alloca does not become part of the coro frame since that
+// may be destroyed before the destructor call.
+void msabi(MSParm p) {
+  // MSABI: define{{.*}} void @"?msabi@@YAXUMSParm@@@Z"(i32 %[[Param:.+]])
+
+  // The parameter's local alloca is marked not part of the frame.
+  // MSABI: %[[ParamAlloca:.+]] = alloca %struct.MSParm
+  // MSABI-SAME: !coro.outside.frame
+
+  // MSABI: %[[ParamCopy:.+]] = alloca %struct.MSParm
+
+  consume(p.val);
+  // The parameter's copy is used by the coroutine.
+  // MSABI: %[[ValPtr:.+]] = getelementptr inbounds nuw %struct.MSParm, ptr %[[ParamCopy]], i32 0, i32 0
+  // MSABI: %[[Val:.+]] = load i32, ptr %[[ValPtr]]
+  // MSABI: call void @"?consume@@YAXH@Z"(i32{{.*}} %[[Val]])
+
+  co_return;
+
+  // The local alloca is used for the destructor call at the end of the ramp.
+  // MSABI: call i1 @llvm.coro.end
+  // MSABI: call void @"??1MSParm@@QEAA@XZ"(ptr{{.*}} %[[ParamAlloca]])
+}