[coroutines] Add support for coroutines with non-scalar parameters

[GorNishanov]

Summary

Simple types like int are handled by LLVM Coroutines just fine.
But for non-scalar parameters we need to create copies of those parameters in the coroutine frame and make all uses of those parameters to refer to parameter copies.

[GorNishanov]

LGTM

lib/CodeGen/CGCoroutine.cpp

 // This contains code dealing with C++ code generation of coroutines.
 //
 //===----------------------------------------------------------------------===//
 
 #include "CGCleanup.h"
 #include "CodeGenFunction.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "clang/AST/StmtCXX.h"
+#include "clang/AST/StmtVisitor.h"
 
 using namespace clang;
 using namespace CodeGen;
 
 using llvm::Value;
 using llvm::BasicBlock;
 
 namespace {

 }
 
 void CodeGenFunction::EmitCoreturnStmt(CoreturnStmt const &S) {
   ++CurCoro.Data->CoreturnCount;
   EmitStmt(S.getPromiseCall());
   EmitBranchThroughCleanup(CurCoro.Data->FinalJD);
 }
 
-// For WinEH exception representation backend need to know what funclet coro.end
+// Hunts for the parameter reference in the parameter copy/move declaration.
+namespace {
+struct GetParamRef : public StmtVisitor<GetParamRef> {
+public:
+  DeclRefExpr *Expr = nullptr;
+  GetParamRef() {}
+  void VisitDeclRefExpr(DeclRefExpr *E) {
+    assert(Expr == nullptr && "multilple declref in param move");
+    Expr = E;
+  }
+  void VisitStmt(Stmt *S) {
+    for (auto *C : S->children()) {
+      if (C)
+        Visit(C);
+    }
+  }
+};
+}
+
+// This class replaces references to parameters to their copies by changing
+// the addresses in CGF.LocalDeclMap and restoring back the original values in
+// its destructor.
+
+namespace {
+  struct ParamReferenceReplacerRAII {
+    CodeGenFunction::DeclMapTy SavedLocals;
+    CodeGenFunction::DeclMapTy& LocalDeclMap;
+
+    ParamReferenceReplacerRAII(CodeGenFunction::DeclMapTy &LocalDeclMap)
+        : LocalDeclMap(LocalDeclMap) {}
+
+    void addCopy(DeclStmt const *PM) {
+      // Figure out what param it refers to.
+
+      assert(PM->isSingleDecl());
+      VarDecl const*VD = static_cast<VarDecl const*>(PM->getSingleDecl());
+      Expr const *InitExpr = VD->getInit();
+      GetParamRef Visitor;
+      Visitor.Visit(const_cast<Expr*>(InitExpr));
+      assert(Visitor.Expr);
+      auto *DREOrig = cast<DeclRefExpr>(Visitor.Expr);
+      auto *PD = DREOrig->getDecl();
+
+      auto it = LocalDeclMap.find(PD);
+      assert(it != LocalDeclMap.end() && "parameter is not found");
+      SavedLocals.insert({ PD, it->second });
+
+      auto copyIt = LocalDeclMap.find(VD);
+      assert(copyIt != LocalDeclMap.end() && "parameter copy is not found");
+      it->second = copyIt->getSecond();
+    }
+
+    ~ParamReferenceReplacerRAII() {
+      for (auto&& SavedLocal : SavedLocals) {
+        LocalDeclMap.insert({SavedLocal.first, SavedLocal.second});
+      }
+    }
+  };
+}
+
+// For WinEH exception representation backend needs to know what funclet coro.end
 // belongs to. That information is passed in a funclet bundle.
 static SmallVector<llvm::OperandBundleDef, 1>
 getBundlesForCoroEnd(CodeGenFunction &CGF) {
   SmallVector<llvm::OperandBundleDef, 1> BundleList;
 
   if (llvm::Instruction *EHPad = CGF.CurrentFuncletPad)
     BundleList.emplace_back("funclet", EHPad);
 
auto *CoroBegin = Builder.CreateCall(
       CGM.getIntrinsic(llvm::Intrinsic::coro_begin), {CoroId, Phi});
   CurCoro.Data->CoroBegin = CoroBegin;
 
   GetReturnObjectManager GroManager(*this, S);
   GroManager.EmitGroAlloca();
 
   CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(RetBB);
   {
+    ParamReferenceReplacerRAII ParamReplacer(LocalDeclMap);
     CodeGenFunction::RunCleanupsScope ResumeScope(*this);
     EHStack.pushCleanup<CallCoroDelete>(NormalAndEHCleanup, S.getDeallocate());
 
+    // 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 (auto *PM : S.getParamMoves()) {
+      EmitStmt(PM);
+      ParamReplacer.addCopy(cast<DeclStmt>(PM));
+      // TODO: if(CoroParam(...)) need to surround ctor and dtor
+      // for the copy, so that llvm can elide it if the copy is
+      // not needed.
+    }
+
     EmitStmt(S.getPromiseDeclStmt());
 
+    Address PromiseAddr = GetAddrOfLocalVar(S.getPromiseDecl());
+    auto *PromiseAddrVoidPtr =
+        new llvm::BitCastInst(PromiseAddr.getPointer(), VoidPtrTy, "", CoroId);
+    // Update CoroId to refer to the promise. We could not do it earlier because
+    // promise local variable was not emitted yet.
+    CoroId->setArgOperand(1, PromiseAddrVoidPtr);
+
     // Now we have the promise, initialize the GRO
     GroManager.EmitGroInit();
-    EHStack.pushCleanup<CallCoroEnd>(EHCleanup);
 
-    CurCoro.Data->FinalJD = getJumpDestInCurrentScope(FinalBB);
-
-    // FIXME: Emit param moves.
+    EHStack.pushCleanup<CallCoroEnd>(EHCleanup);
 
     CurCoro.Data->CurrentAwaitKind = AwaitKind::Init;
     EmitStmt(S.getInitSuspendStmt());
+    CurCoro.Data->FinalJD = getJumpDestInCurrentScope(FinalBB);
 
     CurCoro.Data->CurrentAwaitKind = AwaitKind::Normal;
 
     if (auto *OnException = S.getExceptionHandler()) {
       auto Loc = S.getLocStart();
       CXXCatchStmt Catch(Loc, /*exDecl=*/nullptr, OnException);
       auto *TryStmt = CXXTryStmt::Create(getContext(), Loc, S.getBody(), &Catch);
 
else if (IID == llvm::Intrinsic::coro_begin) {
     if (CurCoro.Data)
       CurCoro.Data->CoroBegin = Call;
   }
   else if (IID == llvm::Intrinsic::coro_free) {
     // Remember the last coro_free as we need it to build the conditional
     // deletion of the coroutine frame.
     if (CurCoro.Data)
       CurCoro.Data->LastCoroFree = Call;
-  }  return RValue::get(Call);
+  }
+  return RValue::get(Call);
 }

