[coroutines] Fix diagnostics depending on the first coroutine statement.

[EricWF]

Summary

Some coroutine diagnostics need to point to the location of the first coroutine keyword in the function, like when diagnosing a return inside a coroutine. Previously we did this by storing each *valid* coroutine statement in a list and select the first one to use in diagnostics. However if every coroutine statement is invalid we would have no location to point to.

This patch fixes the storage of the first coroutine statement location, ensuring that it gets stored even when the resulting AST node would be invalid.
This patch also removes the CoroutineStmts list in FunctionScopeInfo because it was unused.

[GorNishanov]

LGTM

[GorNishanov]

One thought I had with respect to removing array of statements: We need to check during Sema that await_ready returns a type contextually convertible to bool and await_suspend returns either void or bool. If we have this array, I can quickly run through it and verify that all await expressions have expected types during CheckCoroutine body. With the yield/await expression vector removed, I am not sure where would be an appropriate spot to check for that.

[EricWF]

Good to know. I'll update this tomorrow.

[GorNishanov]

In D30776#697233, @EricWF wrote:

Good to know. I'll update this tomorrow.

Well, that is just the thought. Possibly we can check for the types of await_ready and await_suspend in BuildResolvedCoawaitExpr. Then, we don't have to keep the vector of await/yield-expressions.

[EricWF]

In D30776#697258, @GorNishanov wrote:

In D30776#697233, @EricWF wrote:

Good to know. I'll update this tomorrow.

Well, that is just the thought. Possibly we can check for the types of await_ready and await_suspend in BuildResolvedCoawaitExpr. Then, we don't have to keep the vector of await/yield-expressions.

I guess it depends on when we want the diagnostics to be issued. And I think emitting them as they are built makes the most sense. If we attempt to issue the diagnostics at the end of the body it may be too late.

Do you still want me to put the vector back in now or when the need actually arises?

[GorNishanov]

Alright, then. LGTM once more!

[EricWF]

Merge with master.

include/clang/Sema/ScopeInfo.h

 #include "clang/AST/Type.h"
 #include "clang/Basic/CapturedStmt.h"
 #include "clang/Basic/PartialDiagnostic.h"
 #include "clang/Sema/CleanupInfo.h"
 #include "clang/Sema/Ownership.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringSwitch.h"
 #include <algorithm>
 
 namespace clang {
 
 class Decl;
 class BlockDecl;
 class CapturedDecl;
 class CXXMethodDecl;
public:
   /// This starts true for a secondary initializer method and will be set to
   /// false if there is an invocation of an initializer on 'self'.
   bool ObjCWarnForNoInitDelegation : 1;
 
   /// \brief True only when this function has not already built, or attempted
   /// to build, the initial and final coroutine suspend points
   bool NeedsCoroutineSuspends : 1;
 
+  /// \brief An enumeration represeting the kind of the first coroutine statement
+  /// in the function. One of co_return, co_await, or co_yield.
+  unsigned char FirstCoroutineStmtKind : 2;
+
+  /// First coroutine statement in the current function.
+  /// (ex co_return, co_await, co_yield)
+  SourceLocation FirstCoroutineStmtLoc;
+
   /// First 'return' statement in the current function.
   SourceLocation FirstReturnLoc;
 
   /// First C++ 'try' statement in the current function.
   SourceLocation FirstCXXTryLoc;
 
   /// First SEH '__try' statement in the current function.
   SourceLocation FirstSEHTryLoc;
public:
   SmallVector<ReturnStmt*, 4> Returns;
 
   /// \brief The promise object for this coroutine, if any.
   VarDecl *CoroutinePromise = nullptr;
 
   /// \brief The initial and final coroutine suspend points.
   std::pair<Stmt *, Stmt *> CoroutineSuspends;
 
-  /// \brief The list of coroutine control flow constructs (co_await, co_yield,
-  /// co_return) that occur within the function or block. Empty if and only if
-  /// this function or block is not (yet known to be) a coroutine.
-  SmallVector<Stmt*, 4> CoroutineStmts;
-
   /// \brief The stack of currently active compound stamement scopes in the
   /// function.
   SmallVector<CompoundScopeInfo, 4> CompoundScopes;
 
   /// \brief A list of PartialDiagnostics created but delayed within the
   /// current function scope.  These diagnostics are vetted for reachability
   /// prior to being emitted.
   SmallVector<PossiblyUnreachableDiag, 4> PossiblyUnreachableDiags;
public:
   }
 
   bool NeedsScopeChecking() const {
     return !HasDroppedStmt &&
         (HasIndirectGoto ||
           (HasBranchProtectedScope && HasBranchIntoScope));
   }
 
