1   
//
1   
//
2   
// Copyright (c) 2025 Vinnie Falco (vinnie.falco@gmail.com)
2   
// Copyright (c) 2025 Vinnie Falco (vinnie.falco@gmail.com)
3   
//
3   
//
4   
// Distributed under the Boost Software License, Version 1.0. (See accompanying
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)
5   
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
6   
//
6   
//
7   
// Official repository: https://github.com/cppalliance/capy
7   
// Official repository: https://github.com/cppalliance/capy
8   
//
8   
//
9  
9  
10   
#ifndef BOOST_CAPY_EX_IO_AWAITABLE_SUPPORT_HPP
10   
#ifndef BOOST_CAPY_EX_IO_AWAITABLE_SUPPORT_HPP
11   
#define BOOST_CAPY_EX_IO_AWAITABLE_SUPPORT_HPP
11   
#define BOOST_CAPY_EX_IO_AWAITABLE_SUPPORT_HPP
12  
12  
13   
#include <boost/capy/detail/config.hpp>
13   
#include <boost/capy/detail/config.hpp>
14   
#include <boost/capy/ex/frame_allocator.hpp>
14   
#include <boost/capy/ex/frame_allocator.hpp>
15   
#include <boost/capy/ex/io_env.hpp>
15   
#include <boost/capy/ex/io_env.hpp>
16   
#include <boost/capy/ex/this_coro.hpp>
16   
#include <boost/capy/ex/this_coro.hpp>
17  
17  
18   
#include <coroutine>
18   
#include <coroutine>
19   
#include <cstddef>
19   
#include <cstddef>
20   
#include <memory_resource>
20   
#include <memory_resource>
21   
#include <stop_token>
21   
#include <stop_token>
22   
#include <type_traits>
22   
#include <type_traits>
23  
23  
24   
namespace boost {
24   
namespace boost {
25   
namespace capy {
25   
namespace capy {
26  
26  
27   
/** CRTP mixin that adds I/O awaitable support to a promise type.
27   
/** CRTP mixin that adds I/O awaitable support to a promise type.
28  
28  
29   
    Inherit from this class to enable these capabilities in your coroutine:
29   
    Inherit from this class to enable these capabilities in your coroutine:
30  
30  
31   
    1. **Frame allocation** — The mixin provides `operator new/delete` that
31   
    1. **Frame allocation** — The mixin provides `operator new/delete` that
32   
       use the thread-local frame allocator set by `run_async`.
32   
       use the thread-local frame allocator set by `run_async`.
33  
33  
34   
    2. **Environment storage** — The mixin stores a pointer to the `io_env`
34   
    2. **Environment storage** — The mixin stores a pointer to the `io_env`
35   
       containing the executor, stop token, and allocator for this coroutine.
35   
       containing the executor, stop token, and allocator for this coroutine.
36  
36  
37   
    3. **Environment access** — Coroutine code can retrieve the environment
37   
    3. **Environment access** — Coroutine code can retrieve the environment
38   
       via `co_await this_coro::environment`, or individual fields via
38   
       via `co_await this_coro::environment`, or individual fields via
39   
       `co_await this_coro::executor`, `co_await this_coro::stop_token`,
39   
       `co_await this_coro::executor`, `co_await this_coro::stop_token`,
40   
       and `co_await this_coro::allocator`.
40   
       and `co_await this_coro::allocator`.
41  
41  
42   
    @tparam Derived The derived promise type (CRTP pattern).
42   
    @tparam Derived The derived promise type (CRTP pattern).
43  
43  
44   
    @par Basic Usage
44   
    @par Basic Usage
45  
45  
46   
    For coroutines that need to access their execution environment:
46   
    For coroutines that need to access their execution environment:
47  
47  
48   
    @code
48   
    @code
49   
    struct my_task
49   
    struct my_task
50   
    {
50   
    {
51   
        struct promise_type : io_awaitable_support<promise_type>
51   
        struct promise_type : io_awaitable_support<promise_type>
52   
        {
52   
        {
53   
            my_task get_return_object();
53   
            my_task get_return_object();
54   
            std::suspend_always initial_suspend() noexcept;
54   
            std::suspend_always initial_suspend() noexcept;
55   
            std::suspend_always final_suspend() noexcept;
55   
            std::suspend_always final_suspend() noexcept;
56   
            void return_void();
56   
            void return_void();
57   
            void unhandled_exception();
57   
            void unhandled_exception();
58   
        };
58   
        };
59  
59  
60   
        // ... awaitable interface ...
60   
        // ... awaitable interface ...
61   
    };
61   
    };
62  
62  
63   
    my_task example()
63   
    my_task example()
64   
    {
64   
    {
65   
        auto env = co_await this_coro::environment;
65   
        auto env = co_await this_coro::environment;
66   
        // Access env->executor, env->stop_token, env->allocator
66   
        // Access env->executor, env->stop_token, env->allocator
67  
67  
68   
        // Or use fine-grained accessors:
68   
        // Or use fine-grained accessors:
69   
        auto ex = co_await this_coro::executor;
69   
        auto ex = co_await this_coro::executor;
70   
        auto token = co_await this_coro::stop_token;
70   
        auto token = co_await this_coro::stop_token;
71   
        auto* alloc = co_await this_coro::allocator;
71   
        auto* alloc = co_await this_coro::allocator;
72   
    }
72   
    }
73   
    @endcode
73   
    @endcode
74  
74  
75   
    @par Custom Awaitable Transformation
75   
    @par Custom Awaitable Transformation
76  
76  
77   
    If your promise needs to transform awaitables (e.g., for affinity or
77   
    If your promise needs to transform awaitables (e.g., for affinity or
78   
    logging), override `transform_awaitable` instead of `await_transform`:
78   
    logging), override `transform_awaitable` instead of `await_transform`:
79  
79  
80   
    @code
80   
    @code
81   
    struct promise_type : io_awaitable_support<promise_type>
81   
    struct promise_type : io_awaitable_support<promise_type>
82   
    {
82   
    {
83   
        template<typename A>
83   
        template<typename A>
84   
        auto transform_awaitable(A&& a)
84   
        auto transform_awaitable(A&& a)
85   
        {
85   
        {
86   
            // Your custom transformation logic
86   
            // Your custom transformation logic
87   
            return std::forward<A>(a);
87   
            return std::forward<A>(a);
88   
        }
88   
        }
89   
    };
89   
    };
90   
    @endcode
90   
    @endcode
91  
91  
92   
    The mixin's `await_transform` intercepts @ref this_coro::environment_tag
92   
    The mixin's `await_transform` intercepts @ref this_coro::environment_tag
93   
    and the fine-grained tag types (@ref this_coro::executor_tag,
93   
    and the fine-grained tag types (@ref this_coro::executor_tag,
94   
    @ref this_coro::stop_token_tag, @ref this_coro::allocator_tag),
94   
    @ref this_coro::stop_token_tag, @ref this_coro::allocator_tag),
95   
    then delegates all other awaitables to your `transform_awaitable`.
95   
    then delegates all other awaitables to your `transform_awaitable`.
96  
96  
97   
    @par Making Your Coroutine an IoAwaitable
97   
    @par Making Your Coroutine an IoAwaitable
98  
98  
99   
    The mixin handles the "inside the coroutine" part—accessing the
99   
    The mixin handles the "inside the coroutine" part—accessing the
100   
    environment. To receive the environment when your coroutine is awaited
100   
    environment. To receive the environment when your coroutine is awaited
101   
    (satisfying @ref IoAwaitable), implement the `await_suspend` overload
101   
    (satisfying @ref IoAwaitable), implement the `await_suspend` overload
102   
    on your coroutine return type:
102   
    on your coroutine return type:
103  
103  
104   
    @code
104   
    @code
105   
    struct my_task
105   
    struct my_task
106   
    {
106   
    {
107   
        struct promise_type : io_awaitable_support<promise_type> { ... };
107   
        struct promise_type : io_awaitable_support<promise_type> { ... };
108  
108  
109   
        std::coroutine_handle<promise_type> h_;
109   
        std::coroutine_handle<promise_type> h_;
110  
110  
111   
        // IoAwaitable await_suspend receives and stores the environment
111   
        // IoAwaitable await_suspend receives and stores the environment
112   
        std::coroutine_handle<> await_suspend(std::coroutine_handle<> cont, io_env const* env)
112   
        std::coroutine_handle<> await_suspend(std::coroutine_handle<> cont, io_env const* env)
113   
        {
113   
        {
114   
            h_.promise().set_environment(env);
114   
            h_.promise().set_environment(env);
115   
            // ... rest of suspend logic ...
115   
            // ... rest of suspend logic ...
116   
        }
116   
        }
117   
    };
117   
    };
118   
    @endcode
118   
    @endcode
119  
119  
120   
    @par Thread Safety
120   
    @par Thread Safety
121   
    The environment is stored during `await_suspend` and read during
121   
    The environment is stored during `await_suspend` and read during
122   
    `co_await this_coro::environment`. These occur on the same logical
122   
    `co_await this_coro::environment`. These occur on the same logical
123   
    thread of execution, so no synchronization is required.
123   
    thread of execution, so no synchronization is required.
124  
124  
125   
    @see this_coro::environment, this_coro::executor,
125   
    @see this_coro::environment, this_coro::executor,
126   
         this_coro::stop_token, this_coro::allocator
126   
         this_coro::stop_token, this_coro::allocator
127   
    @see io_env
127   
    @see io_env
128   
    @see IoAwaitable
128   
    @see IoAwaitable
129   
*/
129   
*/
130   
template<typename Derived>
130   
template<typename Derived>
131   
class io_awaitable_support
131   
class io_awaitable_support
132   
{
132   
{
133   
    io_env const* env_ = &detail::empty_io_env;
133   
    io_env const* env_ = &detail::empty_io_env;
134   
    mutable std::coroutine_handle<> cont_{std::noop_coroutine()};
134   
    mutable std::coroutine_handle<> cont_{std::noop_coroutine()};
135  
135  
136   
public:
136   
public:
137   
    //----------------------------------------------------------
137   
    //----------------------------------------------------------
138   
    // Frame allocation support
138   
    // Frame allocation support
139   
    //----------------------------------------------------------
139   
    //----------------------------------------------------------
140  
140  
141   
private:
141   
private:
142   
    static constexpr std::size_t ptr_alignment = alignof(void*);
142   
    static constexpr std::size_t ptr_alignment = alignof(void*);
143  
143  
144   
    static std::size_t
144   
    static std::size_t
145   
    aligned_offset(std::size_t n) noexcept
145   
    aligned_offset(std::size_t n) noexcept
146   
    {
146   
    {
147   
        return (n + ptr_alignment - 1) & ~(ptr_alignment - 1);
147   
        return (n + ptr_alignment - 1) & ~(ptr_alignment - 1);
148   
    }
148   
    }
149  
149  
150   
public:
150   
public:
151   
    /** Allocate a coroutine frame.
151   
    /** Allocate a coroutine frame.
152  
152  
153   
        Uses the thread-local frame allocator set by run_async.
153   
        Uses the thread-local frame allocator set by run_async.
154   
        Falls back to default memory resource if not set.
154   
        Falls back to default memory resource if not set.
155   
        Stores the allocator pointer at the end of each frame for
155   
        Stores the allocator pointer at the end of each frame for
156   
        correct deallocation even when TLS changes.
156   
        correct deallocation even when TLS changes.
157   
    */
157   
    */
158   
    static void*
158   
    static void*
159   
    operator new(std::size_t size)
159   
    operator new(std::size_t size)
160   
    {
160   
    {
161   
        auto* mr = current_frame_allocator();
161   
        auto* mr = current_frame_allocator();
162   
        if(!mr)
162   
        if(!mr)
163   
            mr = std::pmr::get_default_resource();
163   
            mr = std::pmr::get_default_resource();
164  
164  
165   
        // Allocate extra space for memory_resource pointer
165   
        // Allocate extra space for memory_resource pointer
166   
        std::size_t ptr_offset = aligned_offset(size);
166   
        std::size_t ptr_offset = aligned_offset(size);
167   
        std::size_t total = ptr_offset + sizeof(std::pmr::memory_resource*);
167   
        std::size_t total = ptr_offset + sizeof(std::pmr::memory_resource*);
168   
        void* raw = mr->allocate(total, alignof(std::max_align_t));
168   
        void* raw = mr->allocate(total, alignof(std::max_align_t));
169  
169  
170   
        // Store the allocator pointer at the end
170   
        // Store the allocator pointer at the end
171   
        auto* ptr_loc = reinterpret_cast<std::pmr::memory_resource**>(
171   
        auto* ptr_loc = reinterpret_cast<std::pmr::memory_resource**>(
172   
            static_cast<char*>(raw) + ptr_offset);
172   
            static_cast<char*>(raw) + ptr_offset);
173   
        *ptr_loc = mr;
173   
        *ptr_loc = mr;
174  
174  
175   
        return raw;
175   
        return raw;
176   
    }
176   
    }
177  
177  
178   
    /** Deallocate a coroutine frame.
178   
    /** Deallocate a coroutine frame.
179  
179  
180   
        Reads the allocator pointer stored at the end of the frame
180   
        Reads the allocator pointer stored at the end of the frame
181   
        to ensure correct deallocation regardless of current TLS.
181   
        to ensure correct deallocation regardless of current TLS.
182   
    */
182   
    */
183   
    static void
183   
    static void
184   
    operator delete(void* ptr, std::size_t size)
184   
    operator delete(void* ptr, std::size_t size)
185   
    {
185   
    {
186   
        // Read the allocator pointer from the end of the frame
186   
        // Read the allocator pointer from the end of the frame
187   
        std::size_t ptr_offset = aligned_offset(size);
187   
        std::size_t ptr_offset = aligned_offset(size);
188   
        auto* ptr_loc = reinterpret_cast<std::pmr::memory_resource**>(
188   
        auto* ptr_loc = reinterpret_cast<std::pmr::memory_resource**>(
189   
            static_cast<char*>(ptr) + ptr_offset);
189   
            static_cast<char*>(ptr) + ptr_offset);
190   
        auto* mr = *ptr_loc;
190   
        auto* mr = *ptr_loc;
191  
191  
192   
        std::size_t total = ptr_offset + sizeof(std::pmr::memory_resource*);
192   
        std::size_t total = ptr_offset + sizeof(std::pmr::memory_resource*);
193   
        mr->deallocate(ptr, total, alignof(std::max_align_t));
193   
        mr->deallocate(ptr, total, alignof(std::max_align_t));
194   
    }
194   
    }
195  
195  
196   
    ~io_awaitable_support()
196   
    ~io_awaitable_support()
197   
    {
197   
    {
198   
        // Abnormal teardown: destroy orphaned continuation
198   
        // Abnormal teardown: destroy orphaned continuation
199   
        if(cont_ != std::noop_coroutine())
199   
        if(cont_ != std::noop_coroutine())
200   
            cont_.destroy();
200   
            cont_.destroy();
201   
    }
201   
    }
202  
202  
203   
    //----------------------------------------------------------
203   
    //----------------------------------------------------------
204   
    // Continuation support
204   
    // Continuation support
205   
    //----------------------------------------------------------
205   
    //----------------------------------------------------------
206  
206  
207   
    /** Store the continuation to resume on completion.
207   
    /** Store the continuation to resume on completion.
208  
208  
209   
        Call this from your coroutine type's `await_suspend` overload
209   
        Call this from your coroutine type's `await_suspend` overload
210   
        to set up the completion path. The `final_suspend` awaiter
210   
        to set up the completion path. The `final_suspend` awaiter
211   
        returns this handle via unconditional symmetric transfer.
211   
        returns this handle via unconditional symmetric transfer.
212  
212  
213   
        @param cont The continuation to resume on completion.
213   
        @param cont The continuation to resume on completion.
214   
    */
214   
    */
215   
    void set_continuation(std::coroutine_handle<> cont) noexcept
215   
    void set_continuation(std::coroutine_handle<> cont) noexcept
216   
    {
216   
    {
217   
        cont_ = cont;
217   
        cont_ = cont;
218   
    }
218   
    }
219  
219  
220   
    /** Return and consume the stored continuation handle.
220   
    /** Return and consume the stored continuation handle.
221  
221  
222   
        Resets the stored handle to `noop_coroutine()` so the
222   
        Resets the stored handle to `noop_coroutine()` so the
223   
        destructor will not double-destroy it.
223   
        destructor will not double-destroy it.
224  
224  
225   
        @return The continuation for symmetric transfer.
225   
        @return The continuation for symmetric transfer.
226   
    */
226   
    */
227   
    std::coroutine_handle<> continuation() const noexcept
227   
    std::coroutine_handle<> continuation() const noexcept
228   
    {
228   
    {
229   
        return std::exchange(cont_, std::noop_coroutine());
229   
        return std::exchange(cont_, std::noop_coroutine());
230   
    }
230   
    }
231  
231  
232   
    //----------------------------------------------------------
232   
    //----------------------------------------------------------
233   
    // Environment support
233   
    // Environment support
234   
    //----------------------------------------------------------
234   
    //----------------------------------------------------------
235  
235  
236   
    /** Store a pointer to the execution environment.
236   
    /** Store a pointer to the execution environment.
237  
237  
238   
        Call this from your coroutine type's `await_suspend`
238   
        Call this from your coroutine type's `await_suspend`
239   
        overload to make the environment available via
239   
        overload to make the environment available via
240   
        `co_await this_coro::environment`. The pointed-to
240   
        `co_await this_coro::environment`. The pointed-to
241   
        `io_env` must outlive this coroutine.
241   
        `io_env` must outlive this coroutine.
242  
242  
243   
        @param env The environment to store.
243   
        @param env The environment to store.
244   
    */
244   
    */
245   
    void set_environment(io_env const* env) noexcept
245   
    void set_environment(io_env const* env) noexcept
246   
    {
246   
    {
247   
        env_ = env;
247   
        env_ = env;
248   
    }
248   
    }
249  
249  
250   
    /** Return the stored execution environment.
250   
    /** Return the stored execution environment.
251  
251  
252   
        @return The environment.
252   
        @return The environment.
253   
    */
253   
    */
254   
    io_env const* environment() const noexcept
254   
    io_env const* environment() const noexcept
255   
    {
255   
    {
256   
        return env_;
256   
        return env_;
257   
    }
257   
    }
258  
258  
259   
    /** Transform an awaitable before co_await.
259   
    /** Transform an awaitable before co_await.
260  
260  
261   
        Override this in your derived promise type to customize how
261   
        Override this in your derived promise type to customize how
262   
        awaitables are transformed. The default implementation passes
262   
        awaitables are transformed. The default implementation passes
263   
        the awaitable through unchanged.
263   
        the awaitable through unchanged.
264  
264  
265   
        @param a The awaitable expression from `co_await a`.
265   
        @param a The awaitable expression from `co_await a`.
266  
266  
267   
        @return The transformed awaitable.
267   
        @return The transformed awaitable.
268   
    */
268   
    */
269   
    template<typename A>
269   
    template<typename A>
270   
    decltype(auto) transform_awaitable(A&& a)
270   
    decltype(auto) transform_awaitable(A&& a)
271   
    {
271   
    {
272   
        return std::forward<A>(a);
272   
        return std::forward<A>(a);
273   
    }
273   
    }
274  
274  
275   
    /** Intercept co_await expressions.
275   
    /** Intercept co_await expressions.
276  
276  
277   
        This function handles @ref this_coro::environment_tag and
277   
        This function handles @ref this_coro::environment_tag and
278   
        the fine-grained tags (@ref this_coro::executor_tag,
278   
        the fine-grained tags (@ref this_coro::executor_tag,
279   
        @ref this_coro::stop_token_tag, @ref this_coro::allocator_tag)
279   
        @ref this_coro::stop_token_tag, @ref this_coro::allocator_tag)
280   
        specially, returning an awaiter that yields the stored value.
280   
        specially, returning an awaiter that yields the stored value.
281   
        All other awaitables are delegated to @ref transform_awaitable.
281   
        All other awaitables are delegated to @ref transform_awaitable.
282  
282  
283   
        @param t The awaited expression.
283   
        @param t The awaited expression.
284  
284  
285   
        @return An awaiter for the expression.
285   
        @return An awaiter for the expression.
286   
    */
286   
    */
287   
    template<typename T>
287   
    template<typename T>
288   
    auto await_transform(T&& t)
288   
    auto await_transform(T&& t)
289   
    {
289   
    {
290   
        using Tag = std::decay_t<T>;
290   
        using Tag = std::decay_t<T>;
291  
291  
292   
        if constexpr (std::is_same_v<Tag, this_coro::environment_tag>)
292   
        if constexpr (std::is_same_v<Tag, this_coro::environment_tag>)
293   
        {
293   
        {
294   
            struct awaiter
294   
            struct awaiter
295   
            {
295   
            {
296   
                io_env const* env_;
296   
                io_env const* env_;
297   
                bool await_ready() const noexcept { return true; }
297   
                bool await_ready() const noexcept { return true; }
298   
                void await_suspend(std::coroutine_handle<>) const noexcept { }
298   
                void await_suspend(std::coroutine_handle<>) const noexcept { }
299   
                io_env const* await_resume() const noexcept { return env_; }
299   
                io_env const* await_resume() const noexcept { return env_; }
300   
            };
300   
            };
301   
            return awaiter{env_};
301   
            return awaiter{env_};
302   
        }
302   
        }
303   
        else if constexpr (std::is_same_v<Tag, this_coro::executor_tag>)
303   
        else if constexpr (std::is_same_v<Tag, this_coro::executor_tag>)
304   
        {
304   
        {
305   
            struct awaiter
305   
            struct awaiter
306   
            {
306   
            {
307   
                executor_ref executor_;
307   
                executor_ref executor_;
308   
                bool await_ready() const noexcept { return true; }
308   
                bool await_ready() const noexcept { return true; }
309   
                void await_suspend(std::coroutine_handle<>) const noexcept { }
309   
                void await_suspend(std::coroutine_handle<>) const noexcept { }
310   
                executor_ref await_resume() const noexcept { return executor_; }
310   
                executor_ref await_resume() const noexcept { return executor_; }
311   
            };
311   
            };
312   
            return awaiter{env_->executor};
312   
            return awaiter{env_->executor};
313   
        }
313   
        }
314   
        else if constexpr (std::is_same_v<Tag, this_coro::stop_token_tag>)
314   
        else if constexpr (std::is_same_v<Tag, this_coro::stop_token_tag>)
315   
        {
315   
        {
316   
            struct awaiter
316   
            struct awaiter
317   
            {
317   
            {
318   
                std::stop_token token_;
318   
                std::stop_token token_;
319   
                bool await_ready() const noexcept { return true; }
319   
                bool await_ready() const noexcept { return true; }
320   
                void await_suspend(std::coroutine_handle<>) const noexcept { }
320   
                void await_suspend(std::coroutine_handle<>) const noexcept { }
321   
                std::stop_token await_resume() const noexcept { return token_; }
321   
                std::stop_token await_resume() const noexcept { return token_; }
322   
            };
322   
            };
323   
            return awaiter{env_->stop_token};
323   
            return awaiter{env_->stop_token};
324   
        }
324   
        }
325   
        else if constexpr (std::is_same_v<Tag, this_coro::allocator_tag>)
325   
        else if constexpr (std::is_same_v<Tag, this_coro::allocator_tag>)
326   
        {
326   
        {
327   
            struct awaiter
327   
            struct awaiter
328   
            {
328   
            {
329   
                std::pmr::memory_resource* allocator_;
329   
                std::pmr::memory_resource* allocator_;
330   
                bool await_ready() const noexcept { return true; }
330   
                bool await_ready() const noexcept { return true; }
331   
                void await_suspend(std::coroutine_handle<>) const noexcept { }
331   
                void await_suspend(std::coroutine_handle<>) const noexcept { }
332   
                std::pmr::memory_resource* await_resume() const noexcept { return allocator_; }
332   
                std::pmr::memory_resource* await_resume() const noexcept { return allocator_; }
333   
            };
333   
            };
334   
            return awaiter{env_->allocator};
334   
            return awaiter{env_->allocator};
335   
        }
335   
        }
336   
        else
336   
        else
337   
        {
337   
        {
338   
            return static_cast<Derived*>(this)->transform_awaitable(
338   
            return static_cast<Derived*>(this)->transform_awaitable(
339   
                std::forward<T>(t));
339   
                std::forward<T>(t));
340   
        }
340   
        }
341   
    }
341   
    }
342   
};
342   
};
343  
343  
344   
} // namespace capy
344   
} // namespace capy
345   
} // namespace boost
345   
} // namespace boost
346  
346  
347   
#endif
347   
#endif