Fix codegen for coroutine with function-try-block

[MatzeB]

Summary

This fixes an assertion error when writing a coroutine with a function-try-block. In this case the function body is not a CompoundStmt so the code constructing an artificial CXXTryStmt must also construct a CompoundStmt for it.

While on it adjust the CXXStmt::Create function to only accept CompoundStmt*.

[bruno]

LGTM! @ChuanqiXu wdyt?

clang/test/CodeGenCoroutines/coro-function-try-block.cpp

2

Space after -triple=x86_64

[MatzeB]

clang/test/CodeGenCoroutines/coro-function-try-block.cpp

2

it is -triple=x86_64-- (lazy mans way of writing -triple=x86_64-unknown-unknown)

[ChuanqiXu]

Thanks for finding this! We didn't notice this before.

clang/lib/CodeGen/CGCoroutine.cpp

736–738

Can we try to move the logic to CoroutineStmtBuilder? That makes me feel better. And it will be helpful to add a comment to tell that we're handling the case the function body is function-try-block.

738

It reads better to specify the potential type for Body.

clang/test/CodeGenCoroutines/coro-function-try-block.cpp

22–24

I expect to see the nested try statements in the case.

[ChuanqiXu]

clang/lib/CodeGen/CGCoroutine.cpp

736–738

It looks like you didn't address the comments. Would you like to address it? I don't mind to address it later myself.

[MatzeB]

clang/lib/CodeGen/CGCoroutine.cpp

736–738

I'll add a detailed comment. But would you be fine leaving the statements here as-is? The logic only makes sense in the context of using the Body to create a CXXTryStmt below (it's really an effect of CXXTryStmt only accepting CompountStmt operands).

736–738

Did you mean to create a new function named CoroutineStmtBuilder like I did now?

738

I will mention a single-statement body as an example in the comment now. Putting an assert here feels unnecessary and may be in the way if in the future we ever allow other types of single-statement function bodies.

clang/test/CodeGenCoroutines/coro-function-try-block.cpp

22–24

This was mostly meant as a test for "does not crash the compiler", so it is just checking some parts of the output without caring too much what they are specifically.

I cannot test for nested try statements, because in llvm-ir those are already lowered to unwind tables and landing pad blocks, but there's neither a "try" statement nor an immediate nesting I could test for.

[ChuanqiXu]

Since the patch itself is good and not large. Let me handle the trivial refactoring later.

clang/lib/CodeGen/CGCoroutine.cpp

736–738

Did you mean to create a new function named CoroutineStmtBuilder like I did now?

No, I mean we should construct this in Sema.

Putting an assert here feels unnecessary and may be in the way if in the future we ever allow other types of single-statement function bodies.

Personally I prefer the more precise style.

[MatzeB]

clang/lib/CodeGen/CGCoroutine.cpp

736–738

No, I mean we should construct this in Sema.

I'll land the patch as-is then and leave the refactoring to you if necessary.

clang/include/clang/AST/StmtCXX.h

class CXXTryStmt final : public Stmt,
 
   friend TrailingObjects;
   friend class ASTStmtReader;
 
   SourceLocation TryLoc;
   unsigned NumHandlers;
   size_t numTrailingObjects(OverloadToken<Stmt *>) const { return NumHandlers; }
 