+  void setFirstCoroutineStmt(SourceLocation Loc, StringRef Keyword) {
+    assert(FirstCoroutineStmtLoc.isInvalid() &&
+                   "first coroutine statement location already set");
+    FirstCoroutineStmtLoc = Loc;
+    FirstCoroutineStmtKind = llvm::StringSwitch<unsigned char>(Keyword)
+            .Case("co_return", 0)
+            .Case("co_await", 1)
+            .Case("co_yield", 2);
+  }
+
+  StringRef getFirstCoroutineStmtKeyword() const {
+    assert(FirstCoroutineStmtLoc.isValid()
+                   && "no coroutine statement available");
+    switch (FirstCoroutineStmtKind) {
+    case 0: return "co_return";
+    case 1: return "co_await";
+    case 2: return "co_yield";
+    default:
+      llvm_unreachable("FirstCoroutineStmtKind has an invalid value");
+    };
+  }
+
   void setNeedsCoroutineSuspends(bool value = true) {
     assert((!value || CoroutineSuspends.first == nullptr) &&
             "we already have valid suspend points");
     NeedsCoroutineSuspends = value;
   }
 
   bool hasInvalidCoroutineSuspends() const {
     return !NeedsCoroutineSuspends && CoroutineSuspends.first == nullptr;

lib/Sema/ScopeInfo.cpp

void FunctionScopeInfo::Clear() {
   ObjCIsDesignatedInit = false;
   ObjCWarnForNoDesignatedInitChain = false;
   ObjCIsSecondaryInit = false;
   ObjCWarnForNoInitDelegation = false;
   FirstReturnLoc = SourceLocation();
   FirstCXXTryLoc = SourceLocation();
   FirstSEHTryLoc = SourceLocation();
 
-  SwitchStack.clear();
-  Returns.clear();
+  // Coroutine state
+  FirstCoroutineStmtLoc = SourceLocation();
   CoroutinePromise = nullptr;
   NeedsCoroutineSuspends = true;
   CoroutineSuspends.first = nullptr;
   CoroutineSuspends.second = nullptr;
-  CoroutineStmts.clear();
+
+  SwitchStack.clear();
+  Returns.clear();
   ErrorTrap.reset();
   PossiblyUnreachableDiags.clear();
   WeakObjectUses.clear();
   ModifiedNonNullParams.clear();
 }
 
 static const NamedDecl *getBestPropertyDecl(const ObjCPropertyRefExpr *PropE) {
   if (PropE->isExplicitProperty())

lib/Sema/SemaCoroutine.cpp

if (VD->isInvalidDecl())
     return nullptr;
   ActOnUninitializedDecl(VD);
   assert(!VD->isInvalidDecl());
   return VD;
 }
 
 /// Check that this is a context in which a coroutine suspension can appear.
 static FunctionScopeInfo *checkCoroutineContext(Sema &S, SourceLocation Loc,
-                                                StringRef Keyword) {
+                                                StringRef Keyword,
+                                                bool IsImplicit = false) {
   if (!isValidCoroutineContext(S, Loc, Keyword))
     return nullptr;
 
   assert(isa<FunctionDecl>(S.CurContext) && "not in a function scope");
 
   auto *ScopeInfo = S.getCurFunction();
   assert(ScopeInfo && "missing function scope for function");
 
+  if (ScopeInfo->FirstCoroutineStmtLoc.isInvalid() && !IsImplicit)
+    ScopeInfo->setFirstCoroutineStmt(Loc, Keyword);
+
   if (ScopeInfo->CoroutinePromise)
     return ScopeInfo;
 
   ScopeInfo->CoroutinePromise = S.buildCoroutinePromise(Loc);
   if (!ScopeInfo->CoroutinePromise)
     return nullptr;
 
   return ScopeInfo;
ExprResult Sema::ActOnCoawaitExpr(Scope *S, SourceLocation Loc, Expr *E) {
   ExprResult Lookup = buildOperatorCoawaitLookupExpr(*this, S, Loc);
   if (Lookup.isInvalid())
     return ExprError();
   return BuildUnresolvedCoawaitExpr(Loc, E,
                                    cast<UnresolvedLookupExpr>(Lookup.get()));
 }
 
 ExprResult Sema::BuildUnresolvedCoawaitExpr(SourceLocation Loc, Expr *E,
                                             UnresolvedLookupExpr *Lookup) {
   auto *FSI = checkCoroutineContext(*this, Loc, "co_await");
   if (!FSI)
     return ExprError();
 
   if (E->getType()->isPlaceholderType()) {
     ExprResult R = CheckPlaceholderExpr(E);
     if (R.isInvalid())
       return ExprError();
     E = R.get();
   }
 
   auto *Promise = FSI->CoroutinePromise;
   if (Promise->getType()->isDependentType()) {
     Expr *Res =
         new (Context) DependentCoawaitExpr(Loc, Context.DependentTy, E, Lookup);
-    FSI->CoroutineStmts.push_back(Res);
     return Res;
   }
 
   auto *RD = Promise->getType()->getAsCXXRecordDecl();
   if (lookupMember(*this, "await_transform", RD, Loc)) {
     ExprResult R = buildPromiseCall(*this, Promise, Loc, "await_transform", E);
     if (R.isInvalid()) {
       Diag(Loc,
            diag::note_coroutine_promise_implicit_await_transform_required_here)
           << E->getSourceRange();
       return ExprError();
     }
     E = R.get();
   }
   ExprResult Awaitable = buildOperatorCoawaitCall(*this, Loc, E, Lookup);
   if (Awaitable.isInvalid())
     return ExprError();
 
   return BuildResolvedCoawaitExpr(Loc, Awaitable.get());
 }
 
 ExprResult Sema::BuildResolvedCoawaitExpr(SourceLocation Loc, Expr *E,
                                   bool IsImplicit) {
-  auto *Coroutine = checkCoroutineContext(*this, Loc, "co_await");
+  auto *Coroutine = checkCoroutineContext(*this, Loc, "co_await", IsImplicit);
   if (!Coroutine)
     return ExprError();
 
   if (E->getType()->isPlaceholderType()) {
     ExprResult R = CheckPlaceholderExpr(E);
     if (R.isInvalid()) return ExprError();
     E = R.get();
   }
 
   if (E->getType()->isDependentType()) {
     Expr *Res = new (Context)
         CoawaitExpr(Loc, Context.DependentTy, E, IsImplicit);
-    if (!IsImplicit)
-      Coroutine->CoroutineStmts.push_back(Res);
     return Res;
   }
 
   // If the expression is a temporary, materialize it as an lvalue so that we
   // can use it multiple times.
   if (E->getValueKind() == VK_RValue)
     E = CreateMaterializeTemporaryExpr(E->getType(), E, true);
 
   // Build the await_ready, await_suspend, await_resume calls.
   ReadySuspendResumeResult RSS =
       buildCoawaitCalls(*this, Coroutine->CoroutinePromise, Loc, E);
   if (RSS.IsInvalid)
     return ExprError();
 
   Expr *Res =
       new (Context) CoawaitExpr(Loc, E, RSS.Results[0], RSS.Results[1],
                                 RSS.Results[2], RSS.OpaqueValue, IsImplicit);
-  if (!IsImplicit)
-    Coroutine->CoroutineStmts.push_back(Res);
+
   return Res;
 }
 
 ExprResult Sema::ActOnCoyieldExpr(Scope *S, SourceLocation Loc, Expr *E) {
   if (!actOnCoroutineBodyStart(*this, S, Loc, "co_yield")) {
     CorrectDelayedTyposInExpr(E);
     return ExprError();
   }
ExprResult Sema::BuildCoyieldExpr(SourceLocation Loc, Expr *E) {
   if (E->getType()->isPlaceholderType()) {
     ExprResult R = CheckPlaceholderExpr(E);
     if (R.isInvalid()) return ExprError();
     E = R.get();
   }
 
   if (E->getType()->isDependentType()) {
     Expr *Res = new (Context) CoyieldExpr(Loc, Context.DependentTy, E);
-    Coroutine->CoroutineStmts.push_back(Res);
     return Res;
   }
 
   // If the expression is a temporary, materialize it as an lvalue so that we
   // can use it multiple times.
   if (E->getValueKind() == VK_RValue)
     E = CreateMaterializeTemporaryExpr(E->getType(), E, true);
 
   // Build the await_ready, await_suspend, await_resume calls.
   ReadySuspendResumeResult RSS =
       buildCoawaitCalls(*this, Coroutine->CoroutinePromise, Loc, E);
   if (RSS.IsInvalid)
     return ExprError();
 
   Expr *Res = new (Context) CoyieldExpr(Loc, E, RSS.Results[0], RSS.Results[1],
                                         RSS.Results[2], RSS.OpaqueValue);
-  Coroutine->CoroutineStmts.push_back(Res);
+
   return Res;
 }
 
 StmtResult Sema::ActOnCoreturnStmt(Scope *S, SourceLocation Loc, Expr *E) {
   if (!actOnCoroutineBodyStart(*this, S, Loc, "co_return")) {
     CorrectDelayedTyposInExpr(E);
     return StmtError();
   }
   return BuildCoreturnStmt(Loc, E);
 }
 
 StmtResult Sema::BuildCoreturnStmt(SourceLocation Loc, Expr *E,
                                    bool IsImplicit) {
-  auto *FSI = checkCoroutineContext(*this, Loc, "co_return");
+  auto *FSI = checkCoroutineContext(*this, Loc, "co_return", IsImplicit);
   if (!FSI)
     return StmtError();
 
   if (E && E->getType()->isPlaceholderType() &&
       !E->getType()->isSpecificPlaceholderType(BuiltinType::Overload)) {
     ExprResult R = CheckPlaceholderExpr(E);
     if (R.isInvalid()) return StmtError();
     E = R.get();
if (E && (isa<InitListExpr>(E) || !E->getType()->isVoidType())) {
     PC = buildPromiseCall(*this, Promise, Loc, "return_void", None);
   }
   if (PC.isInvalid())
     return StmtError();
 
   Expr *PCE = ActOnFinishFullExpr(PC.get()).get();
 
   Stmt *Res = new (Context) CoreturnStmt(Loc, E, PCE, IsImplicit);
-  if (!IsImplicit)
-    FSI->CoroutineStmts.push_back(Res);
   return Res;
 }
 
 static ExprResult buildStdCurrentExceptionCall(Sema &S, SourceLocation Loc) {
   NamespaceDecl *Std = S.getStdNamespace();
   if (!Std) {
     S.Diag(Loc, diag::err_implied_std_current_exception_not_found);
     return ExprError();
if (isa<CoroutineBodyStmt>(Body)) {
     // FIXME(EricWF): Nothing todo. the body is already a transformed coroutine
     // body statement.
     return;
   }
 
   // Coroutines [stmt.return]p1:
   //   A return statement shall not appear in a coroutine.
   if (Fn->FirstReturnLoc.isValid()) {
+    assert(Fn->FirstCoroutineStmtLoc.isValid() &&
+                   "first coroutine location not set");
     Diag(Fn->FirstReturnLoc, diag::err_return_in_coroutine);
-    // FIXME: Every Coroutine statement may be invalid and therefore not added
-    // to CoroutineStmts. Find another way to provide location information.
-    if (!Fn->CoroutineStmts.empty()) {
-      auto *First = Fn->CoroutineStmts[0];
-      Diag(First->getLocStart(), diag::note_declared_coroutine_here)
-          << (isa<CoawaitExpr>(First) ? "co_await" :
-              isa<CoyieldExpr>(First) ? "co_yield" : "co_return");
-    }
+    Diag(Fn->FirstCoroutineStmtLoc, diag::note_declared_coroutine_here)
+            << Fn->getFirstCoroutineStmtKeyword();
   }
   SubStmtBuilder Builder(*this, *FD, *Fn, Body);
   if (Builder.isInvalid())
     return FD->setInvalidDecl();
 
   // Build body for the coroutine wrapper statement.
   Body = CoroutineBodyStmt::Create(Context, Builder);
 }

lib/Sema/TreeTransform.h

   CoroutineBodyStmt::CtorArgs BodyArgs;
 
   auto *ScopeInfo = SemaRef.getCurFunction();
   auto *FD = cast<FunctionDecl>(SemaRef.CurContext);
   assert(ScopeInfo && !ScopeInfo->CoroutinePromise &&
          ScopeInfo->NeedsCoroutineSuspends &&
          ScopeInfo->CoroutineSuspends.first == nullptr &&
          ScopeInfo->CoroutineSuspends.second == nullptr &&
-         ScopeInfo->CoroutineStmts.empty() && "expected clean scope info");
+         "expected clean scope info");
 
   // Set that we have (possibly-invalid) suspend points before we do anything
   // that may fail.
   ScopeInfo->setNeedsCoroutineSuspends(false);
 
   // The new CoroutinePromise object needs to be built and put into the current
   // FunctionScopeInfo before any transformations or rebuilding occurs.
   auto *Promise = S->getPromiseDecl();

test/SemaCXX/coroutines.cpp

 }
 template void co_await_non_dependent_arg(int);
 
 void mixed_yield() {
   co_yield 0; // expected-note {{use of 'co_yield'}}
   return; // expected-error {{not allowed in coroutine}}
 }
 
+void mixed_yield_invalid() {
+  co_yield blah; // expected-error {{use of undeclared identifier}}
+  // expected-note@-1 {{function is a coroutine due to use of 'co_yield'}}
+  return; // expected-error {{return statement not allowed in coroutine}}
+}
+
+template <class T>
+void mixed_yield_template(T) {
+  co_yield blah; // expected-error {{use of undeclared identifier}}
+  // expected-note@-1 {{function is a coroutine due to use of 'co_yield'}}
+  return; // expected-error {{return statement not allowed in coroutine}}
+}
+
+template <class T>
+void mixed_yield_template2(T) {
+  co_yield 42;
+  // expected-note@-1 {{function is a coroutine due to use of 'co_yield'}}
+  return; // expected-error {{return statement not allowed in coroutine}}
+}
+
+template <class T>
+void mixed_yield_template3(T v) {
+  co_yield blah(v);
+  // expected-note@-1 {{function is a coroutine due to use of 'co_yield'}}
+  return; // expected-error {{return statement not allowed in coroutine}}
+}
+
 void mixed_await() {
   co_await a; // expected-note {{use of 'co_await'}}
   return; // expected-error {{not allowed in coroutine}}
 }
 
+void mixed_await_invalid() {
+  co_await 42; // expected-error {{'int' is not a structure or union}}
+  // expected-note@-1 {{function is a coroutine due to use of 'co_await'}}
+  return; // expected-error {{not allowed in coroutine}}
+}
+
+template <class T>
+void mixed_await_template(T) {
+  co_await 42;
+  // expected-note@-1 {{function is a coroutine due to use of 'co_await'}}
+  return; // expected-error {{not allowed in coroutine}}
+}
+
+template <class T>
+void mixed_await_template2(T v) {
+  co_await v; // expected-error {{'long' is not a structure or union}}
+  // expected-note@-1 {{function is a coroutine due to use of 'co_await'}}
+  return; // expected-error {{not allowed in coroutine}}
+}
+template void mixed_await_template2(long); // expected-note {{requested here}}
+
 void only_coreturn(void_tag) {
   co_return; // OK
 }
 
 void mixed_coreturn(void_tag, bool b) {
   if (b)
     co_return; // expected-note {{use of 'co_return'}}
   else
     return; // expected-error {{not allowed in coroutine}}
 }
 
+void mixed_coreturn_invalid(bool b) {
+  if (b)
+    co_return; // expected-note {{use of 'co_return'}}
+    // expected-error@-1 {{no member named 'return_void' in 'promise'}}
+  else
+    return; // expected-error {{not allowed in coroutine}}
+}
+
+template <class T>
+void mixed_coreturn_template(void_tag, bool b, T v) {
+  if (b)
+    co_return v; // expected-note {{use of 'co_return'}}
+    // expected-error@-1 {{no member named 'return_value' in 'promise_void'}}
+  else
+    return; // expected-error {{not allowed in coroutine}}
+}
+template void mixed_coreturn_template(void_tag, bool, int); // expected-note {{requested here}}
+
+template <class T>
+void mixed_coreturn_template2(bool b, T) {
+  if (b)
+    co_return v; // expected-note {{use of 'co_return'}}
+    // expected-error@-1 {{use of undeclared identifier 'v'}}
+  else
+    return; // expected-error {{not allowed in coroutine}}
+}
+
 struct CtorDtor {
   CtorDtor() {
     co_yield 0; // expected-error {{'co_yield' cannot be used in a constructor}}
   }
   CtorDtor(awaitable a) {
     // The spec doesn't say this is ill-formed, but it must be.
     co_await a; // expected-error {{'co_await' cannot be used in a constructor}}
   }