[coroutines] Fix co_await for range statement

Summary:
Currently we build the co_await expressions on the wrong implicit statements of the implicit ranged for; Specifically we build the co_await expression wrapping the range declaration, but it should wrap the begin expression.

This patch fixes co_await on range for.

Reviewers: rsmith, GorNishanov

Reviewed By: GorNishanov

Subscribers: cfe-commits

Differential Revision: https://reviews.llvm.org/D34021

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@305363 91177308-0d34-0410-b5e6-96231b3b80d8

4 files changed, 217 insertions, 37 deletions

diff --git a/include/clang/Sema/Sema.h b/include/clang/Sema/Sema.h
index 8025668e66..1dedf4ba24 100644
--- a/include/clang/Sema/Sema.h
+++ b/include/clang/Sema/Sema.h
@@ -8364,6 +8364,8 @@ public:
   //===--------------------------------------------------------------------===//
   // C++ Coroutines TS
   //
+  bool ActOnCoroutineBodyStart(Scope *S, SourceLocation KwLoc,
+                               StringRef Keyword);
   ExprResult ActOnCoawaitExpr(Scope *S, SourceLocation KwLoc, Expr *E);
   ExprResult ActOnCoyieldExpr(Scope *S, SourceLocation KwLoc, Expr *E);
   StmtResult ActOnCoreturnStmt(Scope *S, SourceLocation KwLoc, Expr *E);
diff --git a/lib/Sema/SemaCoroutine.cpp b/lib/Sema/SemaCoroutine.cpp
index 06ae66076e..b05c0998d3 100644
--- a/lib/Sema/SemaCoroutine.cpp
+++ b/lib/Sema/SemaCoroutine.cpp
@@ -470,11 +470,11 @@ static FunctionScopeInfo *checkCoroutineContext(Sema &S, SourceLocation Loc,
   return ScopeInfo;
 }
 
