Add async-stack support to co_awaitTry()

Summary:
Updates the co_awaitTry() adapters to support the co_withAsyncStack()
CPO.

This required updating the WithAsyncStackAwaitable to also support
co_awaitTry() but having its wrapper awaiter type also forward through
the .await_resume_try() method if one exists.

Modified the is_awaiter trait check to no longer check for a valid
await_suspend() call as some awaiters returned from co_withAsyncStack()
now require coroutine_handle<Promise> and no longer accept the
coroutine_handle<void> parameter previously checked by the trait.

Reviewed By: andriigrynenko

Differential Revision: D24499833

fbshipit-source-id: fdcb99a2b85c4235bcd4b5dcc12cf79cb414f034

folly/experimental/coro/Traits.h

@@ -65,10 +65,10 @@ struct _is_valid_await_suspend_return_type : folly::Disjunction<
 ///     std::experimental::coroutine_handle<T> for some T
 /// - awaiter.await_resume()
 ///
-/// Note that the requirement to accept coroutine_handle<void> rather than
-/// just some coroutine_handle<P> is to ensure that the awaitable can be
-/// awaited in any coroutine context where the promise_type does not modify
-/// what is normally awaitable through use of await_transform().
+/// Note that we don't check for a valid await_suspend() method here since
+/// we don't yet know the promise type to use and some await_suspend()
+/// implementations have particular requirements on the promise (eg. the
+/// stack-aware awaiters may require the .getAsyncFrame() method)
 template <typename T, typename = void>
 struct is_awaiter : std::false_type {};
 
@@ -77,15 +77,8 @@ struct is_awaiter<
     T,
     folly::void_t<
         decltype(std::declval<T&>().await_ready()),
-        decltype(std::declval<T&>().await_suspend(
-            std::declval<std::experimental::coroutine_handle<void>>())),
         decltype(std::declval<T&>().await_resume())>>
-    : folly::Conjunction<
-          std::is_same<bool, decltype(std::declval<T&>().await_ready())>,
-          detail::_is_valid_await_suspend_return_type<decltype(
-              std::declval<T&>().await_suspend(
-                  std::declval<
-                      std::experimental::coroutine_handle<void>>()))>> {};
+    : std::is_same<bool, decltype(std::declval<T&>().await_ready())> {};
 
 template <typename T>
 constexpr bool is_awaiter_v = is_awaiter<T>::value;

folly/experimental/coro/ViaIfAsync.h

@@ -242,6 +242,15 @@ class StackAwareViaIfAsyncAwaiter {
     return awaiter_.await_resume();
   }
 
+  template <
+      typename Awaiter2 = Awaiter,
+      typename Result = decltype(std::declval<Awaiter2&>().await_resume_try())>
+  Result await_resume_try() noexcept(
+      noexcept(std::declval<Awaiter&>().await_resume_try())) {
+    viaCoroutine_.destroy();
+    return awaiter_.await_resume_try();
+  }
+
  private:
   CoroutineType viaCoroutine_;
   WithAsyncStackAwaitable awaitable_;
@@ -566,6 +575,19 @@ class TryAwaitable {
                                 }};
   }
 
+  template <
+      typename T2 = T,
+      typename Result =
+          decltype(folly::coro::co_withAsyncStack(std::declval<T2>()))>
+  friend TryAwaitable<Result>
+  tag_invoke(cpo_t<co_withAsyncStack>, TryAwaitable&& awaitable) noexcept(
+      noexcept(folly::coro::co_withAsyncStack(std::declval<T2>()))) {
+    return TryAwaitable<Result>{std::in_place, [&]() -> decltype(auto) {
+                                  return folly::coro::co_withAsyncStack(
+                                      static_cast<T&&>(awaitable.inner_));
+                                }};
+  }
+
   template <
       typename T2 = T,
       typename Result = decltype(folly::coro::co_viaIfAsync(

folly/experimental/coro/WithAsyncStack.h

@@ -106,17 +106,20 @@ class WithAsyncStackCoroutine {
 
 template <typename Awaitable>
 class WithAsyncStackAwaiter {
+  using Awaiter = awaiter_type_t<Awaitable>;
+
  public:
   explicit WithAsyncStackAwaiter(Awaitable&& awaitable)
       : awaiter_(folly::coro::get_awaiter(static_cast<Awaitable&&>(awaitable))),
         coroWrapper_(WithAsyncStackCoroutine::create()) {}
 
-  decltype(auto) await_ready() noexcept(noexcept(awaiter_.await_ready())) {
+  auto await_ready() noexcept(noexcept(std::declval<Awaiter&>().await_ready()))
+      -> decltype(std::declval<Awaiter&>().await_ready()) {
     return awaiter_.await_ready();
   }
 
   template <typename Promise>
-  FOLLY_CORO_AWAIT_SUSPEND_NONTRIVIAL_ATTRIBUTES decltype(auto) await_suspend(
+  FOLLY_CORO_AWAIT_SUSPEND_NONTRIVIAL_ATTRIBUTES auto await_suspend(
       std::experimental::coroutine_handle<Promise> h) {
     AsyncStackFrame& callerFrame = h.promise().getAsyncFrame();
     AsyncStackRoot* stackRoot = callerFrame.getStackRoot();
@@ -145,11 +148,21 @@ class WithAsyncStackAwaiter {
     }
   }
 
-  decltype(auto) await_resume() noexcept(noexcept(awaiter_.await_resume())) {
+  auto await_resume() noexcept(
+      noexcept(std::declval<Awaiter&>().await_resume()))
+      -> decltype(std::declval<Awaiter&>().await_resume()) {
     coroWrapper_ = {};
     return awaiter_.await_resume();
   }
 
+  template <typename Awaiter2 = Awaiter>
+  auto await_resume_try() noexcept(
+      noexcept(std::declval<Awaiter2&>().await_resume_try()))
+      -> decltype(std::declval<Awaiter2&>().await_resume_try()) {
+    coroWrapper_ = {};
+    return awaiter_.await_resume_try();
+  }
+
  private:
   awaiter_type_t<Awaitable> awaiter_;
   WithAsyncStackCoroutine coroWrapper_;

folly/experimental/coro/test/TraitsTest.cpp

@@ -57,17 +57,6 @@ struct WrongAwaitReadyReturnType {
   int await_resume();
 };
 
-struct MissingAwaitSuspend {
-  bool await_ready();
-  int await_resume();
-};
-
-struct WrongAwaitSuspendArgType {
-  bool await_ready();
-  void await_suspend(float);
-  int await_resume();
-};
-
 struct MissingAwaitResume {
   bool await_ready();
   void await_suspend(std::experimental::coroutine_handle<void>);
@@ -96,8 +85,6 @@ static_assert(!is_awaiter_v<void>, "");
 static_assert(!is_awaiter_v<int>, "");
 static_assert(!is_awaiter_v<MissingAwaitReady>, "");
 static_assert(!is_awaiter_v<WrongAwaitReadyReturnType>, "");
-static_assert(!is_awaiter_v<MissingAwaitSuspend>, "");
-static_assert(!is_awaiter_v<WrongAwaitSuspendArgType>, "");
 static_assert(!is_awaiter_v<MissingAwaitResume>, "");
 static_assert(!is_awaiter_v<MemberOperatorCoAwait>, "");