lib/Sema/SemaCoroutine.cpp

if (ReturnStmt.isInvalid()) {
     noteMemberDeclaredHere(S, ReturnValue, Fn);
     return false;
   }
 
   this->ReturnStmt = ReturnStmt.get();
   return true;
 }
 
+// Create a static_cast\<T&&>(expr).
+static Expr *castForMoving(Sema &S, Expr *E, QualType T = QualType()) {
+  if (T.isNull())
+    T = E->getType();
+  QualType TargetType = S.BuildReferenceType(
+      T, /*SpelledAsLValue*/ false, SourceLocation(), DeclarationName());
+  SourceLocation ExprLoc = E->getLocStart();
+  TypeSourceInfo *TargetLoc =
+      S.Context.getTrivialTypeSourceInfo(TargetType, ExprLoc);
+
+  return S
+      .BuildCXXNamedCast(ExprLoc, tok::kw_static_cast, TargetLoc, E,
+                         SourceRange(ExprLoc, ExprLoc), E->getSourceRange())
+      .get();
+}
+
+/// \brief Build a variable declaration for move parameter.
+static VarDecl *buildVarDecl(Sema &S, SourceLocation Loc, QualType Type,
+                             StringRef Name) {
+  DeclContext *DC = S.CurContext;
+  IdentifierInfo *II = &S.PP.getIdentifierTable().get(Name);
+  TypeSourceInfo *TInfo = S.Context.getTrivialTypeSourceInfo(Type, Loc);
+  VarDecl *Decl =
+      VarDecl::Create(S.Context, DC, Loc, Loc, II, Type, TInfo, SC_None);
+  Decl->setImplicit();
+  return Decl;
+}
+
 bool CoroutineStmtBuilder::makeParamMoves() {
-  // FIXME: Perform move-initialization of parameters into frame-local copies.
+  for (auto *paramDecl : FD.parameters()) {
+    auto Ty = paramDecl->getType();
+    if (Ty->isDependentType())
+      continue;
+
+    // No need to copy scalars, llvm will take care of them.
+    if (Ty->getAsCXXRecordDecl()) {
+      if (!paramDecl->getIdentifier())
+        continue;
+
+      ExprResult ParamRef =
+          S.BuildDeclRefExpr(paramDecl, paramDecl->getType(),
+                             ExprValueKind::VK_LValue, Loc); // FIXME: scope?
+      if (ParamRef.isInvalid())
+        return false;
+
+      Expr *RCast = castForMoving(S, ParamRef.get());
+
+      auto D = buildVarDecl(S, Loc, Ty, paramDecl->getIdentifier()->getName());
+
+      S.AddInitializerToDecl(D, RCast, /*DirectInit=*/true);
+
+      // Convert decl to a statement.
+      StmtResult Stmt = S.ActOnDeclStmt(S.ConvertDeclToDeclGroup(D), Loc, Loc);
+      if (Stmt.isInvalid())
+        return false;
+
+      ParamMovesVector.push_back(Stmt.get());
+    }
+  }
+
+  // Convert to ArrayRef in CtorArgs structure that builder inherits from.
+  ParamMoves = ParamMovesVector;
   return true;
 }
 
 StmtResult Sema::BuildCoroutineBodyStmt(CoroutineBodyStmt::CtorArgs Args) {
   CoroutineBodyStmt *Res = CoroutineBodyStmt::Create(Context, Args);
   if (!Res)
     return StmtError();
   return Res;
 }

