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