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 "__concepts.hpp"
22 #include "__cpo.hpp"
23 #include "__env.hpp"
24 #include "__senders.hpp"
25 #include "__tag_invoke.hpp"
26 
27 namespace stdexec
28 {
29 /////////////////////////////////////////////////////////////////////////////
30 // [execution.senders.schedule]
31 namespace __sched
32 {
33 struct schedule_t
34 {
35     template <__same_as<schedule_t> _Self, class _Scheduler>
36     STDEXEC_ATTRIBUTE((host, device, always_inline))
tag_invoke(_Self,_Scheduler && __sched)37     friend auto tag_invoke(_Self, _Scheduler&& __sched) //
38         noexcept(noexcept(static_cast<_Scheduler&&>(__sched).schedule()))
39             -> decltype(static_cast<_Scheduler&&>(__sched).schedule())
40     {
41         static_assert(
42             sender<decltype(static_cast<_Scheduler&&>(__sched).schedule())>,
43             "schedule() member functions must return a sender");
44         return static_cast<_Scheduler&&>(__sched).schedule();
45     }
46 
47     template <class _Scheduler>
48         requires tag_invocable<schedule_t, _Scheduler>
49     STDEXEC_ATTRIBUTE((host, device))
operator ()stdexec::__sched::schedule_t50     auto operator()(_Scheduler&& __sched) const
51         noexcept(nothrow_tag_invocable<schedule_t, _Scheduler>)
52     {
53         static_assert(sender<tag_invoke_result_t<schedule_t, _Scheduler>>);
54         return tag_invoke(schedule_t{}, static_cast<_Scheduler&&>(__sched));
55     }
56 
querystdexec::__sched::schedule_t57     static constexpr auto query(forwarding_query_t) noexcept -> bool
58     {
59         return false;
60     }
61 };
62 } // namespace __sched
63 
64 using __sched::schedule_t;
65 inline constexpr schedule_t schedule{};
66 
67 template <class _Scheduler>
68 concept __has_schedule = //
69     requires(_Scheduler&& __sched) {
70         { schedule(static_cast<_Scheduler&&>(__sched)) } -> sender;
71     };
72 
73 template <class _Scheduler>
74 concept __sender_has_completion_scheduler =
75     requires(_Scheduler&& __sched) {
76         {
77             stdexec::__decay_copy(get_completion_scheduler<set_value_t>(
78                 get_env(schedule(static_cast<_Scheduler&&>(__sched)))))
79         } -> same_as<__decay_t<_Scheduler>>;
80     };
81 
82 template <class _Scheduler>
83 concept scheduler =                                  //
84     __has_schedule<_Scheduler>                       //
85     && __sender_has_completion_scheduler<_Scheduler> //
86     && equality_comparable<__decay_t<_Scheduler>>    //
87     && copy_constructible<__decay_t<_Scheduler>>;
88 
89 template <scheduler _Scheduler>
90 using schedule_result_t = __call_result_t<schedule_t, _Scheduler>;
91 
92 template <class _SchedulerProvider>
93 concept __scheduler_provider = //
94     requires(const _SchedulerProvider& __sp) {
95         { get_scheduler(__sp) } -> scheduler;
96     };
97 
98 namespace __queries
99 {
100 template <class _Env>
101     requires tag_invocable<get_scheduler_t, const _Env&>
operator ()(const _Env & __env) const102 inline auto get_scheduler_t::operator()(const _Env& __env) const noexcept
103     -> tag_invoke_result_t<get_scheduler_t, const _Env&>
104 {
105     static_assert(nothrow_tag_invocable<get_scheduler_t, const _Env&>);
106     static_assert(scheduler<tag_invoke_result_t<get_scheduler_t, const _Env&>>);
107     return tag_invoke(get_scheduler_t{}, __env);
108 }
109 
110 template <class _Env>
111     requires tag_invocable<get_delegatee_scheduler_t, const _Env&>
operator ()(const _Env & __env) const112 inline auto get_delegatee_scheduler_t::operator()(const _Env& __env) const
113     noexcept -> tag_invoke_result_t<get_delegatee_scheduler_t, const _Env&>
114 {
115     static_assert(
116         nothrow_tag_invocable<get_delegatee_scheduler_t, const _Env&>);
117     static_assert(
118         scheduler<tag_invoke_result_t<get_delegatee_scheduler_t, const _Env&>>);
119     return tag_invoke(get_delegatee_scheduler_t{}, __env);
120 }
121 
122 template <__completion_tag _Tag>
123 template <__has_completion_scheduler_for<_Tag> _Env>
operator ()(const _Env & __env) const124 auto get_completion_scheduler_t<_Tag>::operator()(
125     const _Env& __env) const noexcept
126     -> tag_invoke_result_t<get_completion_scheduler_t<_Tag>, const _Env&>
127 {
128     static_assert(
129         nothrow_tag_invocable<get_completion_scheduler_t<_Tag>, const _Env&>,
130         "get_completion_scheduler<_Tag> should be noexcept");
131     static_assert(
132         scheduler<tag_invoke_result_t<get_completion_scheduler_t<_Tag>,
133                                       const _Env&>>);
134     return tag_invoke(*this, __env);
135 }
136 } // namespace __queries
137 
138 namespace __detail
139 {
140 // A handy utility for augmenting an environment with a scheduler.
141 template <class _Env, class _Scheduler>
142 STDEXEC_ATTRIBUTE((always_inline))
__mkenv_sched(_Env && __env,_Scheduler __sched)143 auto __mkenv_sched(_Env&& __env, _Scheduler __sched)
144 {
145     auto __env2 =
146         __env::__join(prop{get_scheduler, __sched},
147                       __env::__without(static_cast<_Env&&>(__env), get_domain));
148     using _Env2 = decltype(__env2);
149 
150     struct __env_t : _Env2
151     {};
152 
153     return __env_t{static_cast<_Env2&&>(__env2)};
154 }
155 } // namespace __detail
156 } // namespace stdexec
157