include/cppcoro/schedule_on.hpp

0001 ///////////////////////////////////////////////////////////////////////////////
0002 // Copyright (c) Lewis Baker
0003 // Licenced under MIT license. See LICENSE.txt for details.
0004 ///////////////////////////////////////////////////////////////////////////////
0005 #ifndef CPPCORO_SCHEDULE_ON_HPP_INCLUDED
0006 #define CPPCORO_SCHEDULE_ON_HPP_INCLUDED
0007
0008 #include <cppcoro/task.hpp>
0009 #include <cppcoro/shared_task.hpp>
0010 #include <cppcoro/async_generator.hpp>
0011 #include <cppcoro/awaitable_traits.hpp>
0012
0013 #include <cppcoro/detail/remove_rvalue_reference.hpp>
0014
0015 namespace cppcoro
0016 {
0017     template<typename SCHEDULER>
0018     struct schedule_on_transform
0019     {
0020         explicit schedule_on_transform(SCHEDULER& scheduler) noexcept
0021             : scheduler(scheduler)
0022         {}
0023
0024         SCHEDULER& scheduler;
0025     };
0026
0027     template<typename SCHEDULER>
0028     schedule_on_transform<SCHEDULER> schedule_on(SCHEDULER& scheduler)
0029     {
0030         return schedule_on_transform<SCHEDULER>{ scheduler };
0031     }
0032
0033     template<typename T, typename SCHEDULER>
0034     decltype(auto) operator|(T&& value, schedule_on_transform<SCHEDULER> transform)
0035     {
0036         return schedule_on(transform.scheduler, std::forward<T>(value));
0037     }
0038
0039     template<typename SCHEDULER, typename AWAITABLE>
0040     auto schedule_on(SCHEDULER& scheduler, AWAITABLE awaitable)
0041         -> task<detail::remove_rvalue_reference_t<typename awaitable_traits<AWAITABLE>::await_result_t>>
0042     {
0043         co_await scheduler.schedule();
0044         co_return co_await std::move(awaitable);
0045     }
0046
0047     template<typename T, typename SCHEDULER>
0048     async_generator<T> schedule_on(SCHEDULER& scheduler, async_generator<T> source)
0049     {
0050         // Transfer exection to the scheduler before the implicit calls to
0051         // 'co_await begin()' or subsequent calls to `co_await iterator::operator++()`
0052         // below. This ensures that all calls to the generator's coroutine_handle<>::resume()
0053         // are executed on the execution context of the scheduler.
0054         co_await scheduler.schedule();
0055
0056         const auto itEnd = source.end();
0057         auto it = co_await source.begin();
0058         while (it != itEnd)
0059         {
0060             co_yield *it;
0061
0062             co_await scheduler.schedule();
0063
0064             (void)co_await ++it;
0065         }
0066     }
0067 }
0068
0069 #endif