1//
2// Copyright (c) 2022 Klemens Morgenstern (klemens.morgenstern@gmx.net)
3//
4// Distributed under the Boost Software License, Version 1.0. (See accompanying
5// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
6//
7
8#ifndef BOOST_COBALT_DETAIL_GATHER_HPP
9#define BOOST_COBALT_DETAIL_GATHER_HPP
10
11#include <boost/cobalt/detail/await_result_helper.hpp>
12#include <boost/cobalt/detail/exception.hpp>
13#include <boost/cobalt/detail/fork.hpp>
14#include <boost/cobalt/detail/forward_cancellation.hpp>
15#include <boost/cobalt/detail/util.hpp>
16#include <boost/cobalt/detail/wrapper.hpp>
17#include <boost/cobalt/task.hpp>
18#include <boost/cobalt/this_thread.hpp>
19
20#include <boost/asio/associated_cancellation_slot.hpp>
21#include <boost/asio/bind_cancellation_slot.hpp>
22#include <boost/asio/cancellation_signal.hpp>
23
24
25#include <boost/core/ignore_unused.hpp>
26#include <boost/intrusive_ptr.hpp>
27#include <boost/system/result.hpp>
28#include <boost/variant2/variant.hpp>
29
30#include <array>
31#include <coroutine>
32#include <algorithm>
33
34namespace boost::cobalt::detail
35{
36
37
38template<typename ... Args>
39struct gather_variadic_impl
40{
41 using tuple_type = std::tuple<decltype(get_awaitable_type(std::declval<Args&&>()))...>;
42
43 gather_variadic_impl(Args && ... args)
44 : args{std::forward<Args>(args)...}
45 {
46 }
47
48 std::tuple<Args...> args;
49
50 constexpr static std::size_t tuple_size = sizeof...(Args);
51
52 struct awaitable : fork::static_shared_state<256 * tuple_size>
53 {
54 template<std::size_t ... Idx>
55 awaitable(std::tuple<Args...> & args, std::index_sequence<Idx...>)
56 : aws(awaitable_type_getter<Args>(std::get<Idx>(args))...)
57 {
58 }
59
60 tuple_type aws;
61 std::array<asio::cancellation_signal, tuple_size> cancel;
62
63 template<typename T>
64 using result_store_part = variant2::variant<
65 variant2::monostate,
66 void_as_monostate<co_await_result_t<T>>,
67 std::exception_ptr>;
68
69 std::tuple<result_store_part<Args>...> result;
70
71
72 template<std::size_t Idx>
73 void interrupt_await_step()
74 {
75 using type= std::tuple_element_t<Idx, std::tuple<Args...>>;
76 using t = std::conditional_t<
77 std::is_reference_v<std::tuple_element_t<Idx, decltype(aws)>>,
78 co_awaitable_type<type> &,
79 co_awaitable_type<type> &&>;
80
81 if constexpr (interruptible<t>)
82 static_cast<t>(std::get<Idx>(aws)).interrupt_await();
83 }
84 void interrupt_await()
85 {
86 mp11::mp_for_each<mp11::mp_iota_c<sizeof...(Args)>>
87 ([&](auto idx)
88 {
89 interrupt_await_step<idx>();
90 });
91 }
92
93 // GCC doesn't like member funs
94 template<std::size_t Idx>
95 static detail::fork await_impl(awaitable & this_)
96 try
97 {
98 auto & aw = std::get<Idx>(this_.aws);
99 // check manually if we're ready
100 auto rd = aw.await_ready();
101 if (!rd)
102 {
103 co_await this_.cancel[Idx].slot();
104 // make sure the executor is set
105 co_await detail::fork::wired_up;
106 // do the await - this doesn't call await-ready again
107
108 if constexpr (std::is_void_v<decltype(aw.await_resume())>)
109 {
110 co_await aw;
111 std::get<Idx>(this_.result).template emplace<1u>();
112 }
113 else
114 std::get<Idx>(this_.result).template emplace<1u>(co_await aw);
115 }
116 else
117 {
118 if constexpr (std::is_void_v<decltype(aw.await_resume())>)
119 {
120 aw.await_resume();
121 std::get<Idx>(this_.result).template emplace<1u>();
122 }
123 else
124 std::get<Idx>(this_.result).template emplace<1u>(aw.await_resume());
125 }
126 }
127 catch(...)
128 {
129 std::get<Idx>(this_.result).template emplace<2u>(std::current_exception());
130 }
131
132 std::array<detail::fork(*)(awaitable&), tuple_size> impls {
133 []<std::size_t ... Idx>(std::index_sequence<Idx...>)
134 {
135 return std::array<detail::fork(*)(awaitable&), tuple_size>{&await_impl<Idx>...};
136 }(std::make_index_sequence<tuple_size>{})
137 };
138
139 detail::fork last_forked;
140 std::size_t last_index = 0u;
141
142 bool await_ready()
143 {
144 while (last_index < tuple_size)
145 {
146 last_forked = impls[last_index++](*this);
147 if (!last_forked.done())
148 return false; // one coro didn't immediately complete!
149 }
150 last_forked.release();
151 return true;
152 }
153
154 template<typename H>
155 auto await_suspend(
156 std::coroutine_handle<H> h
157#if defined(BOOST_ASIO_ENABLE_HANDLER_TRACKING)
158 , const boost::source_location & loc = BOOST_CURRENT_LOCATION
159#endif
160 )
161 {
162#if defined(BOOST_ASIO_ENABLE_HANDLER_TRACKING)
163 this->loc = loc;
164#endif
165 this->exec = &cobalt::detail::get_executor(h);
166 last_forked.release().resume();
167 while (last_index < tuple_size)
168 impls[last_index++](*this).release();
169
170 if (!this->outstanding_work()) // already done, resume rightaway.
171 return false;
172
173 // arm the cancel
174 assign_cancellation(
175 h,
176 [&](asio::cancellation_type ct)
177 {
178 for (auto & cs : cancel)
179 cs.emit(ct);
180 });
181
182
183 this->coro.reset(h.address());
184 return true;
185 }
186
187 template<typename T>
188 using result_part = system::result<co_await_result_t<T>, std::exception_ptr>;
189
190#if _MSC_VER
191 BOOST_NOINLINE
192#endif
193 std::tuple<result_part<Args> ...> await_resume()
194 {
195 return mp11::tuple_transform(
196 []<typename T>(variant2::variant<variant2::monostate, T, std::exception_ptr> & var)
197 -> system::result<monostate_as_void<T>, std::exception_ptr>
198 {
199 BOOST_ASSERT(var.index() != 0u);
200 if (var.index() == 1u)
201 {
202 if constexpr (std::is_same_v<T, variant2::monostate>)
203 return {system::in_place_value};
204 else
205 return {system::in_place_value, std::move(get<1>(var))};
206 }
207 else
208 return {system::in_place_error, std::move(get<2>(var))};
209
210 }
211 , result);
212 }
213 };
214 awaitable operator co_await() &&
215 {
216 return awaitable(args, std::make_index_sequence<sizeof...(Args)>{});
217 }
218};
219
220template<typename Range>
221struct gather_ranged_impl
222{
223 Range aws;
224
225 using result_type = system::result<
226 co_await_result_t<std::decay_t<decltype(*std::begin(std::declval<Range>()))>>,
227 std::exception_ptr>;
228
229 using result_storage_type = variant2::variant<
230 variant2::monostate,
231 void_as_monostate<
232 co_await_result_t<std::decay_t<decltype(*std::begin(std::declval<Range>()))>>
233 >,
234 std::exception_ptr>;
235
236 struct awaitable : fork::shared_state
237 {
238 using type = std::decay_t<decltype(*std::begin(std::declval<Range>()))>;
239#if !defined(BOOST_COBALT_NO_PMR)
240 pmr::polymorphic_allocator<void> alloc{&resource};
241 std::conditional_t<awaitable_type<type>, Range &,
242 pmr::vector<co_awaitable_type<type>>> aws;
243
244 pmr::vector<bool> ready{std::size(aws), alloc};
245 pmr::vector<asio::cancellation_signal> cancel{std::size(aws), alloc};
246 pmr::vector<result_storage_type> result{cancel.size(), alloc};
247
248#else
249 std::allocator<void> alloc{};
250 std::conditional_t<awaitable_type<type>, Range &,
251 std::vector<co_awaitable_type<type>>> aws;
252
253 std::vector<bool> ready{std::size(aws), alloc};
254 std::vector<asio::cancellation_signal> cancel{std::size(aws), alloc};
255 std::vector<result_storage_type> result{cancel.size(), alloc};
256#endif
257
258
259 awaitable(Range & aws_, std::false_type /* needs operator co_await */)
260 : fork::shared_state((512 + sizeof(co_awaitable_type<type>)) * std::size(aws_))
261 , aws{alloc}
262 , ready{std::size(aws_), alloc}
263 , cancel{std::size(aws_), alloc}
264 {
265 aws.reserve(std::size(aws_));
266 for (auto && a : aws_)
267 {
268 using a_0 = std::decay_t<decltype(a)>;
269 using a_t = std::conditional_t<
270 std::is_lvalue_reference_v<Range>, a_0 &, a_0 &&>;
271 aws.emplace_back(awaitable_type_getter<a_t>(static_cast<a_t>(a)));
272 }
273
274
275 std::transform(std::begin(this->aws),
276 std::end(this->aws),
277 std::begin(cont&: ready),
278 [](auto & aw) {return aw.await_ready();});
279 }
280 awaitable(Range & aws, std::true_type /* needs operator co_await */)
281 : fork::shared_state((512 + sizeof(co_awaitable_type<type>)) * std::size(aws))
282 , aws(aws)
283 {
284 std::transform(std::begin(aws), std::end(aws), std::begin(cont&: ready), [](auto & aw) {return aw.await_ready();});
285 }
286
287 awaitable(Range & aws)
288 : awaitable(aws, std::bool_constant<awaitable_type<type>>{})
289 {
290 }
291
292 void interrupt_await()
293 {
294 using t = std::conditional_t<std::is_reference_v<Range>,
295 co_awaitable_type<type> &,
296 co_awaitable_type<type> &&>;
297
298 if constexpr (interruptible<t>)
299 for (auto & aw : aws)
300 static_cast<t>(aw).interrupt_await();
301 }
302
303 static detail::fork await_impl(awaitable & this_, std::size_t idx)
304 try
305 {
306 auto & aw = *std::next(std::begin(this_.aws), idx);
307 auto rd = aw.await_ready();
308 if (!rd)
309 {
310 co_await this_.cancel[idx].slot();
311 co_await detail::fork::wired_up;
312 if constexpr (std::is_void_v<decltype(aw.await_resume())>)
313 {
314 co_await aw;
315 this_.result[idx].template emplace<1u>();
316 }
317 else
318 this_.result[idx].template emplace<1u>(co_await aw);
319 }
320 else
321 {
322 if constexpr (std::is_void_v<decltype(aw.await_resume())>)
323 {
324 aw.await_resume();
325 this_.result[idx].template emplace<1u>();
326 }
327 else
328 this_.result[idx].template emplace<1u>(aw.await_resume());
329 }
330 }
331 catch(...)
332 {
333 this_.result[idx].template emplace<2u>(std::current_exception());
334
335 }
336
337 detail::fork last_forked;
338 std::size_t last_index = 0u;
339
340 bool await_ready()
341 {
342 while (last_index < cancel.size())
343 {
344 last_forked = await_impl(this_&: *this, idx: last_index++);
345 if (!last_forked.done())
346 return false; // one coro didn't immediately complete!
347 }
348 last_forked.release();
349 return true;
350 }
351
352 template<typename H>
353 auto await_suspend(
354 std::coroutine_handle<H> h
355#if defined(BOOST_ASIO_ENABLE_HANDLER_TRACKING)
356 , const boost::source_location & loc = BOOST_CURRENT_LOCATION
357#endif
358 )
359 {
360#if defined(BOOST_ASIO_ENABLE_HANDLER_TRACKING)
361 this->loc = loc;
362#endif
363 exec = &detail::get_executor(h);
364
365 last_forked.release().resume();
366 while (last_index < cancel.size())
367 await_impl(this_&: *this, idx: last_index++).release();
368
369 if (!this->outstanding_work()) // already done, resume rightaway.
370 return false;
371
372 // arm the cancel
373 assign_cancellation(
374 h,
375 [&](asio::cancellation_type ct)
376 {
377 for (auto & cs : cancel)
378 cs.emit(type: ct);
379 });
380
381 this->coro.reset(handle: h.address());
382 return true;
383 }
384
385#if _MSC_VER
386 BOOST_NOINLINE
387#endif
388 auto await_resume()
389 {
390#if !defined(BOOST_COBALT_NO_PMR)
391 pmr::vector<result_type> res{result.size(), this_thread::get_allocator()};
392#else
393 std::vector<result_type> res(result.size());
394#endif
395
396 std::transform(
397 result.begin(), result.end(), res.begin(),
398 [](result_storage_type & res) -> result_type
399 {
400 BOOST_ASSERT(res.index() != 0u);
401 if (res.index() == 1u)
402 {
403 if constexpr (std::is_void_v<typename result_type::value_type>)
404 return system::in_place_value;
405 else
406 return {system::in_place_value, std::move(get<1u>(res))};
407 }
408 else
409 return {system::in_place_error, get<2u>(res)};
410 });
411
412 return res;
413 }
414 };
415 awaitable operator co_await() && {return awaitable{aws};}
416};
417
418}
419
420#endif //BOOST_COBALT_DETAIL_GATHER_HPP
421

source code of boost/libs/cobalt/include/boost/cobalt/detail/gather.hpp