Add co_result support to AsyncGenerator

Summary: Mirrors Task. One behavioral difference is that yielding an empty Try from a Task converts into a UsingUninitializedTry exception, while for AsyncGenerator it converts into a return. This is for consistency with Thrift's ClientBufferedStream to avoid confusing users, but there is also an argument for converting to UsingUninitializedTry instead.

Reviewed By: yfeldblum

Differential Revision: D25379361

fbshipit-source-id: 2a2878185696566b29345c99fee7ba74bef4ae8a

folly/experimental/coro/AsyncGenerator.h

@@ -133,6 +133,17 @@ class AsyncGeneratorPromise {
     return {};
   }
 
+  YieldAwaiter yield_value(co_result<Value>&& res) noexcept {
+    if (res.result().hasValue()) {
+      return yield_value(std::move(res.result().value()));
+    } else if (res.result().hasException()) {
+      return yield_value(co_error(res.result().exception()));
+    } else {
+      return_void();
+      return {};
+    }
+  }
+
   void unhandled_exception() noexcept {
     DCHECK(state_ == State::INVALID);
     folly::coro::detail::activate(exceptionPtr_, std::current_exception());

folly/experimental/coro/Error.h

@@ -20,6 +20,7 @@
 #include <type_traits>
 
 #include <folly/ExceptionWrapper.h>
+#include <folly/Try.h>
 
 namespace folly {
 namespace coro {
@@ -45,5 +46,20 @@ class co_error final {
   exception_wrapper ex_;
 };
 
+template <typename T>
+class co_result final {
+ public:
+  explicit co_result(Try<T>&& result) noexcept(
+      std::is_nothrow_move_constructible<T>::value)
+      : result_(std::move(result)) {}
+
+  const Try<T>& result() const { return result_; }
+
+  Try<T>& result() { return result_; }
+
+ private:
+  Try<T> result_;
+};
+
 } // namespace coro
 } // namespace folly

folly/experimental/coro/Task.h

@@ -53,24 +53,6 @@ class Task;
 template <typename T = void>
 class TaskWithExecutor;
 
-template <typename T>
-class co_result final {
- public:
-  explicit co_result(Try<T>&& result) noexcept(
-      std::is_nothrow_move_constructible<T>::value)
-      : result_(std::move(result)) {}
-
-  const Try<T>& result() const { return result_; }
-
-  Try<T>& result() { return result_; }
-
- private:
-  Try<T> result_;
-};
-
-template <class T>
-co_result(Try<T>)->co_result<T>;
-
 namespace detail {
 
 class TaskPromiseBase {

folly/experimental/coro/test/AsyncGeneratorTest.cpp

@@ -505,6 +505,78 @@ TEST(AsyncGenerator, YieldCoError) {
   }());
 }
 
+TEST(AsyncGenerator, YieldCoResult) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    auto gen = []() -> folly::coro::AsyncGenerator<std::string> {
+      co_yield folly::coro::co_result(folly::Try<std::string>("foo"));
+      co_yield folly::coro::co_result(folly::Try<std::string>("bar"));
+      co_yield folly::coro::co_result(folly::Try<std::string>(SomeError{}));
+      ADD_FAILURE();
+    }();
+
+    auto item1 = co_await gen.next();
+    CHECK(item1.value() == "foo");
+    auto item2 = co_await gen.next();
+    CHECK(item2.value() == "bar");
+
+    EXPECT_THROW(co_await gen.next(), SomeError);
+
+    gen = []() -> folly::coro::AsyncGenerator<std::string> {
+      co_yield folly::coro::co_result(folly::Try<std::string>("foo"));
+      co_yield folly::coro::co_result(folly::Try<std::string>("bar"));
+      co_yield folly::coro::co_result(folly::Try<std::string>());
+      ADD_FAILURE();
+    }();
+
+    item1 = co_await gen.next();
+    CHECK(item1.value() == "foo");
+    item2 = co_await gen.next();
+    CHECK(item2.value() == "bar");
+    EXPECT_FALSE(co_await gen.next());
+  }());
+}
+
+TEST(AsyncGenerator, CoResultProxyExample) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    auto makeProxy = [](folly::coro::AsyncGenerator<int> gen)
+        -> folly::coro::AsyncGenerator<int> {
+      while (true) {
+        auto t = co_await folly::coro::co_awaitTry(gen.next());
+        co_yield folly::coro::co_result(std::move(t));
+      }
+    };
+
+    auto gen = []() -> folly::coro::AsyncGenerator<int> {
+      for (int i = 0; i < 5; ++i) {
+        co_yield i;
+      }
+    }();
+
+    auto proxy = makeProxy(std::move(gen));
+
+    for (int i = 0; i < 5; ++i) {
+      auto val = co_await proxy.next();
+      EXPECT_EQ(*val, i);
+    }
+    EXPECT_FALSE(co_await proxy.next());
+
+    gen = []() -> folly::coro::AsyncGenerator<int> {
+      for (int i = 0; i < 5; ++i) {
+        co_yield i;
+      }
+      throw SomeError();
+    }();
+
+    proxy = makeProxy(std::move(gen));
+
+    for (int i = 0; i < 5; ++i) {
+      auto val = co_await proxy.next();
+      EXPECT_EQ(*val, i);
+    }
+    EXPECT_THROW(co_await proxy.next(), SomeError);
+  }());
+}
+
 TEST(AsyncGeneraor, CoAwaitTry) {
   folly::coro::blockingWait([]() -> folly::coro::Task<void> {
     auto gen = []() -> folly::coro::AsyncGenerator<std::string> {

folly/experimental/coro/test/TaskTest.cpp

@@ -494,6 +494,12 @@ TEST_F(TaskTest, YieldTry) {
           co_await co_awaitTry(std::move(innerTaskInt)));
     }());
     EXPECT_TRUE(retInt.hasValue());
+
+    EXPECT_THROW(
+        co_await[&]()->folly::coro::Task<int> {
+          co_yield folly::coro::co_result(folly::Try<int>());
+        }(),
+        UsingUninitializedTry);
   }());
 }