[coroutines] Fix re-building of dependent coroutine promise_type

[EricWF]

Summary

When using TreeTransform to rebuild a coroutine the coroutine promise_type variable is not transformed because it's stored in the FunctionScopeInfo and not as a regular sub-statement of the function.

This patch attempts to fix this by rebuilding the promise variable at the start of transforming a CoroutineBodyStmt. Additionally this patch changes TransformCoroutineBodyStmt so that it always re-builds the coroutine body. Is there a better alternative to always rebuilding the body? Or is always rebuilding sufficient for now?

Note: This patch also moves the location of some diagnostics from the location of a co_foo keyword to the function declaration.

[EricWF]

This is all wrong. Sorry for the noise.

include/clang/Sema/Sema.h

   ExprResult ActOnCoawaitExpr(Scope *S, SourceLocation KwLoc, Expr *E);
   ExprResult ActOnCoyieldExpr(Scope *S, SourceLocation KwLoc, Expr *E);
   StmtResult ActOnCoreturnStmt(SourceLocation KwLoc, Expr *E);
 
   ExprResult BuildCoawaitExpr(SourceLocation KwLoc, Expr *E);
   ExprResult BuildCoyieldExpr(SourceLocation KwLoc, Expr *E);
   StmtResult BuildCoreturnStmt(SourceLocation KwLoc, Expr *E);
 
-  void CheckCompletedCoroutineBody(FunctionDecl *FD, Stmt *&Body);
+  VarDecl *buildCoroutinePromise(SourceLocation KWLoc);
+
+  StmtResult ActOnFinishCoroutineBody(FunctionDecl *FD, Stmt *Body);
 
   //===--------------------------------------------------------------------===//
   // OpenMP directives and clauses.
   //
 private:
   void *VarDataSharingAttributesStack;
   /// Set to true inside '#pragma omp declare target' region.
   bool IsInOpenMPDeclareTargetContext = false;

lib/Sema/SemaCoroutine.cpp

 #include "clang/Sema/Initialization.h"
 #include "clang/Sema/Overload.h"
 using namespace clang;
 using namespace sema;
 
 /// Look up the std::coroutine_traits<...>::promise_type for the given
 /// function type.
 static QualType lookupPromiseType(Sema &S, const FunctionProtoType *FnType,
-                                  SourceLocation Loc) {
+                                  SourceLocation KWLoc,
+                                  SourceLocation FuncLoc) {
   // FIXME: Cache std::coroutine_traits once we've found it.
   NamespaceDecl *StdExp = S.lookupStdExperimentalNamespace();
   if (!StdExp) {
-    S.Diag(Loc, diag::err_implied_std_coroutine_traits_not_found);
+    S.Diag(KWLoc, diag::err_implied_std_coroutine_traits_not_found);
     return QualType();
   }
 
   LookupResult Result(S, &S.PP.getIdentifierTable().get("coroutine_traits"),
-                      Loc, Sema::LookupOrdinaryName);
+                      FuncLoc, Sema::LookupOrdinaryName);
   if (!S.LookupQualifiedName(Result, StdExp)) {
-    S.Diag(Loc, diag::err_implied_std_coroutine_traits_not_found);
+    S.Diag(KWLoc, diag::err_implied_std_coroutine_traits_not_found);
     return QualType();
   }
 
   ClassTemplateDecl *CoroTraits = Result.getAsSingle<ClassTemplateDecl>();
   if (!CoroTraits) {
     Result.suppressDiagnostics();
     // We found something weird. Complain about the first thing we found.
     NamedDecl *Found = *Result.begin();
     S.Diag(Found->getLocation(), diag::err_malformed_std_coroutine_traits);
     return QualType();
   }
 
   // Form template argument list for coroutine_traits<R, P1, P2, ...>.
-  TemplateArgumentListInfo Args(Loc, Loc);
+  TemplateArgumentListInfo Args(FuncLoc, FuncLoc);
   Args.addArgument(TemplateArgumentLoc(
       TemplateArgument(FnType->getReturnType()),
-      S.Context.getTrivialTypeSourceInfo(FnType->getReturnType(), Loc)));
+      S.Context.getTrivialTypeSourceInfo(FnType->getReturnType(), FuncLoc)));
   // FIXME: If the function is a non-static member function, add the type
   // of the implicit object parameter before the formal parameters.
   for (QualType T : FnType->getParamTypes())
     Args.addArgument(TemplateArgumentLoc(
-        TemplateArgument(T), S.Context.getTrivialTypeSourceInfo(T, Loc)));
+        TemplateArgument(T), S.Context.getTrivialTypeSourceInfo(T, FuncLoc)));
 
   // Build the template-id.
   QualType CoroTrait =
