Make folly::coro::co_invoke work with folly::coro::Generator

Summary: This diff makes folly::coro::co_invoke work with folly::coro::Generator.

Reviewed By: yfeldblum

Differential Revision: D23332912

fbshipit-source-id: b3e25d160a94aed8892bf55ad5007f4018786e47

folly/experimental/coro/Generator.h

@@ -22,6 +22,8 @@
 #include <type_traits>
 #include <utility>
 
+#include <folly/experimental/coro/Invoke.h>
+
 namespace folly {
 namespace coro {
 
@@ -269,6 +271,14 @@ class Generator {
     std::swap(m_promise, other.m_promise);
   }
 
+  template <typename F, typename... A, typename F_, typename... A_>
+  friend Generator folly_co_invoke(tag_t<Generator, F, A...>, F_ f, A_... a) {
+    auto&& r = invoke(static_cast<F&&>(f), static_cast<A&&>(a)...);
+    for (auto&& v : r) {
+      co_yield std::move(v);
+    }
+  }
+
  private:
   friend class promise_type;
 

folly/experimental/coro/test/GeneratorTest.cpp

@@ -292,6 +292,20 @@ TEST_F(GeneratorTest, UsageInStandardAlgorithms) {
     EXPECT_FALSE(std::equal(a.begin(), a.end(), b.begin(), b.end()));
   }
 }
+
+TEST_F(GeneratorTest, InvokeLambda) {
+  auto ptr = std::make_unique<int>(123);
+  auto gen = folly::coro::co_invoke(
+      [p = std::move(
+           ptr)]() mutable -> folly::coro::Generator<std::unique_ptr<int>&&> {
+        co_yield std::move(p);
+      });
+
+  auto it = gen.begin();
+  auto result = std::move(*it);
+  EXPECT_NE(result, nullptr);
+  EXPECT_EQ(*result, 123);
+}
 } // namespace coro
 } // namespace folly