add the get_await_completion_adaptor query and use it in as_awaitable

Author: ericniebler

include/stdexec/__detail/__as_awaitable.hpp

@@ -22,12 +22,14 @@
 #include "__concepts.hpp"
 #include "__connect.hpp"
 #include "__meta.hpp"
+#include "__queries.hpp"
 #include "__tag_invoke.hpp"
 #include "__type_traits.hpp"
+#include "__variant.hpp"
 
 #include <exception>
+#include <functional>  // for std::identity
 #include <system_error>
-#include <variant>
 
 namespace STDEXEC
 {
@@ -49,6 +51,23 @@ namespace STDEXEC
     template <class _Sender, class _Promise>
     using __value_t = __decay_t<
       __value_types_of_t<_Sender, env_of_t<_Promise&>, __q<__single_value>, __msingle_or<void>>>;
+
+    inline constexpr auto __get_await_completion_adaptor =
+      __with_default{get_await_completion_adaptor, std::identity{}};
+
+    template <class _Sender>
+    using __adapt_completion_t = __result_of<__get_await_completion_adaptor, env_of_t<_Sender>>;
+
+    template <class _Sender>
+    constexpr auto __adapt_sender_for_await(_Sender&& __sndr)
+      noexcept(__nothrow_callable<__adapt_completion_t<_Sender>, _Sender>) -> decltype(auto)
+    {
+      return __get_await_completion_adaptor(get_env(__sndr))(static_cast<_Sender&&>(__sndr));
+    }
+
+    template <class _Sender>
+    using __adapted_sender_t =
+      __remove_rvalue_reference_t<__call_result_t<__adapt_completion_t<_Sender>, _Sender>>;
   }  // namespace __detail
 
   /////////////////////////////////////////////////////////////////////////////
@@ -62,8 +81,7 @@ namespace STDEXEC
     using __value_or_void_t = __if_c<__same_as<_Value, void>, __void, _Value>;
 
     template <class _Value>
-    using __expected_t =
-      std::variant<std::monostate, __value_or_void_t<_Value>, std::exception_ptr>;
+    using __expected_t = __variant<__value_or_void_t<_Value>, std::exception_ptr>;
 
     // Helper to cast a coroutine_handle<void> to coroutine_handle<_Promise>
     template <class _Promise>
@@ -84,7 +102,7 @@ namespace STDEXEC
       {
         STDEXEC_TRY
         {
-          __result_->template emplace<1>(static_cast<_Us&&>(__us)...);
+          __result_->template emplace<0>(static_cast<_Us&&>(__us)...);
           __continuation_.resume();
         }
         STDEXEC_CATCH_ALL
@@ -97,11 +115,11 @@ namespace STDEXEC
       void set_error(_Error&& __err) noexcept
       {
         if constexpr (__decays_to<_Error, std::exception_ptr>)
-          __result_->template emplace<2>(static_cast<_Error&&>(__err));
+          __result_->template emplace<1>(static_cast<_Error&&>(__err));
         else if constexpr (__decays_to<_Error, std::error_code>)
-          __result_->template emplace<2>(std::make_exception_ptr(std::system_error(__err)));
+          __result_->template emplace<1>(std::make_exception_ptr(std::system_error(__err)));
         else
-          __result_->template emplace<2>(std::make_exception_ptr(static_cast<_Error&&>(__err)));
+          __result_->template emplace<1>(std::make_exception_ptr(static_cast<_Error&&>(__err)));
         __continuation_.resume();
       }
 
@@ -145,22 +163,22 @@ namespace STDEXEC
       {
         switch (__result_.index())
         {
-        case 0:  // receiver contract not satisfied
+        case __variant_npos:  // receiver contract not satisfied
           STDEXEC_ASSERT(false && +"_Should never get here" == nullptr);
           break;
-        case 1:  // set_value
+        case 0:  // set_value
           if constexpr (!__same_as<_Value, void>)
-            return static_cast<_Value&&>(std::get<1>(__result_));
+            return static_cast<_Value&&>(__var::__get<0>(__result_));
           else
             return;
-        case 2:  // set_error
-          std::rethrow_exception(std::get<2>(__result_));
+        case 1:  // set_error
+          std::rethrow_exception(__var::__get<1>(__result_));
         }
         std::terminate();
       }
 
      protected:
-      __expected_t<_Value> __result_;
+      __expected_t<_Value> __result_{__no_init};
     };
 
     template <class _Promise, class _Sender>
@@ -185,14 +203,23 @@ namespace STDEXEC
     };
 
     template <class _Sender, class _Promise>
-    concept __awaitable_sender = sender_in<_Sender, env_of_t<_Promise&>>
-                              && __minvocable_q<__detail::__value_t, _Sender, _Promise>
-                              && sender_to<_Sender, __receiver_t<_Sender, _Promise>>
-                              && requires(_Promise& __promise) {
-                                   {
-                                     __promise.unhandled_stopped()
-                                   } -> __std::convertible_to<__std::coroutine_handle<>>;
-                                 };
+    STDEXEC_HOST_DEVICE_DEDUCTION_GUIDE
+    __sender_awaitable(_Sender&&, __std::coroutine_handle<_Promise>)
+      -> __sender_awaitable<_Promise, _Sender>;
+
+    template <class _Sender, class _Promise>
+    concept __awaitable_adapted_sender = sender_in<_Sender, env_of_t<_Promise&>>
+                                      && __minvocable_q<__detail::__value_t, _Sender, _Promise>
+                                      && sender_to<_Sender, __receiver_t<_Sender, _Promise>>
+                                      && requires(_Promise& __promise) {
+                                           {
+                                             __promise.unhandled_stopped()
+                                           } -> __std::convertible_to<__std::coroutine_handle<>>;
+                                         };
+
+    template <class _Sender, class _Promise>
+    concept __awaitable_sender =
+      __awaitable_adapted_sender<__detail::__adapted_sender_t<_Sender>, _Promise>;
 
     struct __unspecified
     {
@@ -214,32 +241,33 @@ namespace STDEXEC
     template <class _Tp, class _Promise>
     static consteval auto __get_declfn() noexcept
     {
-      using __as_awaitable::__unspecified;
+      using namespace __as_awaitable;
       if constexpr (__connect_await::__has_as_awaitable_member<_Tp, _Promise>)
       {
         using __result_t = decltype(__declval<_Tp>().as_awaitable(__declval<_Promise&>()));
         constexpr bool __is_nothrow = noexcept(
           __declval<_Tp>().as_awaitable(__declval<_Promise&>()));
         return __declfn<__result_t, __is_nothrow>();
-        // NOLINTNEXTLINE(bugprone-branch-clone)
       }
-      else if constexpr (__awaitable<_Tp, __unspecified>)
-      {  // NOT __awaitable<_Tp, _Promise> !!
+      else if constexpr (__awaitable<_Tp, __unspecified>)  // NOT __awaitable<_Tp, _Promise> !!
+      {                                                    // NOLINT(bugprone-branch-clone)
         return __declfn<_Tp&&>();
       }
-      else if constexpr (__as_awaitable::__awaitable_sender<_Tp, _Promise>)
+      else if constexpr (__awaitable_sender<_Tp, _Promise>)
       {
-        using __result_t = __as_awaitable::__sender_awaitable<_Promise, _Tp>;
-        constexpr bool __is_nothrow =
-          __nothrow_constructible_from<__result_t, _Tp, __std::coroutine_handle<_Promise>>;
+        using __result_t            = decltype(  //
+          __sender_awaitable{__detail::__adapt_sender_for_await(__declval<_Tp>()),
+                             __std::coroutine_handle<_Promise>()});
+        constexpr bool __is_nothrow = noexcept(
+          __sender_awaitable{__detail::__adapt_sender_for_await(__declval<_Tp>()),
+                             __std::coroutine_handle<_Promise>()});
         return __declfn<__result_t, __is_nothrow>();
-        // NOT TO SPEC
       }
-      else if constexpr (__as_awaitable::__incompatible_sender<_Tp, _Promise>)
+      else if constexpr (__incompatible_sender<_Tp, _Promise>)
       {
-        // It's a sender, but it isn't a sender in the current promise's environment, so
-        // we can return the error type that results from trying to compute the sender's
-        // value type:
+        // NOT TO SPEC: It's a sender, but it isn't a sender in the current promise's
+        // environment, so we can return the error type that results from trying to
+        // compute the sender's value type:
         return __declfn<__detail::__value_t<_Tp, _Promise>>();
       }
       else
@@ -253,24 +281,24 @@ namespace STDEXEC
     auto operator()(_Tp&& __t, _Promise& __promise) const noexcept(noexcept(_DeclFn()))
       -> decltype(_DeclFn())
     {
-      using __as_awaitable::__unspecified;
+      using namespace __as_awaitable;
       if constexpr (__connect_await::__has_as_awaitable_member<_Tp, _Promise>)
       {
         using __result_t = decltype(static_cast<_Tp&&>(__t).as_awaitable(__promise));
         static_assert(__awaitable<__result_t, _Promise>);
         return static_cast<_Tp&&>(__t).as_awaitable(__promise);
-        // NOLINTNEXTLINE(bugprone-branch-clone)
       }
-      else if constexpr (__awaitable<_Tp, __unspecified>)
-      {  // NOT __awaitable<_Tp, _Promise> !!
+      else if constexpr (__awaitable<_Tp, __unspecified>)  // NOT __awaitable<_Tp, _Promise> !!
+      {                                                    // NOLINT(bugprone-branch-clone)
         return static_cast<_Tp&&>(__t);
       }
-      else if constexpr (__as_awaitable::__awaitable_sender<_Tp, _Promise>)
+      else if constexpr (__awaitable_sender<_Tp, _Promise>)
       {
         auto __hcoro = __std::coroutine_handle<_Promise>::from_promise(__promise);
-        return __as_awaitable::__sender_awaitable<_Promise, _Tp>{static_cast<_Tp&&>(__t), __hcoro};
+        return __sender_awaitable{__detail::__adapt_sender_for_await(static_cast<_Tp&&>(__t)),
+                                  __hcoro};
       }
-      else if constexpr (__as_awaitable::__incompatible_sender<_Tp, _Promise>)
+      else if constexpr (__incompatible_sender<_Tp, _Promise>)
       {
         return __detail::__value_t<_Tp, _Promise>();
       }

include/stdexec/__detail/__execution_fwd.hpp

@@ -117,6 +117,7 @@ namespace STDEXEC
   template <__completion_tag _CPO>
   struct get_completion_behavior_t;
   struct get_domain_t;
+  struct get_await_completion_adaptor_t;
 
   struct __debug_env_t;
 
@@ -127,8 +128,9 @@ namespace STDEXEC
   template <__completion_tag _CPO>
   extern get_completion_scheduler_t<_CPO> const get_completion_scheduler;
   template <class _CPO = void>
-  extern get_completion_domain_t<_CPO> const get_completion_domain;
-  extern get_domain_t const                  get_domain;
+  extern get_completion_domain_t<_CPO> const  get_completion_domain;
+  extern get_domain_t const                   get_domain;
+  extern get_await_completion_adaptor_t const get_await_completion_adaptor;
 
   template <class _Env>
   concept __is_debug_env = __callable<__debug_env_t, _Env>;

include/stdexec/__detail/__queries.hpp

@@ -28,6 +28,25 @@ namespace STDEXEC
   //////////////////////////////////////////////////////////////////////////////////////////////////
   // [exec.queries]
 
+  // [exec.get.await.adapt], see https://eel.is/c++draft/exec#get.await.adapt
+  struct get_await_completion_adaptor_t : __query<get_await_completion_adaptor_t>
+  {
+    template <class _Env>
+    STDEXEC_ATTRIBUTE(always_inline, host, device)
+    static constexpr void __validate() noexcept
+    {
+      static_assert(STDEXEC::__nothrow_callable<get_await_completion_adaptor_t, _Env const &>);
+    }
+
+    STDEXEC_ATTRIBUTE(nodiscard, always_inline, host, device)
+    static consteval auto query(forwarding_query_t) noexcept -> bool
+    {
+      return true;
+    }
+  };
+
+  inline constexpr get_await_completion_adaptor_t get_await_completion_adaptor{};
+
   // NOT TO SPEC:
   struct __is_scheduler_affine_t
   {