coroutines : Handle for await expressions in for stmts [PR98480].

The handling of await expressions in the init, condition and iteration
expressions of for loops had been omitted.  Fixed thus.

gcc/cp/ChangeLog:

	PR c++/98480
	* coroutines.cc (replace_continue): Rewrite continue into
	'goto label'.
	(await_statement_walker): Handle await expressions in the
	initializer, condition and iteration expressions of for
	loops.

gcc/testsuite/ChangeLog:

	PR c++/98480
	* g++.dg/coroutines/pr98480.C: New test.
	* g++.dg/coroutines/torture/co-await-24-for-init.C: New test.
	* g++.dg/coroutines/torture/co-await-25-for-condition.C: New test.
	* g++.dg/coroutines/torture/co-await-26-for-iteration-expr.C: New test.

4 files changed, 302 insertions, 0 deletions

diff --git a/gcc/testsuite/g++.dg/coroutines/pr98480.C b/gcc/testsuite/g++.dg/coroutines/pr98480.C
new file mode 100644
index 0000000..2d87261
--- /dev/null
+++ b/gcc/testsuite/g++.dg/coroutines/pr98480.C
@@ -0,0 +1,20 @@
+#include <coroutine>
+
+struct future {
+  struct promise_type {
+    void return_value(int) {}
+    auto initial_suspend() { return std::suspend_never{}; }
+    auto final_suspend() noexcept { return std::suspend_never{}; }
+    void unhandled_exception() {}
+    future get_return_object() { return {}; }
+  };
+  bool await_ready() { return true; }
+  void await_suspend(std::coroutine_handle<>) {}
+  int await_resume() { return 0; }
+};
+
+future co_foo() {
+  for( int i = 0; i < co_await future{}; ++i );
+  // ICE -------------^
+  co_return 0;
+}
diff --git a/gcc/testsuite/g++.dg/coroutines/torture/co-await-24-for-init.C b/gcc/testsuite/g++.dg/coroutines/torture/co-await-24-for-init.C
new file mode 100644
index 0000000..1bf2f6d
--- /dev/null
+++ b/gcc/testsuite/g++.dg/coroutines/torture/co-await-24-for-init.C
@@ -0,0 +1,101 @@
+// { dg-do run }
+
+// Test co-await in while condition.
+
+#include "../coro.h"
+
+// boiler-plate for tests of codegen
+#include "../coro1-ret-int-yield-int.h"
+
+/* An awaiter that suspends always and returns an int as the
+   await_resume output.  */
+struct IntAwaiter {
+  int v;
+  IntAwaiter (int _v) : v(_v) {}
+  bool await_ready () { return false; }
+  void await_suspend (coro::coroutine_handle<>) {}
+  int await_resume () { return v; }
+};
+
+struct coro1
+coro_a (bool t)
+{
+  int accum = 0;
+  for (int x = co_await IntAwaiter (3); x < 10; x++)
+    accum += x;
+
+  co_return accum;
+}
+
+struct coro1
+coro_b (bool t)
+{
+  int accum = 0;
+  int x;
+  for (x = co_await IntAwaiter (3); x < 10; x++)
+    accum += x;
+
+  co_return accum;
+}
+
+struct coro1
+coro_c (bool t)
+{
+  int accum = 0;
+  int x = 3;
+  for (co_await IntAwaiter (3); x < 10; x++)
+    accum += x;
+
+  co_return accum;
+}
+
+void
+check_a_coro (struct coro1& x)
+{
+  if (x.handle.done())
+    {
+      PRINT ("check_a_coro: apparently done when we shouldn't be...");
+      abort ();
+    }
+
+  PRINT ("check_a_coro: resume initial suspend");
+  x.handle.resume();
+
+  // will be false - so no yield expected.
+  PRINT ("check_a_coro: resume for init");
+  x.handle.resume();
+
+  int y = x.handle.promise().get_value();
+  if ( y != 42 )
+    {
+      PRINTF ("check_a_coro: apparently wrong value : %d\n", y);
+      abort ();
+    }
+
+  if (!x.handle.done())
+    {
+      PRINT ("check_a_coro: apparently not done...");
+      abort ();
+    }
+}
+
+int main ()
+{
+  {
+    struct coro1 x = coro_a (false);
+    check_a_coro (x);
+  }
+
+  {
+    struct coro1 x = coro_b (false);
+    check_a_coro (x);
+  }
+
+  {
+    struct coro1 x = coro_c (false);
+    check_a_coro (x);
+  }
+
+  PRINT ("main: done");
+  return 0;
+}
diff --git a/gcc/testsuite/g++.dg/coroutines/torture/co-await-25-for-condition.C b/gcc/testsuite/g++.dg/coroutines/torture/co-await-25-for-condition.C
new file mode 100644
index 0000000..2208e34
--- /dev/null
+++ b/gcc/testsuite/g++.dg/coroutines/torture/co-await-25-for-condition.C
@@ -0,0 +1,94 @@
+// { dg-do run }
+
+// Test co-await in while condition.
+
+#include "../coro.h"
+
+// boiler-plate for tests of codegen
+#include "../coro1-ret-int-yield-int.h"
+
+/* An awaiter that suspends always and returns an int as the
+   await_resume output.  */
+struct IntAwaiter {
+  int v;
+  IntAwaiter (int _v) : v(_v) {}
+  bool await_ready () { return false; }
+  void await_suspend (coro::coroutine_handle<>) {}
+  int await_resume () { return v; }
+};
+
+struct coro1
+coro_a (bool t)
+{
+  int accum = 0;
+  for (int x = 3; x < co_await IntAwaiter (10); x++)
+    accum += x;
+
+  co_return accum;
+}
+
+/* An awaiter that suspends always and returns an int as the
+   await_resume output.  */
+struct TenAwaiter {
+  int v;
+  TenAwaiter (int _v) : v(_v) {}
+  bool await_ready () { return false; }
+  void await_suspend (coro::coroutine_handle<>) {}
+  bool await_resume () { return v < 10; }
+};
+
+struct coro1
+coro_b (bool t)
+{
+  int accum = 0;
+  for (int x = 3; co_await TenAwaiter (x); x++)
+    accum += x;
+
+  co_return accum;
+}
+
+void
+check_a_coro (struct coro1& x)
+{
+  if (x.handle.done())
+    {
+      PRINT ("check_a_coro: apparently done when we shouldn't be...");
+      abort ();
+    }
+
+  PRINT ("check_a_coro: resume initial suspend");
+  x.handle.resume();
+
+  // will be false - so no yield expected.
+  PRINT ("check_a_coro: resume loops");
+  while (!x.handle.done())
+    x.handle.resume();
+
+  int y = x.handle.promise().get_value();
+  if ( y != 42 )
+    {
+      PRINTF ("check_a_coro: apparently wrong value : %d\n", y);
+      abort ();
+    }
+
+  if (!x.handle.done())
+    {
+      PRINT ("check_a_coro: apparently not done...");
+      abort ();
+    }
+}
+
+int main ()
+{
+  {
+    struct coro1 x = coro_a (false);
+    check_a_coro (x);
+  }
+  {
+    struct coro1 x = coro_b (false);
+    check_a_coro (x);
+  }
+
+  PRINT ("main: returning");
+  return 0;
+}
diff --git a/gcc/testsuite/g++.dg/coroutines/torture/co-await-26-for-iteration-expr.C b/gcc/testsuite/g++.dg/coroutines/torture/co-await-26-for-iteration-expr.C
new file mode 100644
index 0000000..f361fb5
--- /dev/null
+++ b/gcc/testsuite/g++.dg/coroutines/torture/co-await-26-for-iteration-expr.C
@@ -0,0 +1,87 @@
+// { dg-do run }
+
+// Test co-await in while condition.
+
+#include "../coro.h"
+
+// boiler-plate for tests of codegen
+#include "../coro1-ret-int-yield-int.h"
+
+/* An awaiter that suspends always and returns an int as the
+   await_resume output.  */
+struct IntAwaiter {
+  int v;
+  IntAwaiter (int _v) : v(_v) {}
+  bool await_ready () { return false; }
+  void await_suspend (coro::coroutine_handle<>) {}
+  int await_resume () { return v; }
+};
+
+coro1
+coro_a (bool t)
+{
+  int accum = 0;
+  for (int x = 3; x < 10; x += co_await IntAwaiter (1))
+    accum += x;
+
+  co_return accum;
+}
+
+coro1
+coro_b (bool t)
+{
+  int accum = 0;
+  for (int x = 3; x < 10; x += co_await IntAwaiter (1))
+    {
+       if (x & 1)
+         continue;
+       accum += x;
+    }
+
+  co_return accum;
+}
+
+void check_a_coro (coro1& x, int expected_answer)
+{
+  if (x.handle.done())
+    {
+      PRINT ("check_a_coro: apparently done when we shouldn't be...");
+      abort ();
+    }
+
+  PRINT ("check_a_coro: resume initial suspend");
+  x.handle.resume();
+
+  // will be false - so no yield expected.
+  PRINT ("check_a_coro: resume for init");
+  while (!x.handle.done())
+    x.handle.resume();
+
+  int y = x.handle.promise().get_value();
+  if ( y != expected_answer )
+    {
+      PRINTF ("check_a_coro: apparently wrong value : %d\n", y);
+      abort ();
+    }
+
+  if (!x.handle.done())
+    {
+      PRINT ("check_a_coro: apparently not done...");
+      abort ();
+    }
+}
+
+int main ()
+{
+  {
+    coro1 x = coro_a (false);
+    check_a_coro (x, 42);
+  }
+  {
+    coro1 x = coro_b (false);
+    check_a_coro (x, 18);
+  }
+  
+  PRINT ("main: done");
+  return 0;
+}