Let co_invoke be a CPO

Summary: [Folly] Implement `co_invoke` as a generic CPO with customizations found via ADL.

Reviewed By: lewissbaker

Differential Revision: D18753006

fbshipit-source-id: 16239f0bd6140f56bd23effbf092f2b8279c1350

folly/experimental/coro/AsyncGenerator.h

@@ -19,6 +19,7 @@
 #include <folly/CancellationToken.h>
 #include <folly/Traits.h>
 #include <folly/experimental/coro/CurrentExecutor.h>
+#include <folly/experimental/coro/Invoke.h>
 #include <folly/experimental/coro/Utils.h>
 #include <folly/experimental/coro/ViaIfAsync.h>
 #include <folly/experimental/coro/WithCancellation.h>
@@ -470,6 +471,15 @@ class FOLLY_NODISCARD AsyncGenerator {
     return NextSemiAwaitable{coro_};
   }
 
+  template <typename F, typename... A, typename F_, typename... A_>
+  friend AsyncGenerator
+  folly_co_invoke(tag_t<AsyncGenerator, F, A...>, F_ f, A_... a) {
+    auto r = invoke(static_cast<F&&>(f), static_cast<A&&>(a)...);
+    while (auto v = co_await r.next()) {
+      co_yield std::move(v).value();
+    }
+  }
+
  private:
   friend class detail::AsyncGeneratorPromise<Reference, Value>;
 
@@ -489,34 +499,7 @@ AsyncGeneratorPromise<Reference, Value>::get_return_object() noexcept {
       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
 
-// 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)...);
-  while (auto result = co_await asyncRange.next()) {
-    co_yield* std::move(result);
-  }
-}
-
 } // namespace coro
 } // namespace folly

folly/experimental/coro/Invoke.h

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <folly/functional/Invoke.h>
+#include <folly/lang/CustomizationPoint.h>
+
+namespace folly {
+namespace coro {
+
+namespace invoke_detail {
+void folly_co_invoke(folly::detail::invoke_private_overload&);
+}
+
+//  co_invoke
+//
+//  This utility callable is a safe way to instantiate a coroutine using a
+//  coroutine callable. It guarantees that the callable and the arguments
+//  outlive the coroutine which invocation returns. Otherwise, the callable
+//  and the arguments are not safe to be used within the coroutine body.
+//
+//  For example, if the callable is a lambda with captures, the captures would
+//  not otherwise be safe to use in the coroutine body without using co_invoke.
+//
+//  Models invoke for any callable which returns a coroutine type which declares
+//  support for co_invoke, including:
+//    * AsyncGenerator<...>
+//    * Task<...>
+//
+//  Like invoke in that the callable is invoked with the cvref-qual with which
+//  it is passed to co_invoke and the arguments are passed with the cvref-quals
+//  with which they were passed to co_invoke.
+//
+//  Different from invoke in that the callable and all arguments are decay-
+//  copied and it is the copies that are held for the lifetime of the coroutine
+//  and used in the invocation, whereas invoke merely forwards them directly in
+//  the invocation without first constructing any values copied from them.
+//
+//  Example:
+//
+//      auto gen = co_invoke([range]() -> AsyncGenerator<T> {
+//        for (auto value : range) {
+//          co_yield co_await make<T>(value);
+//        }
+//      });
+//
+//  Example:
+//
+//      auto task = co_invoke([name, dob]() -> Task<T> {
+//        co_return co_await make<T>(name, dob);
+//      });
+//
+//  A word of caution. The callable and each argument is decay-copied by the
+//  customizations. No effort is made to coalesce copies when copies would have
+//  been made with direct invocation.
+//
+//      string name = "foobar"; // will be copied twice
+//      auto task = co_invoke([](string n) -> Task<T> {
+//        co_return co_await make<T>(n);
+//      }, name); // passed as &
+//
+//      string name = "foobar"; // will be moved twice and copied zero times
+//      auto task = co_invoke([](string n) -> Task<T> {
+//        co_return co_await make<T>(n);
+//      }, std::move(name)); // passed as &&
+struct co_invoke_type {
+  template <typename F, typename... A>
+  FOLLY_ERASE constexpr auto operator()(F&& f, A&&... a) const
+      noexcept(noexcept(folly_co_invoke(
+          tag<invoke_result_t<F, A...>, F, A...>,
+          static_cast<F&&>(f),
+          static_cast<A&&>(a)...)))
+          -> decltype(folly_co_invoke(
+              tag<invoke_result_t<F, A...>, F, A...>,
+              static_cast<F&&>(f),
+              static_cast<A&&>(a)...)) {
+    return folly_co_invoke(
+        tag<invoke_result_t<F, A...>, F, A...>,
+        static_cast<F&&>(f),
+        static_cast<A&&>(a)...);
+  }
+};
+FOLLY_DEFINE_CPO(co_invoke_type, co_invoke)
+
+} // namespace coro
+} // namespace folly

folly/experimental/coro/Task.h

@@ -28,6 +28,7 @@
 #include <folly/Traits.h>
 #include <folly/Try.h>
 #include <folly/experimental/coro/CurrentExecutor.h>
+#include <folly/experimental/coro/Invoke.h>
 #include <folly/experimental/coro/Traits.h>
 #include <folly/experimental/coro/Utils.h>
 #include <folly/experimental/coro/ViaIfAsync.h>
@@ -498,6 +499,11 @@ class FOLLY_NODISCARD Task {
     return std::move(task);
   }
 
