pw_async2: Replace AwaitableBase with pw::Function

This reduces the per-`co_await` overhead by a single
pointer (the vtable pointer) in exchange for a single
extra pointer (the function pointer) in the promise
type.

Change-Id: I1a0b3778dcb438c1731f530747ae43e103efa639
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/234392
Commit-Queue: Taylor Cramer <cramertj@google.com>
Reviewed-by: Wyatt Hepler <hepler@google.com>
Lint: Lint 🤖 <android-build-ayeaye@system.gserviceaccount.com>

pw_async2/BUILD.bazel

@@ -219,6 +219,7 @@ cc_library(
     deps = [
         ":dispatcher",
         "//pw_allocator:allocator",
+        "//pw_function",
     ],
 )
 

pw_async2/BUILD.gn

@@ -174,6 +174,7 @@ if (pw_toolchain_CXX_STANDARD >= pw_toolchain_STANDARD.CXX20) {
     public_deps = [
       ":dispatcher",
       "$dir_pw_allocator:allocator",
+      dir_pw_function,
     ]
   }
 

pw_async2/CMakeLists.txt

@@ -155,6 +155,7 @@ if(NOT "cxx_std_20" IN_LIST CMAKE_CXX_COMPILE_FEATURES)
     PUBLIC_DEPS
       pw_allocator.allocator
       pw_async2.dispatcher
+      pw_function
     PUBLIC_INCLUDES
       public
   )

pw_async2/public/pw_async2/coro.h

@@ -19,6 +19,7 @@
 #include "pw_allocator/allocator.h"
 #include "pw_allocator/layout.h"
 #include "pw_async2/dispatcher.h"
+#include "pw_function/function.h"
 #include "pw_log/log.h"
 #include "pw_status/status.h"
 #include "pw_status/try.h"
@@ -201,23 +202,17 @@ struct InOut final {
   OptionalOrDefault<T>* output;
 };
 
-// A base class for `Awaitable` instances.
+// Attempt to complete the current pendable value passed to `co_await`,
+// storing its return value inside the `Awaitable` object so that it can
+// be retrieved by the coroutine.
 //
 // Each `co_await` statement creates an `Awaitable` object whose `Pend`
 // method must be completed before the coroutine's `resume()` function can
 // be invoked.
-class AwaitableBase {
- public:
-  // Attempt to complete the current pendable value passed to `co_await`,
-  // storing its return value inside the `Awaitable` object so that it can
-  // be retrieved by the coroutine.
-  virtual Poll<> PendFillReturnValue(Context& cx) = 0;
-
- protected:
-  // A protected destructor ensures that child classes are never destroyed
-  // through a base pointer, so no virtual destructor is needed.
-  ~AwaitableBase() {}
-};
+//
+// `sizeof(void*)` is used as the size since only one pointer capture is
+// required in all cases.
+using PendFillReturnValueFn = pw::Function<Poll<>(Context&), sizeof(void*)>;
 
 // The `promise_type` of `Coro<T>`.
 //
@@ -233,12 +228,12 @@ class CoroPromiseType final {
   // arguments are unused, but must be accepted in order for this to compile.
   template <typename... Args>
   CoroPromiseType(CoroContext& cx, const Args&...)
-      : dealloc_(cx.alloc()), current_awaitable_(nullptr), in_out_(nullptr) {}
+      : dealloc_(cx.alloc()), currently_pending_(nullptr), in_out_(nullptr) {}
 
   // Method-receiver version.
   template <typename MethodReceiver, typename... Args>
   CoroPromiseType(const MethodReceiver&, CoroContext& cx, const Args&...)
-      : dealloc_(cx.alloc()), current_awaitable_(nullptr), in_out_(nullptr) {}
+      : dealloc_(cx.alloc()), currently_pending_(nullptr), in_out_(nullptr) {}
 
   // Get the `Coro<T>` after successfully allocating the coroutine space
   // and constructing `this`.
@@ -319,7 +314,7 @@ class CoroPromiseType final {
   Context& cx() { return *in_out_->input_cx; }
 
   pw::allocator::Deallocator& dealloc_;
-  AwaitableBase* current_awaitable_;
+  PendFillReturnValueFn currently_pending_;
   InOut<T>* in_out_;
 };
 
@@ -328,7 +323,7 @@ class CoroPromiseType final {
 // This wraps a `Pendable` type and implements the awaitable interface
 // expected by the standard coroutine API.
 template <typename Pendable, typename PromiseType>
-class Awaitable final : AwaitableBase {
+class Awaitable final {
  public:
   // The `OutputType` in `Poll<OutputType> Pendable::Pend(Context&)`.
   using OutputType = std::remove_cvref_t<
@@ -352,7 +347,9 @@ class Awaitable final : AwaitableBase {
     Context& cx = promise.promise().cx();
     if (PendFillReturnValue(cx).IsPending()) {
       /// The coroutine should suspend since the await-ed thing is pending.
-      promise.promise().current_awaitable_ = this;
+      promise.promise().currently_pending_ = [this](Context& lambda_cx) {
+        return PendFillReturnValue(lambda_cx);
+      };
       return true;
     }
     return false;
@@ -380,7 +377,7 @@ class Awaitable final : AwaitableBase {
   // This method must return `Ready()` before the coroutine can be safely
   // resumed, as otherwise the return value will not be available when
   // `await_resume` is called to produce the result of `co_await`.
-  Poll<> PendFillReturnValue(Context& cx) final {
+  Poll<> PendFillReturnValue(Context& cx) {
     Poll<OutputType> poll_res(PendableNoPtr().Pend(cx));
     if (poll_res.IsPending()) {
       return Pending();
@@ -466,10 +463,8 @@ class Coro final {
     // If an `Awaitable` value is currently being processed, it must be
     // allowed to complete and store its return value before we can resume
     // the coroutine.
-    if (promise_handle_.promise().current_awaitable_ != nullptr &&
-        promise_handle_.promise()
-            .current_awaitable_->PendFillReturnValue(cx)
-            .IsPending()) {
+    if (promise_handle_.promise().currently_pending_ != nullptr &&
+        promise_handle_.promise().currently_pending_(cx).IsPending()) {
       return Pending();
     }
     // Create the arguments (and output storage) for the coroutine.