[CIR] Emit promise declaration in coroutine (#166683)

This PR adds support for emitting the promise declaration in coroutines
and obtaining the `get_return_object()`.

Author: Andres-Salamanca

clang/include/clang/CIR/MissingFeatures.h

@@ -153,6 +153,7 @@ struct MissingFeatures {
   static bool coroEndBuiltinCall() { return false; }
   static bool coroutineFrame() { return false; }
   static bool emitBodyAndFallthrough() { return false; }
+  static bool coroOutsideFrameMD() { return false; }
 
   // Various handling of deferred processing in CIRGenModule.
   static bool cgmRelease() { return false; }

clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp

@@ -13,6 +13,7 @@
 #include "CIRGenFunction.h"
 #include "mlir/Support/LLVM.h"
 #include "clang/AST/StmtCXX.h"
+#include "clang/AST/StmtVisitor.h"
 #include "clang/Basic/TargetInfo.h"
 #include "clang/CIR/Dialect/IR/CIRTypes.h"
 #include "clang/CIR/MissingFeatures.h"
@@ -33,6 +34,65 @@ struct clang::CIRGen::CGCoroData {
 CIRGenFunction::CGCoroInfo::CGCoroInfo() {}
 CIRGenFunction::CGCoroInfo::~CGCoroInfo() {}
 
+namespace {
+// FIXME: both GetParamRef and ParamReferenceReplacerRAII are good template
+// candidates to be shared among LLVM / CIR codegen.
+
+// Hunts for the parameter reference in the parameter copy/move declaration.
+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 (Stmt *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.
+struct ParamReferenceReplacerRAII {
+  CIRGenFunction::DeclMapTy savedLocals;
+  CIRGenFunction::DeclMapTy &localDeclMap;
+
+  ParamReferenceReplacerRAII(CIRGenFunction::DeclMapTy &localDeclMap)
+      : localDeclMap(localDeclMap) {}
+
+  void addCopy(const DeclStmt *pm) {
+    // Figure out what param it refers to.
+
+    assert(pm->isSingleDecl());
+    const VarDecl *vd = static_cast<const VarDecl *>(pm->getSingleDecl());
+    const Expr *initExpr = vd->getInit();
+    GetParamRef visitor;
+    visitor.Visit(const_cast<Expr *>(initExpr));
+    assert(visitor.expr);
+    DeclRefExpr *dreOrig = 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});
+    }
+  }
+};
+} // namespace
 static void createCoroData(CIRGenFunction &cgf,
                            CIRGenFunction::CGCoroInfo &curCoro,
                            cir::CallOp coroId) {
@@ -149,7 +209,47 @@ CIRGenFunction::emitCoroutineBody(const CoroutineBodyStmt &s) {
   if (s.getReturnStmtOnAllocFailure())
     cgm.errorNYI("handle coroutine return alloc failure");
 
-  assert(!cir::MissingFeatures::generateDebugInfo());
-  assert(!cir::MissingFeatures::emitBodyAndFallthrough());
+  {
+    assert(!cir::MissingFeatures::generateDebugInfo());
+    ParamReferenceReplacerRAII paramReplacer(localDeclMap);
+    // Create mapping between parameters and copy-params for coroutine
+    // function.
+    llvm::ArrayRef<const Stmt *> paramMoves = s.getParamMoves();
+    assert((paramMoves.size() == 0 || (paramMoves.size() == fnArgs.size())) &&
+           "ParamMoves and FnArgs should be the same size for coroutine "
+           "function");
+    // For zipping the arg map into debug info.
+    assert(!cir::MissingFeatures::generateDebugInfo());
+
+    // Create parameter copies. We do it before creating a promise, since an
+    // evolution of coroutine TS may allow promise constructor to observe
+    // parameter copies.
+    assert(!cir::MissingFeatures::coroOutsideFrameMD());
+    for (auto *pm : paramMoves) {
+      if (emitStmt(pm, /*useCurrentScope=*/true).failed())
+        return mlir::failure();
+      paramReplacer.addCopy(cast<DeclStmt>(pm));
+    }
+
+    if (emitStmt(s.getPromiseDeclStmt(), /*useCurrentScope=*/true).failed())
+      return mlir::failure();
+    // returnValue should be valid as long as the coroutine's return type
+    // is not void. The assertion could help us to reduce the check later.
+    assert(returnValue.isValid() == (bool)s.getReturnStmt());
+    // Now we have the promise, initialize the GRO.
+    // We need to emit `get_return_object` first. According to:
+    // [dcl.fct.def.coroutine]p7
+    // The call to get_return_­object is sequenced before the call to
+    // initial_suspend and is invoked at most once.
+    //
+    // So we couldn't emit return value when we emit return statment,
+    // otherwise the call to get_return_object wouldn't be in front
+    // of initial_suspend.
+    if (returnValue.isValid())
+      emitAnyExprToMem(s.getReturnValue(), returnValue,
+                       s.getReturnValue()->getType().getQualifiers(),
+                       /*isInit*/ true);
+    assert(!cir::MissingFeatures::emitBodyAndFallthrough());
+  }
   return mlir::success();
 }

clang/lib/CIR/CodeGen/CIRGenFunction.cpp

@@ -632,6 +632,10 @@ cir::FuncOp CIRGenFunction::generateCode(clang::GlobalDecl gd, cir::FuncOp fn,
 
     startFunction(gd, retTy, fn, funcType, args, loc, bodyRange.getBegin());
 
+    // Save parameters for coroutine function.
+    if (body && isa_and_nonnull<CoroutineBodyStmt>(body))
+      llvm::append_range(fnArgs, funcDecl->parameters());
+
     if (isa<CXXDestructorDecl>(funcDecl)) {
       emitDestructorBody(args);
     } else if (isa<CXXConstructorDecl>(funcDecl)) {

clang/lib/CIR/CodeGen/CIRGenFunction.h

@@ -152,6 +152,9 @@ class CIRGenFunction : public CIRGenTypeCache {
   /// global initializers.
   mlir::Operation *curFn = nullptr;
 
+  /// Save Parameter Decl for coroutine.
+  llvm::SmallVector<const ParmVarDecl *> fnArgs;
+
   using DeclMapTy = llvm::DenseMap<const clang::Decl *, Address>;
   /// This keeps track of the CIR allocas or globals for local C
   /// declarations.

clang/test/CIR/CodeGen/coro-task.cpp

@@ -36,6 +36,12 @@ struct suspend_never {
   void await_resume() noexcept {}
 };
 
+struct string {
+  int size() const;
+  string();
+  string(char const *s);
+};
+
 } // namespace std
 
 namespace folly {
@@ -101,7 +107,10 @@ co_invoke_fn co_invoke;
 }} // namespace folly::coro
 
 // CIR-DAG: ![[VoidTask:.*]] = !cir.record<struct "folly::coro::Task<void>" padded {!u8i}>
-
+// CIR-DAG: ![[IntTask:.*]] = !cir.record<struct "folly::coro::Task<int>" padded {!u8i}>
+// CIR-DAG: ![[VoidPromisse:.*]] = !cir.record<struct "folly::coro::Task<void>::promise_type" padded {!u8i}>
+// CIR-DAG: ![[IntPromisse:.*]] = !cir.record<struct "folly::coro::Task<int>::promise_type" padded {!u8i}>
+// CIR-DAG: ![[StdString:.*]] = !cir.record<struct "std::string" padded {!u8i}>
 // CIR: module {{.*}} {
 // CIR-NEXT: cir.global external @_ZN5folly4coro9co_invokeE = #cir.zero : !rec_folly3A3Acoro3A3Aco_invoke_fn
 
@@ -119,6 +128,7 @@ VoidTask silly_task() {
 // CIR: cir.func coroutine dso_local @_Z10silly_taskv() -> ![[VoidTask]]
 // CIR: %[[VoidTaskAddr:.*]] = cir.alloca ![[VoidTask]], {{.*}}, ["__retval"]
 // CIR: %[[SavedFrameAddr:.*]] = cir.alloca !cir.ptr<!void>, !cir.ptr<!cir.ptr<!void>>, ["__coro_frame_addr"]
+// CIR: %[[VoidPromisseAddr:.*]] = cir.alloca ![[VoidPromisse]], {{.*}}, ["__promise"]
 
 // Get coroutine id with __builtin_coro_id.
 
@@ -138,3 +148,27 @@ VoidTask silly_task() {
 // CIR: }
 // CIR: %[[Load0:.*]] = cir.load{{.*}} %[[SavedFrameAddr]] : !cir.ptr<!cir.ptr<!void>>, !cir.ptr<!void>
 // CIR: %[[CoroFrameAddr:.*]] = cir.call @__builtin_coro_begin(%[[CoroId]], %[[Load0]])
+
+// Call promise.get_return_object() to retrieve the task object.
+
+// CIR: %[[RetObj:.*]] = cir.call @_ZN5folly4coro4TaskIvE12promise_type17get_return_objectEv(%[[VoidPromisseAddr]]) nothrow : {{.*}} -> ![[VoidTask]]
+// CIR: cir.store{{.*}} %[[RetObj]], %[[VoidTaskAddr]] : ![[VoidTask]]
+
+folly::coro::Task<int> byRef(const std::string& s) {
+  co_return s.size();
+}
+
+// CIR:  cir.func coroutine dso_local @_Z5byRefRKSt6string(%[[ARG:.*]]: !cir.ptr<![[StdString]]> {{.*}}) -> ![[IntTask]]
+// CIR:    %[[AllocaParam:.*]] = cir.alloca !cir.ptr<![[StdString]]>, {{.*}}, ["s", init, const]
+// CIR:    %[[IntTaskAddr:.*]] = cir.alloca ![[IntTask]], {{.*}}, ["__retval"]
+// CIR:    %[[SavedFrameAddr:.*]]  = cir.alloca !cir.ptr<!void>, !cir.ptr<!cir.ptr<!void>>, ["__coro_frame_addr"]
+// CIR:    %[[AllocaFnUse:.*]] = cir.alloca !cir.ptr<![[StdString]]>, {{.*}}, ["s", init, const]
+// CIR:    %[[IntPromisseAddr:.*]] = cir.alloca ![[IntPromisse]], {{.*}}, ["__promise"]
+// CIR:    cir.store %[[ARG]], %[[AllocaParam]] : !cir.ptr<![[StdString]]>, {{.*}}
+
+// Call promise.get_return_object() to retrieve the task object.
+
+// CIR:    %[[LOAD:.*]] = cir.load %[[AllocaParam]] : !cir.ptr<!cir.ptr<![[StdString]]>>, !cir.ptr<![[StdString]]>
+// CIR:    cir.store {{.*}} %[[LOAD]], %[[AllocaFnUse]] : !cir.ptr<![[StdString]]>, !cir.ptr<!cir.ptr<![[StdString]]>>
+// CIR:    %[[RetObj:.*]] = cir.call @_ZN5folly4coro4TaskIiE12promise_type17get_return_objectEv(%4) nothrow : {{.*}} -> ![[IntTask]]
+// CIR:    cir.store {{.*}} %[[RetObj]], %[[IntTaskAddr]] : ![[IntTask]]