+  template <typename F, typename... A, typename F_, typename... A_>
+  friend Task folly_co_invoke(tag_t<Task, F, A...>, F_ f, A_... a) {
+    co_return co_await invoke(static_cast<F&&>(f), static_cast<A&&>(a)...);
+  }
+
  private:
   friend class detail::TaskPromiseBase;
   friend class detail::TaskPromise<T>;
@@ -558,43 +564,5 @@ inline Task<void> detail::TaskPromise<void>::get_return_object() noexcept {
       detail::TaskPromise<void>>::from_promise(*this)};
 }
 
-namespace detail {
-template <typename T>
-using is_task = folly::detail::is_instantiation_of<Task, T>;
-
-template <typename F, typename... A, typename F_, typename... A_>
-invoke_result_t<F, A...> co_invoke_(F_ f, A_... a) {
-  co_return co_await folly::invoke(static_cast<F&&>(f), static_cast<A&&>(a)...);
-}
-} // namespace detail
-
-/// co_invoke
-///
-/// This utility function is a safe way to instantiate a coroutine using a
-/// coroutine callable. It guarantees that the callable and the arguments
-/// outlive the coroutine which invocation returns. Otherwise, the callable
-/// and the arguments are not safe to be used within the coroutine body.
-///
-/// For example, if the callable is a lambda with captures, the captures would
-/// not otherwise be safe to use in the coroutine body without using co_invoke.
-///
-/// Models std::invoke for any callable object which returns Task<_>.
-///
-/// Like std::invoke in that the callable is invoked with the cvref-qual with
-/// which it is passed to co_invoke and the arguments are passed with the cvref-
-/// quals with which they were passed to co_invoke.
-///
-/// Different from std::invoke in that the callable and all arguments are decay-
-/// copied and held for the lifetime of the coroutine, whereas std::invoke never
-/// never constructs anything from the callable or the arguments.
-template <typename F, typename... A>
-std::enable_if_t<
-    detail::is_task<invoke_result_t<F, A...>>::value,
-    invoke_result_t<F, A...>>
-co_invoke(F&& f, A&&... a) {
-  return detail::co_invoke_<F, A...>(
-      static_cast<F&&>(f), static_cast<A&&>(a)...);
-}
-
 } // namespace coro
 } // namespace folly

folly/experimental/coro/TimedWait.h

@@ -21,6 +21,7 @@
 
 #include <folly/Optional.h>
 #include <folly/experimental/coro/Baton.h>
+#include <folly/experimental/coro/Invoke.h>
 #include <folly/experimental/coro/Task.h>
 #include <folly/experimental/coro/Traits.h>
 #include <folly/experimental/coro/detail/Helpers.h>

folly/experimental/coro/test/AsyncGeneratorTest.cpp

@@ -24,6 +24,7 @@
 #include <folly/experimental/coro/Baton.h>
 #include <folly/experimental/coro/BlockingWait.h>
 #include <folly/experimental/coro/Collect.h>
+#include <folly/experimental/coro/Invoke.h>
 #include <folly/experimental/coro/Sleep.h>
 #include <folly/experimental/coro/Task.h>
 #include <folly/experimental/coro/WithCancellation.h>

folly/experimental/coro/test/BlockingWaitTest.cpp

@@ -22,6 +22,7 @@
 #include <folly/ScopeGuard.h>
 #include <folly/experimental/coro/Baton.h>
 #include <folly/experimental/coro/BlockingWait.h>
+#include <folly/experimental/coro/Invoke.h>
 #include <folly/experimental/coro/Utils.h>
 #include <folly/fibers/FiberManager.h>
 #include <folly/fibers/FiberManagerMap.h>

folly/experimental/coro/test/CoroTest.cpp

@@ -25,6 +25,7 @@
 #include <folly/experimental/coro/BlockingWait.h>
 #include <folly/experimental/coro/Collect.h>
 #include <folly/experimental/coro/CurrentExecutor.h>
+#include <folly/experimental/coro/Invoke.h>
 #include <folly/experimental/coro/Sleep.h>
 #include <folly/experimental/coro/Task.h>
 #include <folly/experimental/coro/TimedWait.h>

folly/experimental/coro/test/CurrentExecutorTest.cpp

@@ -21,6 +21,7 @@
 #include <folly/experimental/coro/BlockingWait.h>
 #include <folly/experimental/coro/Collect.h>
 #include <folly/experimental/coro/CurrentExecutor.h>
+#include <folly/experimental/coro/Invoke.h>
 #include <folly/experimental/coro/Task.h>
 
 #include <folly/portability/GTest.h>

folly/experimental/coro/test/MergeTest.cpp

@@ -24,6 +24,7 @@
 #include <folly/experimental/coro/BlockingWait.h>
 #include <folly/experimental/coro/Collect.h>
 #include <folly/experimental/coro/CurrentExecutor.h>
+#include <folly/experimental/coro/Invoke.h>
 #include <folly/experimental/coro/Merge.h>
 #include <folly/experimental/coro/Task.h>
 

folly/experimental/coro/test/TaskTest.cpp

@@ -22,6 +22,7 @@
 #include <folly/executors/ManualExecutor.h>
 #include <folly/experimental/coro/Baton.h>
 #include <folly/experimental/coro/BlockingWait.h>
+#include <folly/experimental/coro/Invoke.h>
 #include <folly/experimental/coro/Mutex.h>
 #include <folly/experimental/coro/SharedMutex.h>
 #include <folly/experimental/coro/Task.h>