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 "__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 __on_data 75 { 76 _Scheduler __sched_; 77 _Closure __clsur_; 78 }; 79 template <class _Scheduler, class _Closure> 80 __on_data(_Scheduler, _Closure) -> __on_data<_Scheduler, _Closure>; 81 82 template <class _Scheduler> 83 struct __with_sched 84 { 85 using __t = __with_sched; 86 using __id = __with_sched; 87 88 _Scheduler __sched_; 89 querystdexec::__on_v2::__with_sched90 auto query(get_scheduler_t) const noexcept -> _Scheduler 91 { 92 return __sched_; 93 } 94 querystdexec::__on_v2::__with_sched95 auto query(get_domain_t) const noexcept 96 { 97 return query_or(get_domain, __sched_, default_domain()); 98 } 99 }; 100 101 template <class _Scheduler> 102 __with_sched(_Scheduler) -> __with_sched<_Scheduler>; 103 104 //////////////////////////////////////////////////////////////////////////////////////////////// 105 struct on_t 106 { 107 template <scheduler _Scheduler, sender _Sender> operator ()stdexec::__on_v2::on_t108 auto operator()(_Scheduler&& __sched, _Sender&& __sndr) const 109 -> __well_formed_sender auto 110 { 111 auto __domain = __get_early_domain(__sndr); 112 return stdexec::transform_sender( 113 __domain, __make_sexpr<on_t>(static_cast<_Scheduler&&>(__sched), 114 static_cast<_Sender&&>(__sndr))); 115 } 116 117 template <sender _Sender, scheduler _Scheduler, 118 __sender_adaptor_closure_for<_Sender> _Closure> operator ()stdexec::__on_v2::on_t119 auto operator()(_Sender&& __sndr, _Scheduler&& __sched, 120 _Closure&& __clsur) const -> __well_formed_sender auto 121 { 122 auto __domain = __get_early_domain(__sndr); 123 return stdexec::transform_sender( 124 __domain, 125 __make_sexpr<on_t>(__on_data{static_cast<_Scheduler&&>(__sched), 126 static_cast<_Closure&&>(__clsur)}, 127 static_cast<_Sender&&>(__sndr))); 128 } 129 130 template <scheduler _Scheduler, __sender_adaptor_closure _Closure> 131 STDEXEC_ATTRIBUTE((always_inline)) operator ()stdexec::__on_v2::on_t132 auto operator()(_Scheduler&& __sched, _Closure&& __clsur) const 133 { 134 return __binder_back<on_t, __decay_t<_Scheduler>, __decay_t<_Closure>>{ 135 {{static_cast<_Scheduler&&>(__sched)}, 136 {static_cast<_Closure&&>(__clsur)}}, 137 {}, 138 {}}; 139 } 140 141 template <class _Env> 142 STDEXEC_ATTRIBUTE((always_inline)) __transform_env_fnstdexec::__on_v2::on_t143 static auto __transform_env_fn(_Env&& __env) noexcept 144 { 145 return [&]<class _Data>(__ignore, _Data&& __data, 146 __ignore) noexcept -> decltype(auto) { 147 if constexpr (scheduler<_Data>) 148 { 149 return __detail::__mkenv_sched(static_cast<_Env&&>(__env), 150 static_cast<_Data&&>(__data)); 151 } 152 else 153 { 154 return static_cast<_Env>(static_cast<_Env&&>(__env)); 155 } 156 }; 157 } 158 159 template <class _Env> 160 STDEXEC_ATTRIBUTE((always_inline)) __transform_sender_fnstdexec::__on_v2::on_t161 static auto __transform_sender_fn(const _Env& __env) noexcept 162 { 163 return [&]<class _Data, class _Child>(__ignore, _Data&& __data, 164 _Child&& __child) { 165 if constexpr (scheduler<_Data>) 166 { 167 // This branch handles the case where `on` was called like 168 // `on(sch, snd)` 169 auto __old = query_or(get_scheduler, __env, __none_such{}); 170 if constexpr (__same_as<decltype(__old), __none_such>) 171 { 172 if constexpr (__is_root_env<_Env>) 173 { 174 return continues_on( 175 starts_on(static_cast<_Data&&>(__data), 176 static_cast<_Child&&>(__child)), 177 __inln::__scheduler{}); 178 } 179 else 180 { 181 return __none_such{}; 182 } 183 } 184 else 185 { 186 return continues_on( 187 starts_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 continues_on( // 208 __forward_like<_Data>(__clsur)( // 209 continues_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