1 /* 2 * Copyright (c) 2021-2024 NVIDIA Corporation 3 * 4 * Licensed under the Apache License Version 2.0 with LLVM Exceptions 5 * (the "License"); you may not use this file except in compliance with 6 * the License. You may obtain a copy of the License at 7 * 8 * https://llvm.org/LICENSE.txt 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 #pragma once 17 18 #include "__execution_fwd.hpp" 19 20 // include these after __execution_fwd.hpp 21 #include "__basic_sender.hpp" 22 #include "__concepts.hpp" 23 #include "__continues_on.hpp" 24 #include "__diagnostics.hpp" 25 #include "__domain.hpp" 26 #include "__env.hpp" 27 #include "__inline_scheduler.hpp" 28 #include "__meta.hpp" 29 #include "__schedulers.hpp" 30 #include "__senders_core.hpp" 31 #include "__sender_adaptor_closure.hpp" 32 #include "__sender_introspection.hpp" 33 #include "__transform_sender.hpp" 34 #include "__type_traits.hpp" 35 #include "__utility.hpp" 36 #include "__write_env.hpp" 37 38 namespace stdexec { 39 ///////////////////////////////////////////////////////////////////////////// 40 // [execution.senders.adaptors.on] 41 namespace __on { 42 inline constexpr __mstring __on_context = "In stdexec::on(Scheduler, Sender)..."_mstr; 43 inline constexpr __mstring __no_scheduler_diag = 44 "stdexec::on() requires a scheduler to transition back to."_mstr; 45 inline constexpr __mstring __no_scheduler_details = 46 "The provided environment lacks a value for the get_scheduler() query."_mstr; 47 48 template < 49 __mstring _Context = __on_context, 50 __mstring _Diagnostic = __no_scheduler_diag, 51 __mstring _Details = __no_scheduler_details 52 > 53 struct _CANNOT_RESTORE_EXECUTION_CONTEXT_AFTER_ON_ { }; 54 55 struct on_t; 56 57 template <class _Sender, class _Env> 58 struct __no_scheduler_in_environment { 59 using sender_concept = sender_t; 60 61 static auto get_completion_signaturesstdexec::__on::__no_scheduler_in_environment62 get_completion_signatures(const __no_scheduler_in_environment&, const auto&) noexcept { 63 return __mexception< 64 _CANNOT_RESTORE_EXECUTION_CONTEXT_AFTER_ON_<>, 65 _WITH_SENDER_<_Sender>, 66 _WITH_ENVIRONMENT_<_Env> 67 >{}; 68 } 69 }; 70 71 template <class _Scheduler, class _Closure> 72 struct __on_data { 73 _Scheduler __sched_; 74 _Closure __clsur_; 75 }; 76 template <class _Scheduler, class _Closure> 77 __on_data(_Scheduler, _Closure) -> __on_data<_Scheduler, _Closure>; 78 79 template <class _Scheduler> 80 struct __with_sched { 81 using __t = __with_sched; 82 using __id = __with_sched; 83 84 _Scheduler __sched_; 85 querystdexec::__on::__with_sched86 auto query(get_scheduler_t) const noexcept -> _Scheduler { 87 return __sched_; 88 } 89 querystdexec::__on::__with_sched90 auto query(get_domain_t) const noexcept { 91 return query_or(get_domain, __sched_, default_domain()); 92 } 93 }; 94 95 template <class _Scheduler> 96 __with_sched(_Scheduler) -> __with_sched<_Scheduler>; 97 98 //////////////////////////////////////////////////////////////////////////////////////////////// 99 struct on_t { 100 template <scheduler _Scheduler, sender _Sender> operator ()stdexec::__on::on_t101 auto operator()(_Scheduler&& __sched, _Sender&& __sndr) const -> __well_formed_sender auto { 102 auto __domain = __get_early_domain(__sndr); 103 return stdexec::transform_sender( 104 __domain, 105 __make_sexpr<on_t>(static_cast<_Scheduler&&>(__sched), static_cast<_Sender&&>(__sndr))); 106 } 107 108 template <sender _Sender, scheduler _Scheduler, __sender_adaptor_closure_for<_Sender> _Closure> operator ()stdexec::__on::on_t109 auto operator()(_Sender&& __sndr, _Scheduler&& __sched, _Closure&& __clsur) const 110 -> __well_formed_sender auto { 111 auto __domain = __get_early_domain(__sndr); 112 return stdexec::transform_sender( 113 __domain, 114 __make_sexpr<on_t>( 115 __on_data{static_cast<_Scheduler&&>(__sched), static_cast<_Closure&&>(__clsur)}, 116 static_cast<_Sender&&>(__sndr))); 117 } 118 119 template <scheduler _Scheduler, __sender_adaptor_closure _Closure> STDEXEC_ATTRIBUTEstdexec::__on::on_t120 STDEXEC_ATTRIBUTE(always_inline) 121 auto operator()(_Scheduler&& __sched, _Closure&& __clsur) const { 122 return __binder_back<on_t, __decay_t<_Scheduler>, __decay_t<_Closure>>{ 123 {{static_cast<_Scheduler&&>(__sched)}, {static_cast<_Closure&&>(__clsur)}}, 124 {}, 125 {} 126 }; 127 } 128 129 template <class _Env> STDEXEC_ATTRIBUTEstdexec::__on::on_t130 STDEXEC_ATTRIBUTE(always_inline) 131 static auto __transform_env_fn(_Env&& __env) noexcept { 132 return [&]<class _Data>(__ignore, _Data&& __data, __ignore) noexcept -> decltype(auto) { 133 if constexpr (scheduler<_Data>) { 134 return __env::__join( 135 __sched_env{static_cast<_Data&&>(__data)}, static_cast<_Env&&>(__env)); 136 } else { 137 return static_cast<_Env>(static_cast<_Env&&>(__env)); 138 } 139 }; 140 } 141 142 template <class _Env> STDEXEC_ATTRIBUTEstdexec::__on::on_t143 STDEXEC_ATTRIBUTE(always_inline) 144 static auto __transform_sender_fn(const _Env& __env) noexcept { 145 return [&]<class _Data, class _Child>(__ignore, _Data&& __data, _Child&& __child) { 146 // If __is_root_env<_Env> is true, then this sender has no parent, so there is 147 // no need to restore the execution context. We can use the inline scheduler 148 // as the scheduler if __env does not have one. 149 using __default_t = __if_c<__is_root_env<_Env>, inline_scheduler, __none_such>; 150 151 // If scheduler<_Data> is true, then this sender was created with `on(sch, sndr)`. 152 // In that case, the child sender is not a predecessor, so its completion 153 // scheduler is not the one we want to restore. If scheduler<_Data> is false, 154 // then this sender was created with `sndr | on(sch, clsur)`. The child sender 155 // *is* a predecessor, so we can use its completion scheduler to restore the 156 // execution context. 157 using __query_t = 158 __if_c<scheduler<_Data>, __none_such, get_completion_scheduler_t<set_value_t>>; 159 160 // Fetch the scheduler on which this operation will be started, and to which 161 // execution should be restored: 162 auto __old = 163 query_or(__query_t{}, get_env(__child), query_or(get_scheduler, __env, __default_t{})); 164 165 if constexpr (__same_as<decltype(__old), __none_such>) { 166 return __none_such{}; 167 } else if constexpr (scheduler<_Data>) { 168 // This branch handles the case where `on` was called like `on(sch, sndr)` 169 return continues_on( 170 starts_on(static_cast<_Data&&>(__data), static_cast<_Child&&>(__child)), 171 static_cast<decltype(__old)&&>(__old)); 172 } else { 173 // This branch handles the case where `on` was called like `sndr | on(sch, clsur)` 174 auto&& [__sched, __clsur] = static_cast<_Data&&>(__data); 175 return write_env( 176 continues_on( 177 __forward_like<_Data>(__clsur)(continues_on( 178 write_env(static_cast<_Child&&>(__child), __with_sched{__old}), __sched)), 179 __old), 180 __with_sched{__sched}); 181 } 182 }; 183 } 184 185 template <class _Sender, class _Env> STDEXEC_ATTRIBUTEstdexec::__on::on_t186 STDEXEC_ATTRIBUTE(always_inline) 187 static auto transform_env(const _Sender& __sndr, _Env&& __env) noexcept { 188 return __sexpr_apply(__sndr, __transform_env_fn(static_cast<_Env&&>(__env))); 189 } 190 191 template <class _Sender, class _Env> STDEXEC_ATTRIBUTEstdexec::__on::on_t192 STDEXEC_ATTRIBUTE(always_inline) 193 static auto transform_sender(_Sender&& __sndr, const _Env& __env) { 194 auto __tfx_sndr_fn = __transform_sender_fn(__env); 195 using _TfxSndrFn = decltype(__tfx_sndr_fn); 196 using _NewSndr = __sexpr_apply_result_t<_Sender, _TfxSndrFn>; 197 if constexpr (__same_as<_NewSndr, __none_such>) { 198 return __no_scheduler_in_environment<_Sender, _Env>{}; 199 } else { 200 return __sexpr_apply( 201 static_cast<_Sender&&>(__sndr), static_cast<_TfxSndrFn&&>(__tfx_sndr_fn)); 202 } 203 } 204 }; 205 } // namespace __on 206 207 using __on::on_t; 208 inline constexpr on_t on{}; 209 210 namespace v2 { 211 using on_t [[deprecated("use stdexec::on_t instead")]] = stdexec::on_t; 212 [[deprecated("use stdexec::on instead")]] 213 inline constexpr stdexec::on_t const & on = stdexec::on; 214 } // namespace v2 215 216 template <> 217 struct __sexpr_impl<on_t> : __sexpr_defaults { 218 static constexpr auto get_completion_signatures = []<class _Sender>(_Sender&&) noexcept 219 -> __merror_or_t< 220 __completion_signatures_of_t<transform_sender_result_t<default_domain, _Sender, env<>>>, 221 dependent_completions 222 > { 223 return {}; 224 }; 225 }; 226 } // namespace stdexec 227