c++, coroutines: Fix handling of bool await_suspend() [PR115905].

As noted in the PR the action of the existing implementation was to
treat a false value from await_suspend () as equivalent to "do not
suspend".  Actually it needs to be the equivalent of "resume" - and
we need to restart the dispatcher - since the await_suspend() body
could have already resumed the coroutine.
See also https://github.com/cplusplus/CWG/issues/601 (NAD) for more
discussion.

Since we need to amend the await expansion and the actor build, take
the opportunity to clean up and modernise the code there.  Note that
we need to make the jump back to the dispatcher without any scope
exit cleanups (so we have to use the .CO_SUSPN IFN to do this).

	PR c++/115905

gcc/cp/ChangeLog:

	* coroutines.cc (struct coro_aw_data): Add a member for the
	restart dispatch label.
	(expand_one_await_expression): Rework to modernise and to
	handle the boolean await_suspend() case.
	(build_actor_fn): Rework the dispatcher and allow for a jump
	back to the dispatcher.

gcc/testsuite/ChangeLog:

	* g++.dg/coroutines/torture/pr115905.C: New test.

Signed-off-by: Iain Sandoe <iain@sandoe.co.uk>

2 files changed, 174 insertions, 105 deletions

diff --git a/gcc/cp/coroutines.cc b/gcc/cp/coroutines.cc
index d6acf09..002d575 100644
--- a/gcc/cp/coroutines.cc
+++ b/gcc/cp/coroutines.cc
@@ -1780,6 +1780,7 @@ struct coro_aw_data
   tree cleanup;    /* This is where to go once we complete local destroy.  */
   tree cororet;    /* This is where to go if we suspend.  */
   tree corocont;   /* This is where to go if we continue.  */
+  tree dispatch;   /* This is where we go if we restart the dispatch.  */
   tree conthand;   /* This is the handle for a continuation.  */
   unsigned index;  /* This is our current resume index.  */
 };
@@ -1824,42 +1825,44 @@ co_await_find_in_subtree (tree *stmt, int *, void *d)
    in either.  */
 
 static tree *