-static bool actOnCoroutineBodyStart(Sema &S, Scope *SC, SourceLocation KWLoc,
-                                    StringRef Keyword) {
-  if (!checkCoroutineContext(S, KWLoc, Keyword))
+bool Sema::ActOnCoroutineBodyStart(Scope *SC, SourceLocation KWLoc,
+                                   StringRef Keyword) {
+  if (!checkCoroutineContext(*this, KWLoc, Keyword))
     return false;
-  auto *ScopeInfo = S.getCurFunction();
+  auto *ScopeInfo = getCurFunction();
   assert(ScopeInfo->CoroutinePromise);
 
   // If we have existing coroutine statements then we have already built
@@ -484,24 +484,24 @@ static bool actOnCoroutineBodyStart(Sema &S, Scope *SC, SourceLocation KWLoc,
 
   ScopeInfo->setNeedsCoroutineSuspends(false);
 
-  auto *Fn = cast<FunctionDecl>(S.CurContext);
+  auto *Fn = cast<FunctionDecl>(CurContext);
   SourceLocation Loc = Fn->getLocation();
   // Build the initial suspend point
   auto buildSuspends = [&](StringRef Name) mutable -> StmtResult {
     ExprResult Suspend =
-        buildPromiseCall(S, ScopeInfo->CoroutinePromise, Loc, Name, None);
+        buildPromiseCall(*this, ScopeInfo->CoroutinePromise, Loc, Name, None);
     if (Suspend.isInvalid())
       return StmtError();
-    Suspend = buildOperatorCoawaitCall(S, SC, Loc, Suspend.get());
+    Suspend = buildOperatorCoawaitCall(*this, SC, Loc, Suspend.get());
     if (Suspend.isInvalid())
       return StmtError();
-    Suspend = S.BuildResolvedCoawaitExpr(Loc, Suspend.get(),
-                                         /*IsImplicit*/ true);
-    Suspend = S.ActOnFinishFullExpr(Suspend.get());
+    Suspend = BuildResolvedCoawaitExpr(Loc, Suspend.get(),
+                                       /*IsImplicit*/ true);
+    Suspend = ActOnFinishFullExpr(Suspend.get());
     if (Suspend.isInvalid()) {
-      S.Diag(Loc, diag::note_coroutine_promise_suspend_implicitly_required)
+      Diag(Loc, diag::note_coroutine_promise_suspend_implicitly_required)
           << ((Name == "initial_suspend") ? 0 : 1);
-      S.Diag(KWLoc, diag::note_declared_coroutine_here) << Keyword;
+      Diag(KWLoc, diag::note_declared_coroutine_here) << Keyword;
       return StmtError();
     }
     return cast<Stmt>(Suspend.get());
@@ -521,7 +521,7 @@ static bool actOnCoroutineBodyStart(Sema &S, Scope *SC, SourceLocation KWLoc,
 }
 
 ExprResult Sema::ActOnCoawaitExpr(Scope *S, SourceLocation Loc, Expr *E) {
-  if (!actOnCoroutineBodyStart(*this, S, Loc, "co_await")) {
+  if (!ActOnCoroutineBodyStart(S, Loc, "co_await")) {
     CorrectDelayedTyposInExpr(E);
     return ExprError();
   }
@@ -613,7 +613,7 @@ ExprResult Sema::BuildResolvedCoawaitExpr(SourceLocation Loc, Expr *E,
 }
 
 ExprResult Sema::ActOnCoyieldExpr(Scope *S, SourceLocation Loc, Expr *E) {
-  if (!actOnCoroutineBodyStart(*this, S, Loc, "co_yield")) {
+  if (!ActOnCoroutineBodyStart(S, Loc, "co_yield")) {
     CorrectDelayedTyposInExpr(E);
     return ExprError();
   }
@@ -658,14 +658,15 @@ ExprResult Sema::BuildCoyieldExpr(SourceLocation Loc, Expr *E) {
   if (RSS.IsInvalid)
     return ExprError();
 
-  Expr *Res = new (Context) CoyieldExpr(Loc, E, RSS.Results[0], RSS.Results[1],
-                                        RSS.Results[2], RSS.OpaqueValue);
+  Expr *Res =
+      new (Context) CoyieldExpr(Loc, E, RSS.Results[0], RSS.Results[1],
+                                RSS.Results[2], RSS.OpaqueValue);
 
   return Res;
 }
 
 StmtResult Sema::ActOnCoreturnStmt(Scope *S, SourceLocation Loc, Expr *E) {
-  if (!actOnCoroutineBodyStart(*this, S, Loc, "co_return")) {
+  if (!ActOnCoroutineBodyStart(S, Loc, "co_return")) {
     CorrectDelayedTyposInExpr(E);
     return StmtError();
   }
diff --git a/lib/Sema/SemaStmt.cpp b/lib/Sema/SemaStmt.cpp
index 0c35f9a571..eed10b077e 100644
--- a/lib/Sema/SemaStmt.cpp
+++ b/lib/Sema/SemaStmt.cpp
@@ -1989,11 +1989,11 @@ StmtResult Sema::ActOnCXXForRangeStmt(Scope *S, SourceLocation ForLoc,
     return StmtError();
   }
 
-  // Coroutines: 'for co_await' implicitly co_awaits its range.
-  if (CoawaitLoc.isValid()) {
-    ExprResult Coawait = ActOnCoawaitExpr(S, CoawaitLoc, Range);
-    if (Coawait.isInvalid()) return StmtError();
-    Range = Coawait.get();
+  // Build the coroutine state immediately and not later during template
+  // instantiation
+  if (!CoawaitLoc.isInvalid()) {
+    if (!ActOnCoroutineBodyStart(S, CoawaitLoc, "co_await"))
+      return StmtError();
   }
 
   // Build  auto && __range = range-init
@@ -2031,16 +2031,12 @@ StmtResult Sema::ActOnCXXForRangeStmt(Scope *S, SourceLocation ForLoc,
 /// BeginExpr and EndExpr are set and FRS_Success is returned on success;
 /// CandidateSet and BEF are set and some non-success value is returned on
 /// failure.
-static Sema::ForRangeStatus BuildNonArrayForRange(Sema &SemaRef,
-                                            Expr *BeginRange, Expr *EndRange,
-                                            QualType RangeType,
-                                            VarDecl *BeginVar,
-                                            VarDecl *EndVar,
-                                            SourceLocation ColonLoc,
-                                            OverloadCandidateSet *CandidateSet,
-                                            ExprResult *BeginExpr,
-                                            ExprResult *EndExpr,
-                                            BeginEndFunction *BEF) {
+static Sema::ForRangeStatus
+BuildNonArrayForRange(Sema &SemaRef, Expr *BeginRange, Expr *EndRange,
+                      QualType RangeType, VarDecl *BeginVar, VarDecl *EndVar,
+                      SourceLocation ColonLoc, SourceLocation CoawaitLoc,
+                      OverloadCandidateSet *CandidateSet, ExprResult *BeginExpr,
+                      ExprResult *EndExpr, BeginEndFunction *BEF) {
   DeclarationNameInfo BeginNameInfo(
       &SemaRef.PP.getIdentifierTable().get("begin"), ColonLoc);
   DeclarationNameInfo EndNameInfo(&SemaRef.PP.getIdentifierTable().get("end"),
@@ -2087,6 +2083,15 @@ static Sema::ForRangeStatus BuildNonArrayForRange(Sema &SemaRef,
           << ColonLoc << BEF_begin << BeginRange->getType();
     return RangeStatus;
   }
+  if (!CoawaitLoc.isInvalid()) {
+    // FIXME: getCurScope() should not be used during template instantiation.
+    // We should pick up the set of unqualified lookup results for operator
+    // co_await during the initial parse.
+    *BeginExpr = SemaRef.ActOnCoawaitExpr(SemaRef.getCurScope(), ColonLoc,
+                                          BeginExpr->get());
+    if (BeginExpr->isInvalid())
+      return Sema::FRS_DiagnosticIssued;
+  }
   if (FinishForRangeVarDecl(SemaRef, BeginVar, BeginExpr->get(), ColonLoc,
                             diag::err_for_range_iter_deduction_failure)) {
     NoteForRangeBeginEndFunction(SemaRef, BeginExpr->get(), *BEF);
@@ -2253,6 +2258,11 @@ Sema::BuildCXXForRangeStmt(SourceLocation ForLoc, SourceLocation CoawaitLoc,
 
       // begin-expr is __range.
       BeginExpr = BeginRangeRef;
+      if (!CoawaitLoc.isInvalid()) {
+        BeginExpr = ActOnCoawaitExpr(S, ColonLoc, BeginExpr.get());
+        if (BeginExpr.isInvalid())
+          return StmtError();
+      }
       if (FinishForRangeVarDecl(*this, BeginVar, BeginRangeRef.get(), ColonLoc,
                                 diag::err_for_range_iter_deduction_failure)) {
         NoteForRangeBeginEndFunction(*this, BeginExpr.get(), BEF_begin);
@@ -2335,11 +2345,10 @@ Sema::BuildCXXForRangeStmt(SourceLocation ForLoc, SourceLocation CoawaitLoc,
       OverloadCandidateSet CandidateSet(RangeLoc,
                                         OverloadCandidateSet::CSK_Normal);
       BeginEndFunction BEFFailure;
-      ForRangeStatus RangeStatus =
-          BuildNonArrayForRange(*this, BeginRangeRef.get(),
-                                EndRangeRef.get(), RangeType,
-                                BeginVar, EndVar, ColonLoc, &CandidateSet,
-                                &BeginExpr, &EndExpr, &BEFFailure);
+      ForRangeStatus RangeStatus = BuildNonArrayForRange(
+          *this, BeginRangeRef.get(), EndRangeRef.get(), RangeType, BeginVar,
+          EndVar, ColonLoc, CoawaitLoc, &CandidateSet, &BeginExpr, &EndExpr,
+          &BEFFailure);
 
       if (Kind == BFRK_Build && RangeStatus == FRS_NoViableFunction &&
           BEFFailure == BEF_begin) {
@@ -2436,6 +2445,9 @@ Sema::BuildCXXForRangeStmt(SourceLocation ForLoc, SourceLocation CoawaitLoc,
 
     IncrExpr = ActOnUnaryOp(S, ColonLoc, tok::plusplus, BeginRef.get());
     if (!IncrExpr.isInvalid() && CoawaitLoc.isValid())
+      // FIXME: getCurScope() should not be used during template instantiation.
+      // We should pick up the set of unqualified lookup results for operator
+      // co_await during the initial parse.
       IncrExpr = ActOnCoawaitExpr(S, CoawaitLoc, IncrExpr.get());
     if (!IncrExpr.isInvalid())
       IncrExpr = ActOnFinishFullExpr(IncrExpr.get());
diff --git a/test/SemaCXX/co_await-range-for.cpp b/test/SemaCXX/co_await-range-for.cpp
new file mode 100644
index 0000000000..4d999ea7db
--- /dev/null
+++ b/test/SemaCXX/co_await-range-for.cpp
@@ -0,0 +1,165 @@
+// RUN: %clang_cc1 -triple x86_64-apple-darwin9 %s -std=c++14 -fcoroutines-ts \
+// RUN:    -fsyntax-only -Wignored-qualifiers -Wno-error=return-type -verify \
+// RUN:    -fblocks
+#include "Inputs/std-coroutine.h"
+
+using namespace std::experimental;
+
+
+template <class Begin>
+struct Awaiter {
+  bool await_ready();
+  void await_suspend(coroutine_handle<>);
+  Begin await_resume();
+};
+
+template <class Iter> struct BeginTag { BeginTag() = delete; };
+template <class Iter> struct IncTag { IncTag() = delete; };
+
+template <class Iter, bool Delete = false>
+struct CoawaitTag { CoawaitTag() = delete; };
+
+template <class T>
+struct Iter {
+  using value_type = T;
+  using reference = T &;
+  using pointer = T *;
+
+  IncTag<Iter> operator++();
+  reference operator*();
+  pointer operator->();
+};
+template <class T> bool operator==(Iter<T>, Iter<T>);
+template <class T> bool operator!=(Iter<T>, Iter<T>);
+
+template <class T>
+struct Range {
+  BeginTag<Iter<T>> begin();
+  Iter<T> end();
+};
+
+struct MyForLoopArrayAwaiter {
+  struct promise_type {
+    MyForLoopArrayAwaiter get_return_object() { return {}; }
+    void return_void();
+    void unhandled_exception();
+    suspend_never initial_suspend();
+    suspend_never final_suspend();
+    template <class T>
+    Awaiter<T *> await_transform(T *) = delete; // expected-note {{explicitly deleted}}
+  };
+};
+MyForLoopArrayAwaiter g() {
+  int arr[10] = {0};
+  for co_await(auto i : arr) {}
+  // expected-error@-1 {{call to deleted member function 'await_transform'}}
+  // expected-note@-2 {{'await_transform' implicitly required by 'co_await' here}}
+}
+
+struct ForLoopAwaiterBadBeginTransform {
+  struct promise_type {
+    ForLoopAwaiterBadBeginTransform get_return_object();
+    void return_void();
+    void unhandled_exception();
+    suspend_never initial_suspend();
+    suspend_never final_suspend();
+
+    template <class T>
+    Awaiter<T> await_transform(BeginTag<T>) = delete; // expected-note 1+ {{explicitly deleted}}
+
+    template <class T>
+    CoawaitTag<T> await_transform(IncTag<T>); // expected-note 1+ {{candidate}}
+  };
+};
+ForLoopAwaiterBadBeginTransform bad_begin() {
+  Range<int> R;
+  for co_await(auto i : R) {}
+  // expected-error@-1 {{call to deleted member function 'await_transform'}}
+  // expected-note@-2 {{'await_transform' implicitly required by 'co_await' here}}
+}
+template <class Dummy>
+ForLoopAwaiterBadBeginTransform bad_begin_template(Dummy) {
+  Range<Dummy> R;
+  for co_await(auto i : R) {}
+  // expected-error@-1 {{call to deleted member function 'await_transform'}}
+  // expected-note@-2 {{'await_transform' implicitly required by 'co_await' here}}
+}
+template ForLoopAwaiterBadBeginTransform bad_begin_template(int); // expected-note {{requested here}}
+
+template <class Iter>
+Awaiter<Iter> operator co_await(CoawaitTag<Iter, true>) = delete;
+// expected-note@-1 1+ {{explicitly deleted}}
+
+struct ForLoopAwaiterBadIncTransform {
+  struct promise_type {
+    ForLoopAwaiterBadIncTransform get_return_object();
+    void return_void();
+    void unhandled_exception();
+    suspend_never initial_suspend();
+    suspend_never final_suspend();
+
+    template <class T>
+    Awaiter<T> await_transform(BeginTag<T> e);
+
+    template <class T>
+    CoawaitTag<T, true> await_transform(IncTag<T>);
+  };
+};
+ForLoopAwaiterBadIncTransform bad_inc_transform() {
+  Range<float> R;
+  for co_await(auto i : R) {}
+  // expected-error@-1 {{overload resolution selected deleted operator 'co_await'}}
+  // expected-note@-2 {{in implicit call to 'operator++' for iterator of type 'Range<float>'}}
+}
+
+template <class Dummy>
+ForLoopAwaiterBadIncTransform bad_inc_transform_template(Dummy) {
+  Range<Dummy> R;
+  for co_await(auto i : R) {}
+  // expected-error@-1 {{overload resolution selected deleted operator 'co_await'}}
+  // expected-note@-2 {{in implicit call to 'operator++' for iterator of type 'Range<long>'}}
+}
+template ForLoopAwaiterBadIncTransform bad_inc_transform_template(long); // expected-note {{requested here}}
+
+// Ensure we mark and check the function as a coroutine even if it's
+// never instantiated.
+template <class T>
+constexpr void never_instant(T) {
+  static_assert(sizeof(T) != sizeof(T), "function should not be instantiated");
+  for co_await(auto i : foo(T{})) {}
+  // expected-error@-1 {{'co_await' cannot be used in a constexpr function}}
+}
+
+namespace NS {
+struct ForLoopAwaiterCoawaitLookup {
+  struct promise_type {
+    ForLoopAwaiterCoawaitLookup get_return_object();
+    void return_void();
+    void unhandled_exception();
+    suspend_never initial_suspend();
+    suspend_never final_suspend();
+    template <class T>
+    CoawaitTag<T, false> await_transform(BeginTag<T> e);
+    template <class T>
+    Awaiter<T> await_transform(IncTag<T>);
+  };
+};
+} // namespace NS
+using NS::ForLoopAwaiterCoawaitLookup;
+
+template <class T>
+ForLoopAwaiterCoawaitLookup test_coawait_lookup(T) {
+  Range<T> R;
+  for co_await(auto i : R) {}
+  // expected-error@-1 {{no member named 'await_ready' in 'CoawaitTag<Iter<int>, false>'}}
+}
+template ForLoopAwaiterCoawaitLookup test_coawait_lookup(int); // expected-note {{requested here}}
+
+// FIXME: This test should fail as well since the newly declared operator co_await
+// should not be found by lookup.
+namespace NS2 {
+template <class Iter>
+Awaiter<Iter> operator co_await(CoawaitTag<Iter, false>);
+}
+using NS2::operator co_await;
+template ForLoopAwaiterCoawaitLookup test_coawait_lookup(long);