Commit 5513c472 authored by Lewis Baker's avatar Lewis Baker Committed by Facebook Github Bot

Add overload of folly::coro::co_invoke() for AsyncGenerator

Summary: The `co_invoke()` helper function now supports callables that return `AsyncGenerator`.

Reviewed By: rhodo

Differential Revision: D15054793

fbshipit-source-id: d7edc7db7d2fa44115f2c25872a6849b1e07bd4c
parent e5569cdd
...@@ -480,7 +480,37 @@ AsyncGeneratorPromise<Reference, Value>::get_return_object() noexcept { ...@@ -480,7 +480,37 @@ AsyncGeneratorPromise<Reference, Value>::get_return_object() noexcept {
AsyncGeneratorPromise<Reference, Value>>::from_promise(*this)}; AsyncGeneratorPromise<Reference, Value>>::from_promise(*this)};
} }
template <typename T>
inline constexpr bool is_async_generator_v = false;
template <typename Reference, typename Value>
inline constexpr bool is_async_generator_v<AsyncGenerator<Reference, Value>> =
true;
} // namespace detail } // namespace detail
// Helper for immediately invoking a lambda with captures that returns an
// AsyncGenerator to keep the lambda alive until the generator completes.
//
// Example:
// auto gen = co_invoke([min, max]() -> AsyncGenerator<T> {
// for (int i = min; i <= max; ++i) {
// co_yield co_await doSomething(i);
// }
// });
template <typename Func, typename... Args>
auto co_invoke(Func func, Args... args) -> std::enable_if_t<
detail::is_async_generator_v<invoke_result_t<Func, Args...>>,
invoke_result_t<Func, Args...>> {
auto asyncRange =
folly::invoke(static_cast<Func&&>(func), static_cast<Args&&>(args)...);
const auto itEnd = asyncRange.end();
auto it = co_await asyncRange.begin();
while (it != itEnd) {
co_yield* it;
co_await++ it;
}
}
} // namespace coro } // namespace coro
} // namespace folly } // namespace folly
...@@ -391,4 +391,21 @@ TEST(AsyncGenerator, ExplicitValueType) { ...@@ -391,4 +391,21 @@ TEST(AsyncGenerator, ExplicitValueType) {
CHECK_EQ("au revoir", items["bar"]); CHECK_EQ("au revoir", items["bar"]);
} }
TEST(AsyncGenerator, InvokeLambda) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
auto ptr = std::make_unique<int>(123);
auto gen = folly::coro::co_invoke(
[p = std::move(ptr)]() mutable
-> folly::coro::AsyncGenerator<std::unique_ptr<int>&&> {
co_yield std::move(p);
});
auto it = co_await gen.begin();
CHECK(it != gen.end());
ptr = *it;
CHECK(ptr);
CHECK(*ptr == 123);
}());
}
#endif #endif
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment