Merge pull request NVIDIA#1890 from ericniebler/allocator-aware-forward

implement _`allocator-aware-forward`_ from [exec.snd.expos]

Author: ericniebler

include/stdexec/__detail/__associate.hpp

@@ -185,7 +185,7 @@ namespace STDEXEC
       };
 
       explicit __op_state(std::pair<__assoc_t, __sender_ref_t> __parts, _Receiver&& __rcvr)
-        : __assoc_(std::move(__parts.first))
+        : __assoc_(STDEXEC::__allocator_aware_forward(std::move(__parts.first), __rcvr))
       {
         if (__assoc_)
         {
@@ -264,7 +264,7 @@ namespace STDEXEC
         auto& [__tag, __data] = __self;
 
         using op_state_t = __op_state<std::remove_cvref_t<_Self>, _Receiver>;
-        return op_state_t{__forward_like<_Self>(__data), std::move(__rcvr)};
+        return op_state_t{STDEXEC::__forward_like<_Self>(__data), std::move(__rcvr)};
       };
 
       static constexpr auto __start = [](auto& __state) noexcept -> void

include/stdexec/__detail/__basic_sender.hpp

@@ -22,6 +22,7 @@
 #include "__connect.hpp"
 #include "__diagnostics.hpp"
 #include "__env.hpp"
+#include "__memory.hpp"
 #include "__meta.hpp"
 #include "__operation_states.hpp"
 #include "__receivers.hpp"
@@ -104,6 +105,14 @@ namespace STDEXEC
       using __receiver_t = _Receiver;
       using __data_t     = _Data;
 
+      template <class _CvData>
+      STDEXEC_ATTRIBUTE(host, device)
+      constexpr __state(_Receiver __rcvr, _CvData&& __data)
+        noexcept(__nothrow_decay_copyable<_CvData>)
+        : __rcvr_(static_cast<_Receiver&&>(__rcvr))
+        , __data_(STDEXEC::__allocator_aware_forward(static_cast<_CvData&&>(__data), __rcvr_))
+      { }
+
       STDEXEC_IMMOVABLE_NO_UNIQUE_ADDRESS
       _Receiver __rcvr_;
       STDEXEC_IMMOVABLE_NO_UNIQUE_ADDRESS
@@ -160,17 +169,17 @@ namespace STDEXEC
                        STDEXEC::__get<1>(static_cast<_Sender&&>(__sndr))};
       };
 
-      static constexpr auto __get_env =  //
-        []<class _State>(__ignore,
-                         _State const & __state) noexcept -> env_of_t<decltype(_State::__rcvr_)>
+      static constexpr auto __get_env =                              //
+        []<class _State>(__ignore, _State const & __state) noexcept  //
+        -> env_of_t<decltype(_State::__rcvr_)>
       {
         return STDEXEC::get_env(__state.__rcvr_);
       };
 
       static constexpr auto __connect =  //
         []<class _Receiver, __connectable_to<_Receiver> _Sender>(_Sender&&   __sndr,
-                                                                 _Receiver&& __rcvr)  //
-        noexcept(__nothrow_constructible_from<__opstate<_Sender, _Receiver>, _Sender, _Receiver>)
+                                                                 _Receiver&& __rcvr) noexcept(  //
+          __nothrow_constructible_from<__opstate<_Sender, _Receiver>, _Sender, _Receiver>)
         -> __opstate<_Sender, _Receiver>
       {
         return __opstate<_Sender, _Receiver>(static_cast<_Sender&&>(__sndr),

include/stdexec/__detail/__let.hpp

@@ -295,7 +295,7 @@ namespace STDEXEC
                                         _Fun           __fn,
                                         _Receiver&&    __rcvr) noexcept
         : __rcvr_(static_cast<_Receiver&&>(__rcvr))
-        , __fn_(static_cast<_Fun&&>(__fn))
+        , __fn_(STDEXEC::__allocator_aware_forward(static_cast<_Fun&&>(__fn), __rcvr_))
         // TODO(ericniebler): this needs a fallback:
         , __env2_(__let::__mk_env2<_SetTag>(__attrs, STDEXEC::get_env(__rcvr_)))
       {}
@@ -641,7 +641,7 @@ namespace STDEXEC
     //! Implementation of the `let_*_t` types, where `_SetTag` is, e.g., `set_value_t` for `let_value`.
     template <class _LetTag>
     struct __let_t
-    {  // NOLINT(bugprone-crtp-constructor-accessibility)
+    {
       using __t = decltype(__set_tag_from_let_v<_LetTag>());
 
       template <sender _Sender, __movable_value _Fun>
@@ -656,6 +656,10 @@ namespace STDEXEC
       {
         return __closure(*this, static_cast<_Fun&&>(__fn));
       }
+
+     private:
+      friend _LetTag;
+      __let_t() = default;
     };
 
     template <class _LetTag>
@@ -722,11 +726,17 @@ namespace STDEXEC
   }  // namespace __let
 
   struct let_value_t : __let::__let_t<let_value_t>
-  {};
+  {
+    let_value_t() = default;
+  };
   struct let_error_t : __let::__let_t<let_error_t>
-  {};
+  {
+    let_error_t() = default;
+  };
   struct let_stopped_t : __let::__let_t<let_stopped_t>
-  {};
+  {
+    let_stopped_t() = default;
+  };
 
   inline constexpr let_value_t   let_value{};
   inline constexpr let_error_t   let_error{};

include/stdexec/__detail/__memory.hpp

@@ -17,6 +17,7 @@
 
 #include "__execution_fwd.hpp"
 #include "__scope.hpp"
+#include "__tuple.hpp"
 
 // include these after __execution_fwd.hpp
 #include <memory>  // IWYU pragma: export
@@ -87,4 +88,39 @@ namespace STDEXEC
     requires __same_as<_Ty, typename _Alloc::value_type>
   [[nodiscard]]
   constexpr auto __rebind_allocator(_Alloc const &&) noexcept = delete;
+
+  /////////////////////////////////////////////////////////////////////////////////////////
+  // __allocator_aware_forward: https://eel.is/c++draft/exec#snd.expos-49
+  template <class _Alloc>
+  [[nodiscard]]
+  constexpr auto __mk_obj_using_alloc_fn(_Alloc const &__alloc) noexcept
+  {
+    return [&__alloc]<class... _Args>(_Args &&...__args)
+    {
+      return __tuple{
+        std::make_obj_using_allocator<__decay_t<_Args>>(__alloc, static_cast<_Args &&>(__args))...};
+    };
+  }
+
+  template <class _Ty, class _Context>
+  [[nodiscard]]
+  constexpr auto __allocator_aware_forward(_Ty &&__obj, _Context const &__ctx) -> decltype(auto)
+  {
+    using __value_t = __decay_t<_Ty>;
+    if constexpr (!__callable<get_allocator_t, env_of_t<_Context>>)
+    {
+      return static_cast<_Ty &&>(__obj);
+    }
+    else if constexpr (__is_instance_of<__value_t, __tuple>)
+    {
+      auto const __alloc = get_allocator(get_env(__ctx));
+      return __apply(STDEXEC::__mk_obj_using_alloc_fn(__alloc), static_cast<_Ty &&>(__obj));
+    }
+    else
+    {
+      auto const __alloc = get_allocator(get_env(__ctx));
+      return std::make_obj_using_allocator<__decay_t<_Ty>>(__alloc, static_cast<_Ty &&>(__obj));
+    }
+  }
+
 }  // namespace STDEXEC

include/stdexec/__detail/__stop_when.hpp

@@ -184,7 +184,7 @@ namespace STDEXEC
         []<class _Self, class _Receiver>(_Self&& __self, _Receiver __rcvr) noexcept
       {
         auto& [__tag, __token, __child] = __self;
-        auto __new_token                = __make_token_fn{}(__forward_like<_Self>(__token),
+        auto __new_token                = __make_token_fn{}(STDEXEC::__forward_like<_Self>(__token),
                                              get_stop_token(STDEXEC::get_env(__rcvr)));
         return __state{std::move(__new_token), std::move(__rcvr)};
       };

include/stdexec/__detail/__utility.hpp

@@ -197,10 +197,11 @@ namespace STDEXEC
   }
 
   template <class _Ty>
-  constexpr _Ty const & __clamp(_Ty const & v, _Ty const & lo, _Ty const & hi)
+  constexpr _Ty const & __clamp(_Ty const & __val, _Ty const & __low, _Ty const & __high)
   {
-    STDEXEC_ASSERT(!(hi < lo));
-    return v < lo ? lo : hi < v ? hi : v;  // NOLINT(bugprone-return-const-ref-from-parameter)
+    STDEXEC_ASSERT(!(__high < __low));
+    // NOLINTNEXTLINE(bugprone-return-const-ref-from-parameter)
+    return __val < __low ? __low : __high < __val ? __high : __val;
   }
 
   STDEXEC_PRAGMA_PUSH()