From c017fc2d6961962ee87ae387462a099242dfbbd2 Mon Sep 17 00:00:00 2001
From: Ingvar Stepanyan <me@rreverser.com>
Date: Sat, 18 Nov 2023 14:19:38 +0000
Subject: [PATCH] Allow subclassing coroutines (#20682)

Slight refactoring of coroutine glue to allow subclassing via the public API.

This doesn't change how `co_await some_val;` behaves, but makes it possible to create your own subclasses that can also be awaited.
---
 system/include/emscripten/val.h | 22 +++++++++++-----------
 system/lib/embind/bind.cpp      |  2 +-
 test/embind/test_val_coro.cpp   | 20 ++++++++++++++++++++
 3 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/system/include/emscripten/val.h b/system/include/emscripten/val.h
index 8d1e3110744a..30c9cfea6d51 100644
--- a/system/include/emscripten/val.h
+++ b/system/include/emscripten/val.h
@@ -596,7 +596,10 @@ class val {
   constexpr nullptr_t end() const { return nullptr; }
 
 #if __cplusplus >= 202002L
-  struct promise_type;
+  class awaiter;
+  awaiter operator co_await() const;
+
+  class promise_type;
 #endif
 
 private:
@@ -660,12 +663,11 @@ inline val::iterator val::begin() const {
 }
 
 #if __cplusplus >= 202002L
-namespace internal {
 // Awaiter defines a set of well-known methods that compiler uses
 // to drive the argument of the `co_await` operator (regardless
 // of the type of the parent coroutine).
 // This one is used for Promises represented by the `val` type.
-class val_awaiter {
+class val::awaiter {
   // State machine holding awaiter's current state. One of:
   //  - initially created with promise
   //  - waiting with a given coroutine handle
@@ -677,11 +679,11 @@ class val_awaiter {
   constexpr static std::size_t STATE_RESULT = 2;
 
 public:
-  val_awaiter(val&& promise)
-    : state(std::in_place_index<STATE_PROMISE>, std::move(promise)) {}
+  awaiter(const val& promise)
+    : state(std::in_place_index<STATE_PROMISE>, promise) {}
 
   // just in case, ensure nobody moves / copies this type around
-  val_awaiter(val_awaiter&&) = delete;
+  awaiter(awaiter&&) = delete;
 
   // Promises don't have a synchronously accessible "ready" state.
   bool await_ready() { return false; }
@@ -705,6 +707,9 @@ class val_awaiter {
   // of the `co_await ...` expression - in our case, the stored value.
   val await_resume() { return std::move(std::get<STATE_RESULT>(state)); }
 };
+
+inline val::awaiter val::operator co_await() const {
+  return {*this};
 }
 
 // `promise_type` is a well-known subtype with well-known method names
@@ -742,11 +747,6 @@ class val::promise_type {
   void return_value(T&& value) {
     resolve(std::forward<T>(value));
   }
-
-  // Return our awaiter on `co_await promise`.
-  internal::val_awaiter await_transform(val promise) {
-    return {std::move(promise)};
-  }
 };
 #endif
 
diff --git a/system/lib/embind/bind.cpp b/system/lib/embind/bind.cpp
index 676e888b2db6..3ea5f34812b0 100644
--- a/system/lib/embind/bind.cpp
+++ b/system/lib/embind/bind.cpp
@@ -63,7 +63,7 @@ void _embind_register_bindings(InitFunc* f) {
   init_funcs = f;
 }
 
-void _emval_coro_resume(val_awaiter* awaiter, EM_VAL result) {
+void _emval_coro_resume(val::awaiter* awaiter, EM_VAL result) {
   awaiter->resume_with(val::take_ownership(result));
 }
 
diff --git a/test/embind/test_val_coro.cpp b/test/embind/test_val_coro.cpp
index 88eeff543781..442667e2a56c 100644
--- a/test/embind/test_val_coro.cpp
+++ b/test/embind/test_val_coro.cpp
@@ -20,12 +20,32 @@ val promise_sleep(int ms, int result = 0) {
   return val::take_ownership(promise_sleep_impl(ms, result));
 }
 
+// Test that we can subclass and make custom awaitable types.
+template <typename T>
+class typed_promise: public val {
+public:
+  typed_promise(val&& promise): val(std::move(promise)) {}
+
+  auto operator co_await() const {
+    struct typed_awaiter: public val::awaiter {
+      T await_resume() {
+        return val::awaiter::await_resume().template as<T>();
+      }
+    };
+
+    return typed_awaiter(*this);
+  }
+};
+
 val asyncCoro() {
   // check that just sleeping works
   co_await promise_sleep(1);
   // check that sleeping and receiving value works
   val v = co_await promise_sleep(1, 12);
   assert(v.as<int>() == 12);
+  // check that awaiting a subclassed promise works and returns the correct type
+  int x = co_await typed_promise<int>(promise_sleep(1, 23));
+  assert(x == 23);
   // check that returning value works (checked by JS in tests)
   co_return 34;
 }