-expand_one_await_expression (tree *stmt, tree *await_expr, void *d)
+expand_one_await_expression (tree *expr, tree *await_expr, void *d)
 {
   coro_aw_data *data = (coro_aw_data *) d;
 
-  tree saved_statement = *stmt;
+  tree saved_statement = *expr;
   tree saved_co_await = *await_expr;
 
   tree actor = data->actor_fn;
-  location_t loc = EXPR_LOCATION (*stmt);
+  location_t loc = EXPR_LOCATION (*expr);
   tree var = TREE_OPERAND (saved_co_await, 1);  /* frame slot. */
-  tree expr = TREE_OPERAND (saved_co_await, 2); /* initializer.  */
+  tree init_expr = TREE_OPERAND (saved_co_await, 2); /* initializer.  */
   tree awaiter_calls = TREE_OPERAND (saved_co_await, 3);
 
   tree source = TREE_OPERAND (saved_co_await, 4);
   bool is_final = (source
 		   && TREE_INT_CST_LOW (source) == (int) FINAL_SUSPEND_POINT);
-  bool needs_dtor = TYPE_HAS_NONTRIVIAL_DESTRUCTOR (TREE_TYPE (var));
+
+  /* Build labels for the destinations of the control flow when we are resuming
+     or destroying.  */
   int resume_point = data->index;
-  size_t bufsize = sizeof ("destroy.") + 10;
-  char *buf = (char *) alloca (bufsize);
-  snprintf (buf, bufsize, "destroy.%d", resume_point);
+  char *buf = xasprintf ("destroy.%d", resume_point);
   tree destroy_label = create_named_label_with_ctx (loc, buf, actor);
-  snprintf (buf, bufsize, "resume.%d", resume_point);
+  free (buf);
+  buf = xasprintf ("resume.%d", resume_point);
   tree resume_label = create_named_label_with_ctx (loc, buf, actor);
-  tree empty_list = build_empty_stmt (loc);
+  free (buf);
 
-  tree stmt_list = NULL;
-  tree r;
+  /* This will contain our expanded expression and replace the original
+     expression.  */
+  tree stmt_list = push_stmt_list ();
   tree *await_init = NULL;
 
-  if (!expr)
+  bool needs_dtor = TYPE_HAS_NONTRIVIAL_DESTRUCTOR (TREE_TYPE (var));
+  if (!init_expr)
     needs_dtor = false; /* No need, the var's lifetime is managed elsewhere.  */
   else
     {
-      r = coro_build_cvt_void_expr_stmt (expr, loc);
-      append_to_statement_list_force (r, &stmt_list);
+      finish_expr_stmt (init_expr);
       /* We have an initializer, which might itself contain await exprs.  */
       await_init = tsi_stmt_ptr (tsi_last (stmt_list));
     }
@@ -1870,18 +1873,15 @@ expand_one_await_expression (tree *stmt, tree *await_expr, void *d)
   if (TREE_CODE (TREE_TYPE (ready_cond)) != BOOLEAN_TYPE)
     ready_cond = cp_convert (boolean_type_node, ready_cond,
 			     tf_warning_or_error);
+  tree susp_if = begin_if_stmt ();
   /* Be aggressive in folding here, since there are a significant number of
      cases where the ready condition is constant.  */
   ready_cond = invert_truthvalue_loc (loc, ready_cond);
-  ready_cond
-    = build1_loc (loc, CLEANUP_POINT_EXPR, boolean_type_node, ready_cond);
+  finish_if_stmt_cond (ready_cond, susp_if);
 
-  tree body_list = NULL;
   tree susp_idx = build_int_cst (short_unsigned_type_node, data->index);
-  r = build2_loc (loc, MODIFY_EXPR, short_unsigned_type_node, data->resume_idx,
-		  susp_idx);
-  r = coro_build_cvt_void_expr_stmt (r, loc);
-  append_to_statement_list (r, &body_list);
+  tree r = cp_build_init_expr (data->resume_idx, susp_idx);
+  finish_expr_stmt (r);
 
   /* Find out what we have to do with the awaiter's suspend method.
      [expr.await]
@@ -1900,31 +1900,33 @@ expand_one_await_expression (tree *stmt, tree *await_expr, void *d)
   /* NOTE: final suspend can't resume; the "resume" label in that case
      corresponds to implicit destruction.  */
   if (VOID_TYPE_P (susp_type))
-    {
-      /* We just call await_suspend() and hit the yield.  */
-      suspend = coro_build_cvt_void_expr_stmt (suspend, loc);
-      append_to_statement_list (suspend, &body_list);
-    }
+    /* Void return - just proceed to suspend.  */
+    finish_expr_stmt (suspend);
   else if (TREE_CODE (susp_type) == BOOLEAN_TYPE)
     {
-      /* Boolean return, continue if the call returns false.  */
-      suspend = build1_loc (loc, TRUTH_NOT_EXPR, boolean_type_node, suspend);
-      suspend
-	= build1_loc (loc, CLEANUP_POINT_EXPR, boolean_type_node, suspend);
-      tree go_on = build1_loc (loc, GOTO_EXPR, void_type_node, resume_label);
-      r = build3_loc (loc, COND_EXPR, void_type_node, suspend, go_on,
-		      empty_list);
-      append_to_statement_list (r, &body_list);
+      /* Boolean return, "continue" if the call returns false.  */
+      tree restart_if = begin_if_stmt ();
+      suspend = invert_truthvalue_loc (loc, suspend);
+      finish_if_stmt_cond (suspend, restart_if);
+      /* Resume - so restart the dispatcher, since we do not know if this
+	 coroutine was already resumed from within await_suspend.  We must
+	 exit this scope without cleanups.  */
+      r = build_call_expr_internal_loc (loc, IFN_CO_SUSPN, void_type_node, 1,
+					build_address (data->dispatch));
+      /* This will eventually expand to 'goto coro.restart.dispatch'.  */
+      finish_expr_stmt (r);
+      finish_then_clause (restart_if);
+      finish_if_stmt (restart_if);
     }
   else
     {
+      /* Handle return, save it to the continuation.  */
       r = suspend;
       if (!same_type_ignoring_top_level_qualifiers_p (susp_type,
 						      void_coro_handle_type))
 	r = build1_loc (loc, VIEW_CONVERT_EXPR, void_coro_handle_type, r);
       r = cp_build_init_expr (loc, data->conthand, r);
-      r = build1 (CONVERT_EXPR, void_type_node, r);
-      append_to_statement_list (r, &body_list);
+      finish_expr_stmt (r);
       is_cont = true;
     }
 
@@ -1937,55 +1939,32 @@ expand_one_await_expression (tree *stmt, tree *await_expr, void *d)
   susp_idx = build_int_cst (integer_type_node, data->index);
 
   tree sw = begin_switch_stmt ();
-  tree cond = build_decl (loc, VAR_DECL, NULL_TREE, integer_type_node);
-  DECL_ARTIFICIAL (cond) = 1;
-  DECL_IGNORED_P (cond) = 1;
-  layout_decl (cond, 0);
 
   r = build_call_expr_internal_loc (loc, IFN_CO_YIELD, integer_type_node, 5,
 				    susp_idx, final_susp, r_l, d_l,
 				    data->coro_fp);
-  r = cp_build_init_expr (cond, r);
   finish_switch_cond (r, sw);
   finish_case_label (loc, integer_zero_node, NULL_TREE); /*  case 0: */
   /* Implement the suspend, a scope exit without clean ups.  */
   r = build_call_expr_internal_loc (loc, IFN_CO_SUSPN, void_type_node, 1,
 				    is_cont ? cont : susp);
-  r = coro_build_cvt_void_expr_stmt (r, loc);
-  add_stmt (r); /*   goto ret;  */
+  finish_expr_stmt (r); /* This will eventually expand to 'goto return'.  */
   finish_case_label (loc, integer_one_node, NULL_TREE); /*  case 1:  */
-  r = build1_loc (loc, GOTO_EXPR, void_type_node, resume_label);
-  add_stmt (r); /*  goto resume;  */
+  add_stmt (build_stmt (loc, GOTO_EXPR, resume_label)); /*  goto resume;  */
   finish_case_label (loc, NULL_TREE, NULL_TREE); /* default:;  */
-  r = build1_loc (loc, GOTO_EXPR, void_type_node, destroy_label);
-  add_stmt (r); /* goto destroy;  */
-
-  /* part of finish switch.  */
-  SWITCH_STMT_BODY (sw) = pop_stmt_list (SWITCH_STMT_BODY (sw));
-  pop_switch ();
-  tree scope = SWITCH_STMT_SCOPE (sw);
-  SWITCH_STMT_SCOPE (sw) = NULL;
-  r = do_poplevel (scope);
-  append_to_statement_list (r, &body_list);
+  add_stmt (build_stmt (loc, GOTO_EXPR, destroy_label)); /* goto destroy;  */
 
-  destroy_label = build_stmt (loc, LABEL_EXPR, destroy_label);
-  append_to_statement_list (destroy_label, &body_list);
+  finish_switch_stmt (sw);
+  add_stmt (build_stmt (loc, LABEL_EXPR, destroy_label));
   if (needs_dtor)
-    {
-      tree dtor = build_cleanup (var);
-      append_to_statement_list (dtor, &body_list);
-    }
-  r = build1_loc (loc, GOTO_EXPR, void_type_node, data->cleanup);
-  append_to_statement_list (r, &body_list);
-
-  r = build3_loc (loc, COND_EXPR, void_type_node, ready_cond, body_list,
-		  empty_list);
+    finish_expr_stmt (build_cleanup (var));
+  add_stmt (build_stmt (loc, GOTO_EXPR, data->cleanup));
 
-  append_to_statement_list (r, &stmt_list);
+  finish_then_clause (susp_if);
+  finish_if_stmt (susp_if);
 
   /* Resume point.  */
-  resume_label = build_stmt (loc, LABEL_EXPR, resume_label);
-  append_to_statement_list (resume_label, &stmt_list);
+  add_stmt (build_stmt (loc, LABEL_EXPR, resume_label));
 
   /* This will produce the value (if one is provided) from the co_await
      expression.  */
@@ -1999,14 +1978,11 @@ expand_one_await_expression (tree *stmt, tree *await_expr, void *d)
   /* Get a pointer to the revised statement.  */
   tree *revised = tsi_stmt_ptr (tsi_last (stmt_list));
   if (needs_dtor)
-    {
-      tree dtor = build_cleanup (var);
-      append_to_statement_list (dtor, &stmt_list);
-    }
+    finish_expr_stmt (build_cleanup (var));
   data->index += 2;
 
-  /* Replace the original statement with the expansion.  */
-  *stmt = stmt_list;
+  /* Replace the original expression with the expansion.  */
+  *expr = pop_stmt_list (stmt_list);
 
   /* Now, if the awaitable had an initializer, expand any awaits that might
      be embedded in it.  */
@@ -2275,6 +2251,34 @@ build_actor_fn (location_t loc, tree coro_frame_type, tree actor, tree fnbody,
   tree continue_label
     = create_named_label_with_ctx (loc, "actor.continue.ret", actor);
 
+  /* Build the dispatcher; for each await expression there is an odd entry
+     corresponding to the destruction action and an even entry for the resume
+     one:
+     if (resume index is odd)
+	{
+	  switch (resume index)
+	   case 1:
+	     goto cleanup.
+	   case ... odd suspension point number
+	     .CO_ACTOR (... odd suspension point number)
+	     break;
+	   default:
+	     break;
+	}
+      else
+	{
+	  coro.restart.dispatch:
+	   case 0:
+	     goto start.
+	   case ... even suspension point number
+	     .CO_ACTOR (... even suspension point number)
+	     break;
+	   default:
+	     break;
+	}
+    we should not get here unless something is broken badly.
+     __builtin_trap ();
+*/
   tree lsb_if = begin_if_stmt ();
   tree chkb0 = build2 (BIT_AND_EXPR, short_unsigned_type_node, rat,
 		       build_int_cst (short_unsigned_type_node, 1));
@@ -2284,10 +2288,6 @@ build_actor_fn (location_t loc, tree coro_frame_type, tree actor, tree fnbody,
 
   tree destroy_dispatcher = begin_switch_stmt ();
   finish_switch_cond (rat, destroy_dispatcher);
-  tree ddeflab = finish_case_label (loc, NULL_TREE, NULL_TREE);
-  tree b = build_call_expr_loc (loc, builtin_decl_explicit (BUILT_IN_TRAP), 0);
-  b = coro_build_cvt_void_expr_stmt (b, loc);
-  add_stmt (b);
 
   /* The destroy point numbered #1 is special, in that it is reached from a
      coroutine that is suspended after re-throwing from unhandled_exception().
@@ -2304,32 +2304,33 @@ build_actor_fn (location_t loc, tree coro_frame_type, tree actor, tree fnbody,
     {
       tree l_num = build_int_cst (short_unsigned_type_node, lab_num);
       finish_case_label (loc, l_num, NULL_TREE);
-      b = build_call_expr_internal_loc (loc, IFN_CO_ACTOR, void_type_node, 1,
-					l_num);
-      b = coro_build_cvt_void_expr_stmt (b, loc);
-      add_stmt (b);
-      b = build1 (GOTO_EXPR, void_type_node, CASE_LABEL (ddeflab));
-      add_stmt (b);
+      tree c = build_call_expr_internal_loc (loc, IFN_CO_ACTOR, void_type_node,
+					     1, l_num);
+      finish_expr_stmt (c);
+      finish_break_stmt ();
       lab_num += 2;
     }
+  finish_case_label (loc, NULL_TREE, NULL_TREE);
+  finish_break_stmt ();
 
-  /* Insert the prototype dispatcher.  */
+  /* Finish the destroy dispatcher.  */
   finish_switch_stmt (destroy_dispatcher);
 
   finish_then_clause (lsb_if);
   begin_else_clause (lsb_if);
 
+  /* For the case of a boolean await_resume () that returns 'true' we should
+     restart the dispatch, since we cannot know if additional resumes were
+     executed from within the await_suspend function.  */
+  tree restart_dispatch_label
+    = create_named_label_with_ctx (loc, "coro.restart.dispatch", actor);
+  add_stmt (build_stmt (loc, LABEL_EXPR, restart_dispatch_label));
+
   tree dispatcher = begin_switch_stmt ();
   finish_switch_cond (rat, dispatcher);
   finish_case_label (loc, build_int_cst (short_unsigned_type_node, 0),
 		     NULL_TREE);
-  b = build1 (GOTO_EXPR, void_type_node, actor_begin_label);
-  add_stmt (b);
-
-  tree rdeflab = finish_case_label (loc, NULL_TREE, NULL_TREE);
-  b = build_call_expr_loc (loc, builtin_decl_explicit (BUILT_IN_TRAP), 0);
-  b = coro_build_cvt_void_expr_stmt (b, loc);
-  add_stmt (b);
+  add_stmt (build_stmt (loc, GOTO_EXPR, actor_begin_label));
 
   lab_num = 2;
   /* The final resume should be made to hit the default (trap, UB) entry
@@ -2339,23 +2340,27 @@ build_actor_fn (location_t loc, tree coro_frame_type, tree actor, tree fnbody,
     {
       tree l_num = build_int_cst (short_unsigned_type_node, lab_num);
       finish_case_label (loc, l_num, NULL_TREE);
-      b = build_call_expr_internal_loc (loc, IFN_CO_ACTOR, void_type_node, 1,
-					l_num);
-      b = coro_build_cvt_void_expr_stmt (b, loc);
-      add_stmt (b);
-      b = build1 (GOTO_EXPR, void_type_node, CASE_LABEL (rdeflab));
-      add_stmt (b);
+      tree c = build_call_expr_internal_loc (loc, IFN_CO_ACTOR, void_type_node,
+					     1, l_num);
+      finish_expr_stmt (c);
+      finish_break_stmt ();
       lab_num += 2;
     }
+  finish_case_label (loc, NULL_TREE, NULL_TREE);
+  finish_break_stmt ();
 
-  /* Insert the prototype dispatcher.  */
+  /* Finish the resume dispatcher.  */
   finish_switch_stmt (dispatcher);
   finish_else_clause (lsb_if);
 
   finish_if_stmt (lsb_if);
 
-  tree r = build_stmt (loc, LABEL_EXPR, actor_begin_label);
-  add_stmt (r);
+  /* If we reach here then we've hit UB.  */
+  tree t = build_call_expr_loc (loc, builtin_decl_explicit (BUILT_IN_TRAP), 0);
+  finish_expr_stmt (t);
+
+  /* Now we start building the rewritten function body.  */
+  add_stmt (build_stmt (loc, LABEL_EXPR, actor_begin_label));
 
   /* actor's coroutine 'self handle'.  */
   tree ash = coro_build_frame_access_expr (actor_frame, coro_self_handle_id,
@@ -2365,7 +2370,7 @@ build_actor_fn (location_t loc, tree coro_frame_type, tree actor, tree fnbody,
   /* Should have been set earlier by coro_promise_type_found_p.  */
   gcc_assert (hfa_m);
 
-  r = build1 (CONVERT_EXPR, build_pointer_type (void_type_node), actor_fp);
+  tree r = build1 (CONVERT_EXPR, build_pointer_type (void_type_node), actor_fp);
   vec<tree, va_gc> *args = make_tree_vector_single (r);
   tree hfa = build_new_method_call (ash, hfa_m, &args, NULL_TREE, LOOKUP_NORMAL,
 				    NULL, tf_warning_or_error);
@@ -2478,7 +2483,7 @@ build_actor_fn (location_t loc, tree coro_frame_type, tree actor, tree fnbody,
 
   coro_aw_data data = {actor, actor_fp, resume_idx_var, NULL_TREE,
 		       ash, del_promise_label, ret_label,
-		       continue_label, continuation, 2};
+		       continue_label, restart_dispatch_label, continuation, 2};
   cp_walk_tree (&actor_body, await_statement_expander, &data, NULL);
 
   BIND_EXPR_BODY (actor_bind) = pop_stmt_list (actor_body);
diff --git a/gcc/testsuite/g++.dg/coroutines/torture/pr115905.C b/gcc/testsuite/g++.dg/coroutines/torture/pr115905.C
new file mode 100644
index 0000000..bb28779
--- /dev/null
+++ b/gcc/testsuite/g++.dg/coroutines/torture/pr115905.C
@@ -0,0 +1,64 @@
+// { dg-do run }
+// { dg-additional-options "-std=c++20" }
+
+#include <coroutine>
+#ifdef OUTPUT
+#include <iostream>
+#endif
+
+struct Promise;
+
+struct Handle : std::coroutine_handle<Promise> {
+    using promise_type = Promise;
+};
+
+struct Promise {
+    Handle get_return_object() noexcept {
+        return {Handle::from_promise(*this)};
+    }
+    std::suspend_never initial_suspend() const noexcept { return {}; }
+    std::suspend_never final_suspend() const noexcept { return {}; }
+    void return_void() const noexcept {}
+    void unhandled_exception() const noexcept {}
+};
+
+struct Awaiter {
+    auto await_ready() const noexcept {
+        return false;
+    }
+
+    auto await_suspend(auto self) const noexcept {
+        self.resume();
+        return false;
+    }
+
+    int await_resume() const noexcept { return 42; }
+};
+
+bool OK = false;
+
+Handle Coro() {
+#ifdef OUTPUT
+    std::cout << "1\n";
+    std::cout << co_await Awaiter{} << '\n';
+#else
+    (void) co_await Awaiter{};
+#endif
+#ifdef OUTPUT
+    std::cout << "2\n";
+#endif
+    co_await std::suspend_always{};
+#ifdef OUTPUT
+    std::cout << "3\n";
+#endif
+    OK = true;
+    co_return;
+}
+
+int main() {
+    Coro();
+
+    if (!OK)
+      __builtin_abort ();
+    return 0;
+}