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 "__continue_on.hpp" 24 #include "__cpo.hpp" 25 #include "__diagnostics.hpp" 26 #include "__domain.hpp" 27 #include "__env.hpp" 28 #include "__inline_scheduler.hpp" 29 #include "__meta.hpp" 30 #include "__schedulers.hpp" 31 #include "__sender_adaptor_closure.hpp" 32 #include "__sender_introspection.hpp" 33 #include "__senders_core.hpp" 34 #include "__transform_sender.hpp" 35 #include "__type_traits.hpp" 36 #include "__utility.hpp" 37 #include "__write_env.hpp" 38 39 namespace stdexec 40 { 41 ///////////////////////////////////////////////////////////////////////////// 42 // [execution.senders.adaptors.on] 43 namespace __on_v2 44 { 45 inline constexpr __mstring __on_context = 46 "In stdexec::on(Scheduler, Sender)..."_mstr; 47 inline constexpr __mstring __no_scheduler_diag = 48 "stdexec::on() requires a scheduler to transition back to."_mstr; 49 inline constexpr __mstring __no_scheduler_details = 50 "The provided environment lacks a value for the get_scheduler() query."_mstr; 51 52 template <__mstring _Context = __on_context, 53 __mstring _Diagnostic = __no_scheduler_diag, 54 __mstring _Details = __no_scheduler_details> 55 struct _CANNOT_RESTORE_EXECUTION_CONTEXT_AFTER_ON_ 56 {}; 57 58 struct on_t; 59 60 template <class _Sender, class _Env> 61 struct __no_scheduler_in_environment 62 { 63 using sender_concept = sender_t; 64 get_completion_signaturesstdexec::__on_v2::__no_scheduler_in_environment65 static auto get_completion_signatures(const __no_scheduler_in_environment&, 66 const auto&) noexcept 67 { 68 return __mexception<_CANNOT_RESTORE_EXECUTION_CONTEXT_AFTER_ON_<>, 69 _WITH_SENDER_<_Sender>, _WITH_ENVIRONMENT_<_Env>>{}; 70 } 71 }; 72 73 template <class _Scheduler, class _Closure> 74 struct __continue_on_data 75 { 76 _Scheduler __sched_; 77 _Closure __clsur_; 78 }; 79 template <class _Scheduler, class _Closure> 80 __continue_on_data(_Scheduler, 81 _Closure) -> __continue_on_data<_Scheduler, _Closure>; 82 83 template <class _Scheduler> 84 struct __with_sched 85 { 86 using __t = __with_sched; 87 using __id = __with_sched; 88 89 _Scheduler __sched_; 90 querystdexec::__on_v2::__with_sched91 auto query(get_scheduler_t) const noexcept -> _Scheduler 92 { 93 return __sched_; 94 } 95 querystdexec::__on_v2::__with_sched96 auto query(get_domain_t) const noexcept 97 { 98 return query_or(get_domain, __sched_, default_domain()); 99 } 100 }; 101 102 template <class _Scheduler> 103 __with_sched(_Scheduler) -> __with_sched<_Scheduler>; 104 105 //////////////////////////////////////////////////////////////////////////////////////////////// 106 struct on_t 107 { 108 template <scheduler _Scheduler, sender _Sender> operator ()stdexec::__on_v2::on_t109 auto operator()(_Scheduler&& __sched, 110 _Sender&& __sndr) const -> __well_formed_sender auto 111 { 112 auto __domain = __get_early_domain(__sndr); 113 return stdexec::transform_sender( 114 __domain, __make_sexpr<on_t>(static_cast<_Scheduler&&>(__sched), 115 static_cast<_Sender&&>(__sndr))); 116 } 117 118 template <sender _Sender, scheduler _Scheduler, 119 __sender_adaptor_closure_for<_Sender> _Closure> operator ()stdexec::__on_v2::on_t120 auto operator()(_Sender&& __sndr, _Scheduler&& __sched, 121 _Closure&& __clsur) const -> __well_formed_sender auto 122 { 123 auto __domain = __get_early_domain(__sndr); 124 return stdexec::transform_sender( 125 __domain, __make_sexpr<on_t>( 126 __continue_on_data{static_cast<_Scheduler&&>(__sched), 127 static_cast<_Closure&&>(__clsur)}, 128 static_cast<_Sender&&>(__sndr))); 129 } 130 131 template <scheduler _Scheduler, __sender_adaptor_closure _Closure> 132 STDEXEC_ATTRIBUTE((always_inline)) operator ()stdexec::__on_v2::on_t133 auto operator()(_Scheduler&& __sched, _Closure&& __clsur) const 134 { 135 return __binder_back<on_t, __decay_t<_Scheduler>, __decay_t<_Closure>>{ 136 {{static_cast<_Scheduler&&>(__sched)}, 137 {static_cast<_Closure&&>(__clsur)}}, 138 {}, 139 {}}; 140 } 141 142 template <class _Env> 143 STDEXEC_ATTRIBUTE((always_inline)) __transform_env_fnstdexec::__on_v2::on_t144 static auto __transform_env_fn(_Env&& __env) noexcept 145 { 146 return [&]<class _Data>(__ignore, _Data&& __data, 147 __ignore) noexcept -> decltype(auto) { 148 if constexpr (scheduler<_Data>) 149 { 150 return __detail::__mkenv_sched(static_cast<_Env&&>(__env), 151 static_cast<_Data&&>(__data)); 152 } 153 else 154 { 155 return static_cast<_Env>(static_cast<_Env&&>(__env)); 156 } 157 }; 158 } 159 160 template <class _Env> 161 STDEXEC_ATTRIBUTE((always_inline)) __transform_sender_fnstdexec::__on_v2::on_t162 static auto __transform_sender_fn(const _Env& __env) noexcept 163 { 164 return [&]<class _Data, class _Child>(__ignore, _Data&& __data, 165 _Child&& __child) { 166 if constexpr (scheduler<_Data>) 167 { 168 // This branch handles the case where `on` was called like 169 // `on(sch, snd)` 170 auto __old = query_or(get_scheduler, __env, __none_such{}); 171 if constexpr (__same_as<decltype(__old), __none_such>) 172 { 173 if constexpr (__is_root_env<_Env>) 174 { 175 return continue_on( 176 start_on(static_cast<_Data&&>(__data), 177 static_cast<_Child&&>(__child)), 178 __inln::__scheduler{}); 179 } 180 else 181 { 182 return __none_such{}; 183 } 184 } 185 else 186 { 187 return continue_on(start_on(static_cast<_Data&&>(__data), 188 static_cast<_Child&&>(__child)), 189 static_cast<decltype(__old)&&>(__old)); 190 } 191 } 192 else 193 { 194 // This branch handles the case where `on` was called like 195 // `on(snd, sch, clsur)` 196 auto __old = query_or( 197 get_completion_scheduler<set_value_t>, get_env(__child), 198 query_or(get_scheduler, __env, __none_such{})); 199 if constexpr (__same_as<decltype(__old), __none_such>) 200 { 201 return __none_such{}; 202 } 203 else 204 { 205 auto&& [__sched, __clsur] = static_cast<_Data&&>(__data); 206 return __write_env( // 207 continue_on( // 208 __forward_like<_Data>(__clsur)( // 209 continue_on( // 210 __write_env(static_cast<_Child&&>(__child), 211 __with_sched{__old}), // 212 __sched)), // 213 __old), 214 __with_sched{__sched}); 215 } 216 } 217 }; 218 } 219 220 template <class _Sender, class _Env> 221 STDEXEC_ATTRIBUTE((always_inline)) transform_envstdexec::__on_v2::on_t222 static auto transform_env(const _Sender& __sndr, _Env&& __env) noexcept 223 { 224 return __sexpr_apply(__sndr, 225 __transform_env_fn(static_cast<_Env&&>(__env))); 226 } 227 228 template <class _Sender, class _Env> 229 STDEXEC_ATTRIBUTE((always_inline)) transform_senderstdexec::__on_v2::on_t230 static auto transform_sender(_Sender&& __sndr, const _Env& __env) 231 { 232 auto __tfx_sndr_fn = __transform_sender_fn(__env); 233 using _TfxSndrFn = decltype(__tfx_sndr_fn); 234 using _NewSndr = __sexpr_apply_result_t<_Sender, _TfxSndrFn>; 235 if constexpr (__same_as<_NewSndr, __none_such>) 236 { 237 return __no_scheduler_in_environment<_Sender, _Env>{}; 238 } 239 else 240 { 241 return __sexpr_apply(static_cast<_Sender&&>(__sndr), 242 static_cast<_TfxSndrFn&&>(__tfx_sndr_fn)); 243 } 244 } 245 }; 246 } // namespace __on_v2 247 248 namespace v2 249 { 250 using __on_v2::on_t; 251 inline constexpr on_t on{}; 252 253 using continue_on_t = v2::on_t; 254 inline constexpr continue_on_t continue_on{}; // for back-compat 255 } // namespace v2 256 257 template <> 258 struct __sexpr_impl<v2::on_t> : __sexpr_defaults 259 { 260 static constexpr auto get_completion_signatures = // 261 []<class _Sender>(_Sender&&) noexcept // 262 -> __merror_or_t< // 263 __completion_signatures_of_t< // 264 transform_sender_result_t<default_domain, _Sender, empty_env>>, 265 dependent_completions> { return {}; }; 266 }; 267 } // namespace stdexec 268