scoped_task.hpp

// Copyright Take Vos 2022.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at https://www.boost.org/LICENSE_1_0.txt)
 
#pragma once
 
#include "../utility/utility.hpp"
#include "../concurrency/concurrency.hpp"
#include "../macros.hpp"
#include <coroutine>
#include <type_traits>
#include <memory>
#include <exception>
 
 
 
namespace hi::inline v1 {

template<typename T = void>
class scoped_task {
public:
    using value_type = T;
    using notifier_type = notifier<void(value_type)>;
    using callback_token = notifier_type::callback_token;
    using callback_proto = notifier_type::callback_proto;
 
    struct promise_type {
        notifier_type notifier;
        std::optional<value_type> value = {};
        std::exception_ptr exception = nullptr;
 
        void return_value(std::convertible_to<value_type> auto&& new_value) noexcept
        {
            value = hi_forward(new_value);
        }
 
        void unhandled_exception() noexcept
        {
            exception = std::current_exception();
        }
 
        std::suspend_always final_suspend() noexcept
        {
            if (value) {
                // Trigger the notifier with the co_return value.
                notifier(*value);
 
            } else {
                // Notify also in case of exception.
                hi_assert_not_null(exception);
                if constexpr (std::is_default_constructible_v<value_type>) {
                    notifier(value_type{});
                }
            }
 
            return {};
        }
 
        scoped_task get_return_object() noexcept
        {
            return scoped_task{handle_type::from_promise(*this)};
        }

        std::suspend_never initial_suspend() noexcept
        {
            return {};
        }
    };
 
    using handle_type = std::coroutine_handle<promise_type>;
 
    scoped_task(handle_type coroutine) noexcept : _coroutine(coroutine) {}
 
    ~scoped_task()
    {
        if (_coroutine) {
            _coroutine.destroy();
        }
    }
 
    scoped_task() = default;
 
    // scoped_task can not be copied because it tracks if the co-routine must be destroyed by the
    // shared_ptr to the value shared between scoped_task and the promise.
    scoped_task(scoped_task const&) = delete;
    scoped_task& operator=(scoped_task const&) = delete;
 
    scoped_task(scoped_task&& other) noexcept
    {
        _coroutine = std::exchange(other._coroutine, {});
    }
 
    scoped_task& operator=(scoped_task&& other) noexcept
    {
        _coroutine = std::exchange(other._coroutine, {});
        return *this;
    }

    [[nodiscard]] bool done() const noexcept
    {
        hi_axiom(_coroutine);
        return _coroutine.done();
    }

    explicit operator bool() const noexcept
    {
        return done();
    }

    [[nodiscard]] value_type const& value() const
    {
        hi_axiom(done());
 
        hilet& promise = _coroutine.promise();
        if (promise.value) {
            return *promise.value;
 
        } else {
            hi_axiom(promise.exception);
            std::rethrow_exception(promise.exception);
        }
    }

    [[nodiscard]] value_type const& operator*() const
    {
        return value();
    }

    callback_token subscribe(forward_of<callback_proto> auto&& callback, callback_flags flags = callback_flags::synchronous) noexcept
    {
        return _coroutine.promise().notifier.subscribe(hi_forward(callback), flags);
    }
 
private:
    // Optional value type
    handle_type _coroutine;
};

template<>
class scoped_task<void> {
public:
    using value_type = void;
    using notifier_type = notifier<>;
    using callback_token = notifier_type::callback_token;
    using callback_proto = notifier_type::callback_proto;
 
    struct promise_type {
        notifier_type notifier;
        std::exception_ptr exception = nullptr;
 
        void return_void() noexcept {}
 
        void unhandled_exception() noexcept
        {
            exception = std::current_exception();
        }
 
        std::suspend_always final_suspend() noexcept
        {
            notifier();
            return {};
        }
 
        scoped_task get_return_object() noexcept
        {
            return scoped_task{handle_type::from_promise(*this)};
        }
 
        std::suspend_never initial_suspend() noexcept
        {
            return {};
        }
    };
 
    using handle_type = std::coroutine_handle<promise_type>;
 
    scoped_task(handle_type coroutine) noexcept : _coroutine(coroutine) {}
 
    ~scoped_task()
    {
        if (_coroutine) {
            _coroutine.destroy();
        }
    }
 
    scoped_task() = default;
    scoped_task(scoped_task const&) = delete;
    scoped_task& operator=(scoped_task const&) = delete;
 
    scoped_task(scoped_task&& other) noexcept
    {
        _coroutine = std::exchange(other._coroutine, {});
    }
 
    scoped_task& operator=(scoped_task&& other) noexcept
    {
        _coroutine = std::exchange(other._coroutine, {});
        return *this;
    }

    [[nodiscard]] bool done() const noexcept
    {
        hi_axiom(_coroutine);
        return _coroutine.done();
    }

    explicit operator bool() const noexcept
    {
        return done();
    }

    void value() const
    {
        hi_axiom(done());
 
        hilet& promise = _coroutine.promise();
        if (promise.exception) {
            std::rethrow_exception(promise.exception);
        }
    }

    callback_token subscribe(forward_of<callback_proto> auto&& callback, callback_flags flags = callback_flags::synchronous) noexcept
    {
        return _coroutine.promise().notifier.subscribe(hi_forward(callback), flags);
    }
 
private:
    // Optional value type
    handle_type _coroutine;
};
 
} // namespace hi::inline v1