-      S.CheckTemplateIdType(TemplateName(CoroTraits), Loc, Args);
+      S.CheckTemplateIdType(TemplateName(CoroTraits), FuncLoc, Args);
   if (CoroTrait.isNull())
     return QualType();
-  if (S.RequireCompleteType(Loc, CoroTrait,
+  if (S.RequireCompleteType(FuncLoc, CoroTrait,
                             diag::err_coroutine_traits_missing_specialization))
     return QualType();
 
   CXXRecordDecl *RD = CoroTrait->getAsCXXRecordDecl();
   assert(RD && "specialization of class template is not a class?");
 
   // Look up the ::promise_type member.
-  LookupResult R(S, &S.PP.getIdentifierTable().get("promise_type"), Loc,
+  LookupResult R(S, &S.PP.getIdentifierTable().get("promise_type"), FuncLoc,
                  Sema::LookupOrdinaryName);
   S.LookupQualifiedName(R, RD);
   auto *Promise = R.getAsSingle<TypeDecl>();
   if (!Promise) {
-    S.Diag(Loc, diag::err_implied_std_coroutine_traits_promise_type_not_found)
+    S.Diag(FuncLoc,
+           diag::err_implied_std_coroutine_traits_promise_type_not_found)
         << RD;
     return QualType();
   }
 
   // The promise type is required to be a class type.
   QualType PromiseType = S.Context.getTypeDeclType(Promise);
   if (!PromiseType->getAsCXXRecordDecl()) {
     // Use the fully-qualified name of the type.
     auto *NNS = NestedNameSpecifier::Create(S.Context, nullptr, StdExp);
     NNS = NestedNameSpecifier::Create(S.Context, NNS, false,
                                       CoroTrait.getTypePtr());
     PromiseType = S.Context.getElaboratedType(ETK_None, NNS, PromiseType);
 
-    S.Diag(Loc, diag::err_implied_std_coroutine_traits_promise_type_not_class)
+    S.Diag(FuncLoc,
+           diag::err_implied_std_coroutine_traits_promise_type_not_class)
         << PromiseType;
     return QualType();
   }
 
   return PromiseType;
 }
 
+VarDecl *Sema::buildCoroutinePromise(SourceLocation KWLoc) {
+  auto *FD = dyn_cast<FunctionDecl>(CurContext);
+  assert(FD && "Not inside a function context");
+  SourceLocation FuncLoc = FD->getLocation();
+  QualType T =
+      FD->getType()->isDependentType()
+          ? Context.DependentTy
+          : lookupPromiseType(*this, FD->getType()->castAs<FunctionProtoType>(),
+                              KWLoc, FuncLoc);
+  if (T.isNull())
+    return nullptr;
+
+  // Create and default-initialize the promise.
+  VarDecl *VD =
+      VarDecl::Create(Context, FD, FD->getLocation(), FD->getLocation(),
+                      &PP.getIdentifierTable().get("__promise"), T,
+                      Context.getTrivialTypeSourceInfo(T, FuncLoc), SC_None);
+  CheckVariableDeclarationType(VD);
+  if (!VD->isInvalidDecl())
+    ActOnUninitializedDecl(VD, false);
+  return VD;
+}
+
 /// Check that this is a context in which a coroutine suspension can appear.
 static FunctionScopeInfo *
 checkCoroutineContext(Sema &S, SourceLocation Loc, StringRef Keyword) {
   // 'co_await' and 'co_yield' are not permitted in unevaluated operands.
   if (S.isUnevaluatedContext()) {
     S.Diag(Loc, diag::err_coroutine_unevaluated_context) << Keyword;
     return nullptr;
   }
S.Diag(Loc, diag::note_declared_coroutine_here)
       << (Keyword == "co_await" ? 0 :
           Keyword == "co_yield" ? 1 : 2);
   } else {
     auto *ScopeInfo = S.getCurFunction();
     assert(ScopeInfo && "missing function scope for function");
 
     // If we don't have a promise variable, build one now.
     if (!ScopeInfo->CoroutinePromise) {
-      QualType T =
-          FD->getType()->isDependentType()
-              ? S.Context.DependentTy
-              : lookupPromiseType(S, FD->getType()->castAs<FunctionProtoType>(),
-                                  Loc);
-      if (T.isNull())
+      if (!(ScopeInfo->CoroutinePromise = S.buildCoroutinePromise(Loc)))
         return nullptr;
-
-      // Create and default-initialize the promise.
-      ScopeInfo->CoroutinePromise =
-          VarDecl::Create(S.Context, FD, FD->getLocation(), FD->getLocation(),
-                          &S.PP.getIdentifierTable().get("__promise"), T,
-                          S.Context.getTrivialTypeSourceInfo(T, Loc), SC_None);
-      S.CheckVariableDeclarationType(ScopeInfo->CoroutinePromise);
-      if (!ScopeInfo->CoroutinePromise->isInvalidDecl())
-        S.ActOnUninitializedDecl(ScopeInfo->CoroutinePromise, false);
     }
 
     return ScopeInfo;
   }
 
   return nullptr;
 }
 
StmtResult Sema::BuildCoreturnStmt(SourceLocation Loc, Expr *E) {
 
   Expr *PCE = ActOnFinishFullExpr(PC.get()).get();
 
   Stmt *Res = new (Context) CoreturnStmt(Loc, E, PCE);
   Coroutine->CoroutineStmts.push_back(Res);
   return Res;
 }
 
-void Sema::CheckCompletedCoroutineBody(FunctionDecl *FD, Stmt *&Body) {
+StmtResult Sema::ActOnFinishCoroutineBody(FunctionDecl *FD, Stmt *Body) {
   FunctionScopeInfo *Fn = getCurFunction();
   assert(Fn && !Fn->CoroutineStmts.empty() && "not a coroutine");
 
   // Coroutines [stmt.return]p1:
   //   A return statement shall not appear in a coroutine.
   if (Fn->FirstReturnLoc.isValid()) {
     Diag(Fn->FirstReturnLoc, diag::err_return_in_coroutine);
     auto *First = Fn->CoroutineStmts[0];
StmtResult Sema::ActOnFinishCoroutineBody(FunctionDecl *FD, Stmt *Body) {
 
   SourceLocation Loc = FD->getLocation();
 
   // Form a declaration statement for the promise declaration, so that AST
   // visitors can more easily find it.
   StmtResult PromiseStmt =
       ActOnDeclStmt(ConvertDeclToDeclGroup(Fn->CoroutinePromise), Loc, Loc);
   if (PromiseStmt.isInvalid())
-    return FD->setInvalidDecl();
+    return StmtError();
 
   // Form and check implicit 'co_await p.initial_suspend();' statement.
   ExprResult InitialSuspend =
       buildPromiseCall(*this, Fn, Loc, "initial_suspend", None);
   // FIXME: Support operator co_await here.
   if (!InitialSuspend.isInvalid())
     InitialSuspend = BuildCoawaitExpr(Loc, InitialSuspend.get());
   InitialSuspend = ActOnFinishFullExpr(InitialSuspend.get());
   if (InitialSuspend.isInvalid())
-    return FD->setInvalidDecl();
+    return StmtError();
 
   // Form and check implicit 'co_await p.final_suspend();' statement.
   ExprResult FinalSuspend =
       buildPromiseCall(*this, Fn, Loc, "final_suspend", None);
   // FIXME: Support operator co_await here.
   if (!FinalSuspend.isInvalid())
     FinalSuspend = BuildCoawaitExpr(Loc, FinalSuspend.get());
   FinalSuspend = ActOnFinishFullExpr(FinalSuspend.get());
   if (FinalSuspend.isInvalid())
-    return FD->setInvalidDecl();
+    return StmtError();
 
   // FIXME: Perform analysis of set_exception call.
 
   // FIXME: Try to form 'p.return_void();' expression statement to handle
   // control flowing off the end of the coroutine.
 
   // Build implicit 'p.get_return_object()' expression and form initialization
   // of return type from it.
   ExprResult ReturnObject =
     buildPromiseCall(*this, Fn, Loc, "get_return_object", None);
   if (ReturnObject.isInvalid())
-    return FD->setInvalidDecl();
+    return StmtError();
   QualType RetType = FD->getReturnType();
   if (!RetType->isDependentType()) {
     InitializedEntity Entity =
         InitializedEntity::InitializeResult(Loc, RetType, false);
     ReturnObject = PerformMoveOrCopyInitialization(Entity, nullptr, RetType,
                                                    ReturnObject.get());
     if (ReturnObject.isInvalid())
-      return FD->setInvalidDecl();
+      return StmtError();
   }
   ReturnObject = ActOnFinishFullExpr(ReturnObject.get(), Loc);
   if (ReturnObject.isInvalid())
-    return FD->setInvalidDecl();
+    return StmtError();
 
   // FIXME: Perform move-initialization of parameters into frame-local copies.
   SmallVector<Expr*, 16> ParamMoves;
 
   // Build body for the coroutine wrapper statement.
-  Body = new (Context) CoroutineBodyStmt(
+  return new (Context) CoroutineBodyStmt(
       Body, PromiseStmt.get(), InitialSuspend.get(), FinalSuspend.get(),
-      /*SetException*/nullptr, /*Fallthrough*/nullptr,
-      ReturnObject.get(), ParamMoves);
+      /*SetException*/ nullptr, /*Fallthrough*/ nullptr, ReturnObject.get(),
+      ParamMoves);
 }

lib/Sema/SemaDecl.cpp

 }
 
 Decl *Sema::ActOnFinishFunctionBody(Decl *D, Stmt *BodyArg) {
   return ActOnFinishFunctionBody(D, BodyArg, false);
 }
 
 Decl *Sema::ActOnFinishFunctionBody(Decl *dcl, Stmt *Body,
                                     bool IsInstantiation) {
-  FunctionDecl *FD = dcl ? dcl->getAsFunction() : nullptr;
+  if (!dcl)
+    return nullptr;
+  FunctionDecl *FD = dcl->getAsFunction();
 
   sema::AnalysisBasedWarnings::Policy WP = AnalysisWarnings.getDefaultPolicy();
   sema::AnalysisBasedWarnings::Policy *ActivePolicy = nullptr;
 
-  if (getLangOpts().CoroutinesTS && !getCurFunction()->CoroutineStmts.empty())
-    CheckCompletedCoroutineBody(FD, Body);
+  if (getLangOpts().CoroutinesTS && !getCurFunction()->CoroutineStmts.empty()) {
+    // FIXME: support ObjC methods here
+    assert(FD && "Objective C methods are not supported");
+    StmtResult NewBody = ActOnFinishCoroutineBody(FD, Body);
+    if (NewBody.isInvalid())
+      FD->setInvalidDecl();
+    else
+      Body = NewBody.get();
+  }
 
   if (FD) {
     FD->setBody(Body);
 
     if (getLangOpts().CPlusPlus14) {
       if (!FD->isInvalidDecl() && Body && !FD->isDependentContext() &&
           FD->getReturnType()->isUndeducedType()) {
         // If the function has a deduced result type but contains no 'return'

lib/Sema/TreeTransform.h

   /// \brief Build a new co_yield expression.
   ///
   /// By default, performs semantic analysis to build the new expression.
   /// Subclasses may override this routine to provide different behavior.
   ExprResult RebuildCoyieldExpr(SourceLocation CoyieldLoc, Expr *Result) {
     return getSema().BuildCoyieldExpr(CoyieldLoc, Result);
   }
 
+  /// \brief Build a new coroutine body.
+  ///
+  /// By default, performs semantic analysis to build the new body.
+  /// Subclasses may override this routine to provide different behavior.
+  StmtResult RebuildCoroutineBodyStmt(Stmt *Body) {
+    auto *FD = dyn_cast<FunctionDecl>(getSema().CurContext);
+    assert(FD); // FIXME this assertion should never fire
+    return getSema().ActOnFinishCoroutineBody(FD, Body);
+  }
+
   /// \brief Build a new Objective-C \@try statement.
   ///
   /// By default, performs semantic analysis to build the new statement.
   /// Subclasses may override this routine to provide different behavior.
   StmtResult RebuildObjCAtTryStmt(SourceLocation AtLoc,
                                         Stmt *TryBody,
                                         MultiStmtArg CatchStmts,
                                         Stmt *Finally) {

                                        TransformedExprs, S->getEndLoc());
 }
 
 // C++ Coroutines TS
 
 template<typename Derived>
 StmtResult
 TreeTransform<Derived>::TransformCoroutineBodyStmt(CoroutineBodyStmt *S) {
-  // The coroutine body should be re-formed by the caller if necessary.
-  return getDerived().TransformStmt(S->getBody());
+  // FIXME: Don't rebuild the entire coroutine body.
+  // The coroutine body should only be re-formed by the caller if necessary.
+  FunctionScopeInfo *FS = getSema().getCurFunction();
+  assert(FS);
+  VarDecl *VD =
+      getSema().buildCoroutinePromise(S->getPromiseDecl()->getLocation());
+  if (!VD || VD->isInvalidDecl())
+    return StmtError();
+  getDerived().transformedLocalDecl(S->getPromiseDecl(), VD);
+  // FIXME: Re-setting FS->CoroutinePromise feels like a hack. Is there a better
+  // way to do this? Currently this is needed so the rebuilt body uses the
+  // transformed promise type.
+  FS->CoroutinePromise = VD;
+
+  StmtResult Body = getDerived().TransformStmt(S->getBody());
+  if (Body.isInvalid())
+    return StmtError();
+
+  return getDerived().RebuildCoroutineBodyStmt(Body.get());
 }
 
 template<typename Derived>
 StmtResult
 TreeTransform<Derived>::TransformCoreturnStmt(CoreturnStmt *S) {
   ExprResult Result = getDerived().TransformInitializer(S->getOperand(),
                                                         /*NotCopyInit*/false);
   if (Result.isInvalid())

test/SemaCXX/coroutines.cpp

 }
 
 template<typename Promise> struct coro {};
 template <typename Promise, typename... Ps>
 struct std::experimental::coroutine_traits<coro<Promise>, Ps...> {
   using promise_type = Promise;
 };
 
-void no_specialization() {
-  co_await a; // expected-error {{implicit instantiation of undefined template 'std::experimental::coroutine_traits<void>'}}
+void no_specialization() { // expected-error {{implicit instantiation of undefined template 'std::experimental::coroutine_traits<void>'}}
+  co_await a;
 }
 
 template <typename... T>
 struct std::experimental::coroutine_traits<int, T...> {};
 
-int no_promise_type() {
-  co_await a; // expected-error {{this function cannot be a coroutine: 'std::experimental::coroutine_traits<int>' has no member named 'promise_type'}}
+int no_promise_type() { // expected-error {{this function cannot be a coroutine: 'std::experimental::coroutine_traits<int>' has no member named 'promise_type'}}
+  co_await a;
 }
 
 template <>
 struct std::experimental::coroutine_traits<double, double> { typedef int promise_type; };
-double bad_promise_type(double) {
-  co_await a; // expected-error {{this function cannot be a coroutine: 'experimental::coroutine_traits<double, double>::promise_type' (aka 'int') is not a class}}
+double bad_promise_type(double) { // expected-error {{this function cannot be a coroutine: 'experimental::coroutine_traits<double, double>::promise_type' (aka 'int') is not a class}}
+  co_await a;
 }
 
 template <>
 struct std::experimental::coroutine_traits<double, int> {
   struct promise_type {};
 };
 double bad_promise_type_2(int) {
   co_yield 0; // expected-error {{no member named 'yield_value' in 'std::experimental::coroutine_traits<double, int>::promise_type'}}
 }
 
-struct promise; // expected-note 2{{forward declaration}}
+struct promise; // expected-note 3{{forward declaration}}
 template <typename... T>
 struct std::experimental::coroutine_traits<void, T...> { using promise_type = promise; };
 
 namespace std {
 namespace experimental {
 template <typename Promise = void>
 struct coroutine_handle;
 }
 }
 
 // FIXME: This diagnostic is terrible.
 void undefined_promise() { // expected-error {{variable has incomplete type 'promise_type'}}
   // FIXME: This diagnostic doesn't make any sense.
   // expected-error@-2 {{incomplete definition of type 'promise'}}
   co_await a;
 }
+template <class T>
+void undefined_promise_template(T) { // expected-error {{variable has incomplete type 'promise_type'}}
+  // FIXME: This diagnostic doesn't make any sense.
+  co_await a;
+}
+template void undefined_promise_template(int); // expected-note {{requested here}}
 
 struct yielded_thing { const char *p; short a, b; };
 
 struct not_awaitable {};
 
 struct promise {
   void get_return_object();
   suspend_always initial_suspend();
struct bad_promise_5 {
   suspend_always initial_suspend();
   not_awaitable final_suspend();
 };
 // FIXME: This diagnostic is terrible.
 coro<bad_promise_5> bad_final_suspend() { // expected-error {{no member named 'await_ready' in 'not_awaitable'}}
   co_await a;
 }
 
+struct not_class_tag {};
+template <>
+struct std::experimental::coroutine_traits<void, not_class_tag> { using promise_type = int; };
+
+template <class T>
+void promise_type_not_class(T) {
+  // expected-error@-1 {{this function cannot be a coroutine: 'experimental::coroutine_traits<void, not_class_tag>::promise_type' (aka 'int') is not a class}}
+  co_await a;
+}
+template void promise_type_not_class(not_class_tag); // expected-note {{requested here}}
 
 template<> struct std::experimental::coroutine_traits<int, int, const char**>
 { using promise_type = promise; };
 
 int main(int, const char**) { // expected-error {{'main' cannot be a coroutine}}
   co_await a; // expected-note {{function is a coroutine due to use of 'co_await' here}}
 }