1 /*
2  * Copyright (c) 2021-2023 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 "../functional.hpp"
19 #include "../stop_token.hpp"
20 #include "__concepts.hpp"
21 #include "__config.hpp"
22 #include "__execution_fwd.hpp"
23 #include "__meta.hpp"
24 
25 #include <concepts>
26 #include <exception>
27 #include <type_traits>
28 
29 STDEXEC_PRAGMA_PUSH()
30 STDEXEC_PRAGMA_IGNORE_EDG(probable_guiding_friend)
31 STDEXEC_PRAGMA_IGNORE_EDG(type_qualifiers_ignored_on_reference)
32 
33 namespace stdexec
34 {
35 // [exec.queries.queryable]
36 template <class T>
37 concept queryable = destructible<T>;
38 
39 template <class Tag>
40 struct __query
41 {
42     template <class Sig>
43     static inline constexpr Tag (*signature)(Sig) = nullptr;
44 };
45 
46 //////////////////////////////////////////////////////////////////////////////////////////////////
47 // [exec.queries]
48 namespace __queries
49 {
50 struct forwarding_query_t
51 {
52     template <class _Query>
53     constexpr auto operator()(_Query __query) const noexcept -> bool
54     {
55         if constexpr (tag_invocable<forwarding_query_t, _Query>)
56         {
57             return tag_invoke(*this, static_cast<_Query&&>(__query));
58         }
59         else if constexpr (std::derived_from<_Query, forwarding_query_t>)
60         {
61             return true;
62         }
63         else
64         {
65             return false;
66         }
67     }
68 };
69 
70 struct query_or_t
71 {
72     template <class _Query, class _Queryable, class _Default>
73     constexpr auto operator()(_Query, _Queryable&&, _Default&& __default) const
74         noexcept(__nothrow_constructible_from<_Default, _Default&&>) -> _Default
75     {
76         return static_cast<_Default&&>(__default);
77     }
78 
79     template <class _Query, class _Queryable, class _Default>
80         requires __callable<_Query, _Queryable>
81     constexpr auto operator()(_Query __query, _Queryable&& __queryable,
82                               _Default&&) const
83         noexcept(__nothrow_callable<_Query, _Queryable>)
84             -> __call_result_t<_Query, _Queryable>
85     {
86         return static_cast<_Query&&>(__query)(
87             static_cast<_Queryable&&>(__queryable));
88     }
89 };
90 
91 struct execute_may_block_caller_t : __query<execute_may_block_caller_t>
92 {
93     template <class _Tp>
94         requires tag_invocable<execute_may_block_caller_t, __cref_t<_Tp>>
95     constexpr auto operator()(_Tp&& __t) const noexcept -> bool
96     {
97         static_assert(
98             same_as<bool, tag_invoke_result_t<execute_may_block_caller_t,
99                                               __cref_t<_Tp>>>);
100         static_assert(
101             nothrow_tag_invocable<execute_may_block_caller_t, __cref_t<_Tp>>);
102         return tag_invoke(execute_may_block_caller_t{}, std::as_const(__t));
103     }
104 
105     constexpr auto operator()(auto&&) const noexcept -> bool
106     {
107         return true;
108     }
109 };
110 
111 struct get_forward_progress_guarantee_t :
112     __query<get_forward_progress_guarantee_t>
113 {
114     template <class _Tp>
115         requires tag_invocable<get_forward_progress_guarantee_t, __cref_t<_Tp>>
116     constexpr auto operator()(_Tp&& __t) const noexcept(
117         nothrow_tag_invocable<get_forward_progress_guarantee_t, __cref_t<_Tp>>)
118         -> tag_invoke_result_t<get_forward_progress_guarantee_t, __cref_t<_Tp>>
119     {
120         return tag_invoke(get_forward_progress_guarantee_t{},
121                           std::as_const(__t));
122     }
123 
124     constexpr auto operator()(auto&&) const noexcept
125         -> stdexec::forward_progress_guarantee
126     {
127         return stdexec::forward_progress_guarantee::weakly_parallel;
128     }
129 };
130 
131 struct __has_algorithm_customizations_t :
132     __query<__has_algorithm_customizations_t>
133 {
134     template <class _Tp>
135     using __result_t =
136         tag_invoke_result_t<__has_algorithm_customizations_t, __cref_t<_Tp>>;
137 
138     template <class _Tp>
139         requires tag_invocable<__has_algorithm_customizations_t, __cref_t<_Tp>>
140     constexpr auto operator()(_Tp&&) const noexcept(noexcept(__result_t<_Tp>{}))
141         -> __result_t<_Tp>
142     {
143         using _Boolean = tag_invoke_result_t<__has_algorithm_customizations_t,
144                                              __cref_t<_Tp>>;
145         static_assert(_Boolean{}
146                           ? true
147                           : false); // must be contextually convertible to bool
148         return _Boolean{};
149     }
150 
151     constexpr auto operator()(auto&&) const noexcept -> std::false_type
152     {
153         return {};
154     }
155 };
156 
157 // TODO: implement allocator concept
158 template <class _T0>
159 concept __allocator_c = true;
160 
161 struct get_scheduler_t : __query<get_scheduler_t>
162 {
163     friend constexpr auto tag_invoke(forwarding_query_t,
164                                      const get_scheduler_t&) noexcept -> bool
165     {
166         return true;
167     }
168 
169     template <class _Env>
170         requires tag_invocable<get_scheduler_t, const _Env&>
171     auto operator()(const _Env& __env) const noexcept
172         -> tag_invoke_result_t<get_scheduler_t, const _Env&>;
173 
174     template <class _Tag = get_scheduler_t>
175     auto operator()() const noexcept;
176 };
177 
178 struct get_delegatee_scheduler_t : __query<get_delegatee_scheduler_t>
179 {
180     friend constexpr auto tag_invoke(forwarding_query_t,
181                                      const get_delegatee_scheduler_t&) noexcept
182         -> bool
183     {
184         return true;
185     }
186 
187     template <class _Env>
188         requires tag_invocable<get_delegatee_scheduler_t, const _Env&>
189     auto operator()(const _Env& __t) const noexcept
190         -> tag_invoke_result_t<get_delegatee_scheduler_t, const _Env&>;
191 
192     template <class _Tag = get_delegatee_scheduler_t>
193     auto operator()() const noexcept;
194 };
195 
196 struct get_allocator_t : __query<get_allocator_t>
197 {
198     friend constexpr auto tag_invoke(forwarding_query_t,
199                                      const get_allocator_t&) noexcept -> bool
200     {
201         return true;
202     }
203 
204     template <class _Env>
205         requires tag_invocable<get_allocator_t, const _Env&>
206     auto operator()(const _Env& __env) const noexcept
207         -> tag_invoke_result_t<get_allocator_t, const _Env&>
208     {
209         static_assert(nothrow_tag_invocable<get_allocator_t, const _Env&>);
210         static_assert(
211             __allocator_c<tag_invoke_result_t<get_allocator_t, const _Env&>>);
212         return tag_invoke(get_allocator_t{}, __env);
213     }
214 
215     template <class _Tag = get_allocator_t>
216     auto operator()() const noexcept;
217 };
218 
219 struct get_stop_token_t : __query<get_stop_token_t>
220 {
221     friend constexpr auto tag_invoke(forwarding_query_t,
222                                      const get_stop_token_t&) noexcept -> bool
223     {
224         return true;
225     }
226 
227     template <class _Env>
228     auto operator()(const _Env&) const noexcept -> never_stop_token
229     {
230         return {};
231     }
232 
233     template <class _Env>
234         requires tag_invocable<get_stop_token_t, const _Env&>
235     auto operator()(const _Env& __env) const noexcept
236         -> tag_invoke_result_t<get_stop_token_t, const _Env&>
237     {
238         static_assert(nothrow_tag_invocable<get_stop_token_t, const _Env&>);
239         static_assert(stoppable_token<
240                       tag_invoke_result_t<get_stop_token_t, const _Env&>>);
241         return tag_invoke(get_stop_token_t{}, __env);
242     }
243 
244     template <class _Tag = get_stop_token_t>
245     auto operator()() const noexcept;
246 };
247 
248 template <class _Queryable, class _CPO>
249 concept __has_completion_scheduler_for =
250     queryable<_Queryable> && //
251     tag_invocable<get_completion_scheduler_t<_CPO>, const _Queryable&>;
252 
253 template <__completion_tag _CPO>
254 struct get_completion_scheduler_t : __query<get_completion_scheduler_t<_CPO>>
255 {
256     friend constexpr auto
257         tag_invoke(forwarding_query_t,
258                    const get_completion_scheduler_t<_CPO>&) noexcept -> bool
259     {
260         return true;
261     }
262 
263     template <__has_completion_scheduler_for<_CPO> _Queryable>
264     auto operator()(const _Queryable& __queryable) const noexcept
265         -> tag_invoke_result_t<get_completion_scheduler_t<_CPO>,
266                                const _Queryable&>;
267 };
268 
269 struct get_domain_t
270 {
271     template <class _Ty>
272         requires tag_invocable<get_domain_t, const _Ty&>
273     constexpr auto operator()(const _Ty& __ty) const noexcept
274         -> tag_invoke_result_t<get_domain_t, const _Ty&>
275     {
276         static_assert(nothrow_tag_invocable<get_domain_t, const _Ty&>,
277                       "Customizations of get_domain must be noexcept.");
278         static_assert(__class<tag_invoke_result_t<get_domain_t, const _Ty&>>,
279                       "Customizations of get_domain must return a class type.");
280         return tag_invoke(get_domain_t{}, __ty);
281     }
282 
283     friend constexpr auto tag_invoke(forwarding_query_t, get_domain_t) noexcept
284         -> bool
285     {
286         return true;
287     }
288 };
289 } // namespace __queries
290 
291 using __queries::__has_algorithm_customizations_t;
292 using __queries::execute_may_block_caller_t;
293 using __queries::forwarding_query_t;
294 using __queries::get_allocator_t;
295 using __queries::get_completion_scheduler_t;
296 using __queries::get_delegatee_scheduler_t;
297 using __queries::get_domain_t;
298 using __queries::get_forward_progress_guarantee_t;
299 using __queries::get_scheduler_t;
300 using __queries::get_stop_token_t;
301 using __queries::query_or_t;
302 
303 inline constexpr forwarding_query_t forwarding_query{};
304 inline constexpr query_or_t query_or{}; // NOT TO SPEC
305 inline constexpr execute_may_block_caller_t execute_may_block_caller{};
306 inline constexpr __has_algorithm_customizations_t
307     __has_algorithm_customizations{};
308 inline constexpr get_forward_progress_guarantee_t
309     get_forward_progress_guarantee{};
310 inline constexpr get_scheduler_t get_scheduler{};
311 inline constexpr get_delegatee_scheduler_t get_delegatee_scheduler{};
312 inline constexpr get_allocator_t get_allocator{};
313 inline constexpr get_stop_token_t get_stop_token{};
314 #if !STDEXEC_GCC() || defined(__OPTIMIZE_SIZE__)
315 template <__completion_tag _CPO>
316 inline constexpr get_completion_scheduler_t<_CPO> get_completion_scheduler{};
317 #else
318 template <>
319 inline constexpr get_completion_scheduler_t<set_value_t>
320     get_completion_scheduler<set_value_t>{};
321 template <>
322 inline constexpr get_completion_scheduler_t<set_error_t>
323     get_completion_scheduler<set_error_t>{};
324 template <>
325 inline constexpr get_completion_scheduler_t<set_stopped_t>
326     get_completion_scheduler<set_stopped_t>{};
327 #endif
328 
329 template <class _Tag>
330 concept __forwarding_query = forwarding_query(_Tag{});
331 
332 inline constexpr get_domain_t get_domain{};
333 
334 template <class _Tag, class _Queryable, class _Default>
335 using __query_result_or_t =
336     __call_result_t<query_or_t, _Tag, _Queryable, _Default>;
337 
338 /////////////////////////////////////////////////////////////////////////////
339 namespace __env
340 {
341 // For getting an execution environment from a receiver,
342 // or the attributes from a sender.
343 struct get_env_t
344 {
345     template <class _EnvProvider>
346         requires tag_invocable<get_env_t, const _EnvProvider&>
347     STDEXEC_ATTRIBUTE((always_inline)) constexpr auto
348         operator()(const _EnvProvider& __env_provider) const noexcept
349         -> tag_invoke_result_t<get_env_t, const _EnvProvider&>
350     {
351         static_assert(
352             queryable<tag_invoke_result_t<get_env_t, const _EnvProvider&>>);
353         static_assert(nothrow_tag_invocable<get_env_t, const _EnvProvider&>);
354         return tag_invoke(*this, __env_provider);
355     }
356 
357     template <class _EnvProvider>
358     constexpr auto operator()(const _EnvProvider&) const noexcept -> empty_env
359     {
360         return {};
361     }
362 };
363 
364 // To be kept in sync with the promise type used in __connect_awaitable
365 template <class _Env>
366 struct __promise
367 {
368     template <class _Ty>
369     auto await_transform(_Ty&& __value) noexcept -> _Ty&&
370     {
371         return static_cast<_Ty&&>(__value);
372     }
373 
374     template <class _Ty>
375         requires tag_invocable<as_awaitable_t, _Ty, __promise&>
376     auto await_transform(_Ty&& __value) //
377         noexcept(nothrow_tag_invocable<as_awaitable_t, _Ty, __promise&>)
378             -> tag_invoke_result_t<as_awaitable_t, _Ty, __promise&>
379     {
380         return tag_invoke(as_awaitable, static_cast<_Ty&&>(__value), *this);
381     }
382 
383     template <same_as<get_env_t> _Tag>
384     friend auto tag_invoke(_Tag, const __promise&) noexcept -> const _Env&
385     {
386         std::terminate();
387     }
388 };
389 
390 template <class _Value, class _Tag, class... _Tags>
391 struct __with
392 {
393     using __t = __with;
394     using __id = __with;
395     STDEXEC_ATTRIBUTE((no_unique_address))
396     _Value __value_;
397 
398     __with() = default;
399 
400     constexpr explicit __with(_Value __value) noexcept(
401         __nothrow_decay_copyable<_Value>) :
402         __value_(static_cast<_Value&&>(__value))
403     {}
404 
405     constexpr explicit __with(_Value __value, _Tag, _Tags...) noexcept(
406         __nothrow_decay_copyable<_Value>) :
407         __value_(static_cast<_Value&&>(__value))
408     {}
409 
410     template <__one_of<_Tag, _Tags...> _Key>
411     friend auto tag_invoke(_Key, const __with& __self) //
412         noexcept(__nothrow_decay_copyable<_Value>) -> _Value
413     {
414         return __self.__value_;
415     }
416 };
417 
418 template <class _Value, class _Tag, class... _Tags>
419 __with(_Value, _Tag, _Tags...) -> __with<_Value, _Tag, _Tags...>;
420 
421 template <class _Env>
422 struct __fwd
423 {
424     static_assert(__nothrow_move_constructible<_Env>);
425     using __t = __fwd;
426     using __id = __fwd;
427     STDEXEC_ATTRIBUTE((no_unique_address))
428     _Env __env_;
429 
430     template <__forwarding_query _Tag>
431         requires tag_invocable<_Tag, const _Env&>
432     friend auto tag_invoke(_Tag, const __fwd& __self) //
433         noexcept(nothrow_tag_invocable<_Tag, const _Env&>)
434             -> tag_invoke_result_t<_Tag, const _Env&>
435     {
436         return _Tag()(__self.__env_);
437     }
438 };
439 
440 template <class _Env>
441 __fwd(_Env&&) -> __fwd<_Env>;
442 
443 template <class _Env>
444 struct __ref
445 {
446     using __t = __ref;
447     using __id = __ref;
448     const _Env& __env_;
449 
450     template <class _Tag>
451         requires tag_invocable<_Tag, const _Env&>
452     friend auto tag_invoke(_Tag, const __ref& __self) //
453         noexcept(nothrow_tag_invocable<_Tag, const _Env&>)
454             -> tag_invoke_result_t<_Tag, const _Env&>
455     {
456         return _Tag()(__self.__env_);
457     }
458 };
459 
460 template <class _Env>
461 __ref(_Env&) -> __ref<_Env>;
462 
463 struct __ref_fn
464 {
465     template <class _Env>
466     constexpr auto operator()(_Env&& __env) const
467     {
468         if constexpr (same_as<_Env, _Env&>)
469         {
470             return __ref{static_cast<_Env&&>(__env)};
471         }
472         else
473         {
474             return static_cast<_Env&&>(__env);
475         }
476     }
477 };
478 
479 template <class _Env, class _Tag, class... _Tags>
480 struct __without_ : _Env
481 {
482     static_assert(__nothrow_move_constructible<_Env>);
483     using __t = __without_;
484     using __id = __without_;
485 
486     constexpr explicit __without_(_Env&& __env, _Tag, _Tags...) noexcept :
487         _Env(static_cast<_Env&&>(__env))
488     {}
489 
490     template <__one_of<_Tag, _Tags...> _Key, class _Self>
491         requires(std::is_base_of_v<__without_, __decay_t<_Self>>)
492     friend auto tag_invoke(_Key, _Self&&) noexcept = delete;
493 };
494 
495 struct __without_fn
496 {
497     template <class _Env, class _Tag, class... _Tags>
498     constexpr auto operator()(_Env&& __env, _Tag, _Tags...) const noexcept
499         -> decltype(auto)
500     {
501         if constexpr (tag_invocable<_Tag, _Env> ||
502                       (tag_invocable<_Tags, _Env> || ...))
503         {
504             return __without_{__ref_fn()(static_cast<_Env&&>(__env)), _Tag(),
505                               _Tags()...};
506         }
507         else
508         {
509             return static_cast<_Env>(static_cast<_Env&&>(__env));
510         }
511     }
512 };
513 
514 inline constexpr __without_fn __without{};
515 
516 template <class _Env, class _Tag, class... _Tags>
517 using __without_t = __result_of<__without, _Env, _Tag, _Tags...>;
518 
519 template <class _Second, class _First>
520 struct __joined : _Second
521 {
522     static_assert(__nothrow_move_constructible<_First>);
523     static_assert(__nothrow_move_constructible<_Second>);
524     using __t = __joined;
525     using __id = __joined;
526 
527     STDEXEC_ATTRIBUTE((no_unique_address))
528     _First __env_;
529 
530     template <class _Tag>
531         requires tag_invocable<_Tag, const _First&>
532     friend auto tag_invoke(_Tag, const __joined& __self) //
533         noexcept(nothrow_tag_invocable<_Tag, const _First&>)
534             -> tag_invoke_result_t<_Tag, const _First&>
535     {
536         return _Tag()(__self.__env_);
537     }
538 };
539 
540 template <class _Second, class _First>
541 __joined(_Second&&, _First&&) -> __joined<_Second, _First>;
542 
543 template <__nothrow_move_constructible _Fun>
544 struct __from
545 {
546     using __t = __from;
547     using __id = __from;
548     STDEXEC_ATTRIBUTE((no_unique_address))
549     _Fun __fun_;
550 
551     template <class _Tag>
552         requires __callable<const _Fun&, _Tag>
553     friend auto tag_invoke(_Tag, const __from& __self) //
554         noexcept(__nothrow_callable<const _Fun&, _Tag>)
555             -> __call_result_t<const _Fun&, _Tag>
556     {
557         return __self.__fun_(_Tag());
558     }
559 };
560 
561 template <class _Fun>
562 __from(_Fun) -> __from<_Fun>;
563 
564 struct __fwd_fn
565 {
566     template <class Env>
567     auto operator()(Env&& env) const
568     {
569         return __fwd{static_cast<Env&&>(env)};
570     }
571 
572     auto operator()(empty_env) const -> empty_env
573     {
574         return {};
575     }
576 };
577 
578 struct __join_fn
579 {
580     auto operator()() const -> empty_env
581     {
582         return {};
583     }
584 
585     template <class _Env>
586     auto operator()(_Env&& __env) const -> _Env
587     {
588         return static_cast<_Env&&>(__env);
589     }
590 
591     auto operator()(empty_env) const -> empty_env
592     {
593         return {};
594     }
595 
596     template <class _Env>
597     auto operator()(_Env&& __env, empty_env) const -> _Env
598     {
599         return static_cast<_Env&&>(__env);
600     }
601 
602     auto operator()(empty_env, empty_env) const -> empty_env
603     {
604         return {};
605     }
606 
607     template <class... Rest>
608     auto operator()(empty_env, Rest&&... rest) const -> decltype(auto)
609     {
610         return __fwd_fn()(__join_fn()(static_cast<Rest&&>(rest)...));
611     }
612 
613     template <class First, class... Rest>
614     auto operator()(First&& first, Rest&&... rest) const -> decltype(auto)
615     {
616         return __joined{__fwd_fn()(__join_fn()(static_cast<Rest&&>(rest)...)),
617                         static_cast<First&&>(first)};
618     }
619 };
620 
621 inline constexpr __join_fn __join{};
622 
623 template <class... _Envs>
624 using __join_t = __result_of<__join, _Envs...>;
625 } // namespace __env
626 
627 using __env::empty_env;
628 using __env::get_env_t;
629 inline constexpr get_env_t get_env{};
630 
631 template <class _EnvProvider>
632 concept environment_provider = //
633     requires(_EnvProvider& __ep) {
634         {
635             get_env(std::as_const(__ep))
636         } -> queryable;
637     };
638 } // namespace stdexec
639 
640 STDEXEC_PRAGMA_POP()
641