test/CodeGenCoroutines/coro-params.cpp

+// Verifies that parameters are copied with move constructors
+// Verifies that parameter copies are destroyed
+// Vefifies that parameter copies are used in the body of the coroutine
+// RUN: %clang_cc1 -std=c++1z -fcoroutines-ts -triple=x86_64-unknown-linux-gnu -emit-llvm -o - %s -disable-llvm-passes -fexceptions | FileCheck %s
+
+namespace std::experimental {
+template <typename... T> struct coroutine_traits;
+
+template <class Promise = void> struct coroutine_handle {
+  coroutine_handle() = default;
+  static coroutine_handle from_address(void *) noexcept;
+};
+template <> struct coroutine_handle<void> {
+  static coroutine_handle from_address(void *) noexcept;
+  coroutine_handle() = default;
+  template <class PromiseType>
+  coroutine_handle(coroutine_handle<PromiseType>) noexcept;
+};
+}
+
+struct suspend_always {
+  bool await_ready() noexcept;
+  void await_suspend(std::experimental::coroutine_handle<>) noexcept;
+  void await_resume() noexcept;
+};
+
+template <typename... Args> struct std::experimental::coroutine_traits<void, Args...> {
+  struct promise_type {
+    void get_return_object() noexcept;
+    suspend_always initial_suspend() noexcept;
+    suspend_always final_suspend() noexcept;
+    void return_void() noexcept;
+    promise_type();
+    ~promise_type() noexcept;
+    void unhandled_exception() noexcept;
+  };
+};
+
+// TODO: Not supported yet
+struct CopyOnly {
+  int val;
+  CopyOnly(const CopyOnly&) noexcept;
+  CopyOnly(CopyOnly&&) = delete;
+  ~CopyOnly();
+};
+
+struct MoveOnly {
+  int val;
+  MoveOnly(const MoveOnly&) = delete;
+  MoveOnly(MoveOnly&&) noexcept;
+  ~MoveOnly();
+};
+
+struct MoveAndCopy {
+  int val;
+  MoveAndCopy(const MoveAndCopy&)noexcept;
+  MoveAndCopy(MoveAndCopy&&) noexcept;
+  ~MoveAndCopy();
+};
+
+void consume(int,int,int) noexcept;
+
+// TODO: Add support for CopyOnly params
+// CHECK: define void @_Z1fi8MoveOnly11MoveAndCopy(i32 %val, %struct.MoveOnly* %[[MoParam:.+]], %struct.MoveAndCopy* %[[McParam:.+]]) #0 personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*
+void f(int val, MoveOnly moParam, MoveAndCopy mcParam) {
+  // CHECK: %[[MoCopy:.+]] = alloca %struct.MoveOnly
+  // CHECK: %[[McCopy:.+]] = alloca %struct.MoveAndCopy
+  // CHECK: store i32 %val, i32* %[[ValAddr:.+]]
+
+  // CHECK: call i8* @llvm.coro.begin(
+  // CHECK-NEXT: call void @_ZN8MoveOnlyC1EOS_(%struct.MoveOnly* %[[MoCopy]], %struct.MoveOnly* dereferenceable(4) %[[MoParam]])
+  // CHECK-NEXT: call void @_ZN11MoveAndCopyC1EOS_(%struct.MoveAndCopy* %[[McCopy]], %struct.MoveAndCopy* dereferenceable(4) %[[McParam]]) #
+  // CHECK-NEXT: invoke void @_ZNSt12experimental16coroutine_traitsIJvi8MoveOnly11MoveAndCopyEE12promise_typeC1Ev(
+
+  // CHECK: call void @_ZN14suspend_always12await_resumeEv(
+  // CHECK: %[[IntParam:.+]] = load i32, i32* %val.addr
+  // CHECK: %[[MoGep:.+]] = getelementptr inbounds %struct.MoveOnly, %struct.MoveOnly* %[[MoCopy]], i32 0, i32 0
+  // CHECK: %[[MoVal:.+]] = load i32, i32* %[[MoGep]]
+  // CHECK: %[[McGep:.+]] =  getelementptr inbounds %struct.MoveAndCopy, %struct.MoveAndCopy* %[[McCopy]], i32 0, i32 0
+  // CHECK: %[[McVal:.+]] = load i32, i32* %[[McGep]]
+  // CHECK: call void @_Z7consumeiii(i32 %[[IntParam]], i32 %[[MoVal]], i32 %[[McVal]])
+
+  consume(val, moParam.val, mcParam.val);
+  co_return;
+
+  // Skip to final suspend:
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJvi8MoveOnly11MoveAndCopyEE12promise_type13final_suspendEv(
+  // CHECK: call void @_ZN14suspend_always12await_resumeEv(
+
+  // Destroy promise, then parameter copies:
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJvi8MoveOnly11MoveAndCopyEE12promise_typeD1Ev(%"struct.std::experimental::coroutine_traits<void, int, MoveOnly, MoveAndCopy>::promise_type"* %__promise) #2
+  // CHECK-NEXT: call void @_ZN11MoveAndCopyD1Ev(%struct.MoveAndCopy* %[[McCopy]])
+  // CHECK-NEXT: call void @_ZN8MoveOnlyD1Ev(%struct.MoveOnly* %[[MoCopy]]
+  // CHECK-NEXT: call i8* @llvm.coro.free(
+}