-  CXXTryStmt(SourceLocation tryLoc, Stmt *tryBlock, ArrayRef<Stmt*> handlers);
+  CXXTryStmt(SourceLocation tryLoc, CompoundStmt *tryBlock,
+             ArrayRef<Stmt *> handlers);
   CXXTryStmt(EmptyShell Empty, unsigned numHandlers)
     : Stmt(CXXTryStmtClass), NumHandlers(numHandlers) { }
 
   Stmt *const *getStmts() const { return getTrailingObjects<Stmt *>(); }
   Stmt **getStmts() { return getTrailingObjects<Stmt *>(); }
 
 public:
   static CXXTryStmt *Create(const ASTContext &C, SourceLocation tryLoc,
-                            Stmt *tryBlock, ArrayRef<Stmt*> handlers);
+                            CompoundStmt *tryBlock, ArrayRef<Stmt *> handlers);
 
   static CXXTryStmt *Create(const ASTContext &C, EmptyShell Empty,
                             unsigned numHandlers);
 
   SourceLocation getBeginLoc() const LLVM_READONLY { return getTryLoc(); }
 
   SourceLocation getTryLoc() const { return TryLoc; }
   SourceLocation getEndLoc() const {

clang/lib/AST/ASTImporter.cpp

ExpectedStmt ASTNodeImporter::VisitCXXTryStmt(CXXTryStmt *S) {
   for (unsigned HI = 0, HE = S->getNumHandlers(); HI != HE; ++HI) {
     CXXCatchStmt *FromHandler = S->getHandler(HI);
     if (auto ToHandlerOrErr = import(FromHandler))
       ToHandlers[HI] = *ToHandlerOrErr;
     else
       return ToHandlerOrErr.takeError();
   }
 
-  return CXXTryStmt::Create(
-      Importer.getToContext(), *ToTryLocOrErr,*ToTryBlockOrErr, ToHandlers);
+  return CXXTryStmt::Create(Importer.getToContext(), *ToTryLocOrErr,
+                            cast<CompoundStmt>(*ToTryBlockOrErr), ToHandlers);
 }
 
 ExpectedStmt ASTNodeImporter::VisitCXXForRangeStmt(CXXForRangeStmt *S) {
 
   Error Err = Error::success();
   auto ToInit = importChecked(Err, S->getInit());
   auto ToRangeStmt = importChecked(Err, S->getRangeStmt());
   auto ToBeginStmt = importChecked(Err, S->getBeginStmt());

clang/lib/AST/StmtCXX.cpp

 
 QualType CXXCatchStmt::getCaughtType() const {
   if (ExceptionDecl)
     return ExceptionDecl->getType();
   return QualType();
 }
 
 CXXTryStmt *CXXTryStmt::Create(const ASTContext &C, SourceLocation tryLoc,
-                               Stmt *tryBlock, ArrayRef<Stmt *> handlers) {
+                               CompoundStmt *tryBlock,
+                               ArrayRef<Stmt *> handlers) {
   const size_t Size = totalSizeToAlloc<Stmt *>(handlers.size() + 1);
   void *Mem = C.Allocate(Size, alignof(CXXTryStmt));
   return new (Mem) CXXTryStmt(tryLoc, tryBlock, handlers);
 }
 
 CXXTryStmt *CXXTryStmt::Create(const ASTContext &C, EmptyShell Empty,
                                unsigned numHandlers) {
   const size_t Size = totalSizeToAlloc<Stmt *>(numHandlers + 1);
   void *Mem = C.Allocate(Size, alignof(CXXTryStmt));
   return new (Mem) CXXTryStmt(Empty, numHandlers);
 }
 
-CXXTryStmt::CXXTryStmt(SourceLocation tryLoc, Stmt *tryBlock,
+CXXTryStmt::CXXTryStmt(SourceLocation tryLoc, CompoundStmt *tryBlock,
                        ArrayRef<Stmt *> handlers)
     : Stmt(CXXTryStmtClass), TryLoc(tryLoc), NumHandlers(handlers.size()) {
   Stmt **Stmts = getStmts();
   Stmts[0] = tryBlock;
   std::copy(handlers.begin(), handlers.end(), Stmts + 1);
 }
 
 CXXForRangeStmt::CXXForRangeStmt(Stmt *Init, DeclStmt *Range,

clang/lib/CodeGen/CGCoroutine.cpp

static void emitBodyAndFallthrough(CodeGenFunction &CGF,
                                    const CoroutineBodyStmt &S, Stmt *Body) {
   CGF.EmitStmt(Body);
   const bool CanFallthrough = CGF.Builder.GetInsertBlock();
   if (CanFallthrough)
     if (Stmt *OnFallthrough = S.getFallthroughHandler())
       CGF.EmitStmt(OnFallthrough);
 }
 
+static CompoundStmt *CoroutineStmtBuilder(ASTContext &Context,
+                                          const CoroutineBodyStmt &S) {
+  Stmt *Stmt = S.getBody();
+  if (CompoundStmt *Body = dyn_cast<CompoundStmt>(Stmt))
+    return Body;
+  // We are about to create a `CXXTryStmt` which requires a `CompoundStmt`.
+  // If the function body is not a `CompoundStmt` yet then we have to create
+  // a new one. This happens for cases like the "function-try-block" syntax.
+  return CompoundStmt::Create(Context, {Stmt}, FPOptionsOverride(),
+                              SourceLocation(), SourceLocation());
+}
+
 void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
   auto *NullPtr = llvm::ConstantPointerNull::get(Builder.getInt8PtrTy());
   auto &TI = CGM.getContext().getTargetInfo();
   unsigned NewAlign = TI.getNewAlign() / TI.getCharWidth();
 
   auto *EntryBB = Builder.GetInsertBlock();
   auto *AllocBB = createBasicBlock("coro.alloc");
   auto *InitBB = createBasicBlock("coro.init");
if (CurCoro.Data->ExceptionHandler) {
                                                  "coro.resumed.eh");
         Builder.CreateCondBr(SkipBody, ContBB, BodyBB);
         EmitBlock(BodyBB);
       }
 
       auto Loc = S.getBeginLoc();
       CXXCatchStmt Catch(Loc, /*exDecl=*/nullptr,
                          CurCoro.Data->ExceptionHandler);
-      auto *TryStmt =
-          CXXTryStmt::Create(getContext(), Loc, S.getBody(), &Catch);
+      CompoundStmt *Body = CoroutineStmtBuilder(getContext(), S);
+      auto *TryStmt = CXXTryStmt::Create(getContext(), Loc, Body, &Catch);
 
       EnterCXXTryStmt(*TryStmt);
       emitBodyAndFallthrough(*this, S, TryStmt->getTryBlock());
       ExitCXXTryStmt(*TryStmt);
 
       if (ContBB)
         EmitBlock(ContBB);
     }
     else {

clang/lib/Sema/SemaStmt.cpp

if (!R.second) {
       Diag(Problem->getExceptionDecl()->getTypeSpecStartLoc(),
            diag::note_previous_exception_handler)
           << Problem->getCaughtType();
     }
   }
 
   FSI->setHasCXXTry(TryLoc);
 
-  return CXXTryStmt::Create(Context, TryLoc, TryBlock, Handlers);
+  return CXXTryStmt::Create(Context, TryLoc, cast<CompoundStmt>(TryBlock),
+                            Handlers);
 }
 
 StmtResult Sema::ActOnSEHTryBlock(bool IsCXXTry, SourceLocation TryLoc,
                                   Stmt *TryBlock, Stmt *Handler) {
   assert(TryBlock && Handler);
 
   sema::FunctionScopeInfo *FSI = getCurFunction();
 

clang/test/CodeGenCoroutines/coro-function-try-block.cpp

+// RUN: %clang_cc1 -std=c++20 -triple=x86_64-- -emit-llvm -fcxx-exceptions \
+// RUN:            -disable-llvm-passes %s -o - | FileCheck %s
+
+#include "Inputs/coroutine.h"
+
+struct task {
+  struct promise_type {
+    task get_return_object();
+    std::suspend_never initial_suspend();
+    std::suspend_never final_suspend() noexcept;
+    void return_void();
+    void unhandled_exception() noexcept;
+  };
+};
+
+task f() try {
+  co_return;
+} catch(...) {
+}
+
+// CHECK-LABEL: define{{.*}} void @_Z1fv(
+// CHECK: call void @_ZNSt13suspend_never13await_suspendESt16coroutine_handleIvE(
+// CHECK: call void @_ZN4task12promise_type11return_voidEv(