task.h

Go to the documentation of this file.

/**********************************************************************
 * LeechCraft - modular cross-platform feature rich internet client.
 * Copyright (C) 2006-2014  Georg Rudoy
 *
 * Distributed under the Boost Software License, Version 1.0.
 * (See accompanying file LICENSE or copy at https://www.boost.org/LICENSE_1_0.txt)
 **********************************************************************/
 
#pragma once
 
#include <coroutine>
#include <exception>
#include <mutex>
#include <utility>
#include "finalsuspender.h"
#include "defaulthandlers.h"
#include "taskfwd.h"
 
namespace LC::Util
{
    namespace detail
    {
        struct PromiseBase
        {
            std::atomic<size_t> Refs_ = 1;
            std::exception_ptr Exception_ {};
 
            enum class CoroState : std::uint8_t
            {
                Running,
                Suspended,
                Finished,
            };
            CoroState State_ = CoroState::Running;
        };
 
        template<typename R>
        struct PromiseRet : PromiseBase
        {
            using ReturnType_t = R;
 
            constexpr static bool IsVoid = false;
 
            std::optional<R> Ret_;
 
            template<typename U = R>
            void return_value (this auto&& self, U&& val)
            {
                std::lock_guard guard { self };
                self.Ret_.emplace (std::forward<U> (val));
            }
        };
 
        template<>
        struct PromiseRet<void> : PromiseBase
        {
            constexpr static bool IsVoid = true;
 
            bool Done_ = false;
 
            void return_void (this auto&& self) noexcept
            {
                std::lock_guard guard { self };
                self.Done_ = true;
            }
        };
 
        struct EitherFailureAbort final : std::exception {};
 
        template<typename Promise>
        struct TaskAwaiter
        {
            using Handle_t = std::coroutine_handle<Promise>;
            Handle_t Subtask_;
 
            std::coroutine_handle<> OuterTask_;
 
            explicit TaskAwaiter (Handle_t subtask)
            : Subtask_ { subtask }
            {
                if (Subtask_)
                    Subtask_.promise ().IncRef ();
            }
 
            TaskAwaiter (TaskAwaiter&& other) noexcept
            : Subtask_ { std::exchange (other.Subtask_, {}) }
            , OuterTask_ { std::exchange (other.OuterTask_, {}) }
            {
            }
 
            TaskAwaiter (const TaskAwaiter&) = delete;
            TaskAwaiter& operator= (const TaskAwaiter&) = delete;
            TaskAwaiter& operator= (TaskAwaiter&& other) = delete;
 
            ~TaskAwaiter ()
            {
                if (Subtask_ && OuterTask_)
                {
                    auto& promise = Subtask_.promise ();
                    std::lock_guard guard { promise };
                    promise.RemoveAwaiter (OuterTask_);
                }
 
                if (Subtask_)
                    Subtask_.promise ().DecRef ();
            }
 
            bool await_ready () const noexcept
            {
                const auto& promise = Subtask_.promise ();
                std::lock_guard guard { promise };
                return CheckTaskFinishedUnlocked (promise);
            }
 
            bool await_suspend (std::coroutine_handle<> handle)
            {
                auto& promise = Subtask_.promise ();
                std::lock_guard guard { promise };
                if (CheckTaskFinishedUnlocked (promise))
                    return false;
 
                OuterTask_ = handle;
                promise.AddAwaiter (handle);
                return true;
            }
 
            auto await_resume ()
            {
                auto& promise = Subtask_.promise ();
                std::lock_guard guard { promise };
                if (promise.Exception_)
                    try
                    {
                        std::rethrow_exception (promise.Exception_);
                    }
                    catch (const EitherFailureAbort&)
                    {
                    }
 
                if constexpr (!Promise::IsVoid)
                    return Promise::ResumeValue (*promise.Ret_);
            }
        private:
            bool CheckTaskFinishedUnlocked (const Promise& promise) const
            {
                if (promise.Exception_)
                    return true;
 
                if constexpr (Promise::IsVoid)
                    return promise.Done_;
                else
                    return static_cast<bool> (promise.Ret_);
            }
        };
    }
 
    template<typename R, template<typename> typename... Extensions>
    class Task
    {
    public:
        struct promise_type;
    private:
        using Handle_t = std::coroutine_handle<promise_type>;
        Handle_t Handle_;
    public:
        struct promise_type
            : detail::PromiseRet<R>
            , Extensions<promise_type>...
            , detail::DefaultAwaiterHandler<Extensions<promise_type>...>
            , detail::DefaultLockingHandler<Extensions<promise_type>...>
            , detail::DefaultResumeValueHandler<Extensions<promise_type>...>
        {
            auto GetAddress () { return Handle_t::from_promise (*this).address (); }
 
            Task get_return_object ()
            {
                return Task { Handle_t::from_promise (*this) };
            }
 
            std::suspend_never initial_suspend () const noexcept { return {}; }
 
            auto final_suspend () noexcept
            {
                this->State_ = detail::PromiseBase::CoroState::Finished;
                return detail::FinalSuspender<promise_type> { *this };
            }
 
            void unhandled_exception ()
            {
                this->Exception_ = std::current_exception ();
            }
 
            void IncRef ()
            {
                this->Refs_.fetch_add (1);
            }
 
            void DecRef ()
            {
                if (this->Refs_.fetch_sub (1) == 1)
                    Handle_t::from_promise (*this).destroy ();
            }
        };
 
        using ResultType_t = R;
 
        template<typename RR>
        using ReplaceResult_t = Task<RR, Extensions...>;
 
        template<template<typename> typename F>
        using ApplyResult_t = Task<F<R>, Extensions...>;
 
        explicit Task (const std::coroutine_handle<promise_type>& handle)
        : Handle_ { handle }
        {
            if (handle)
                handle.promise ().IncRef ();
        }
 
        ~Task () noexcept
        {
            if (Handle_)
                Handle_.promise ().DecRef ();
        }
 
        Task (const Task& other)
        : Handle_ { other.Handle_ }
        {
            if (Handle_)
                Handle_.promise ().IncRef ();
        }
 
        Task& operator= (const Task& other)
        {
            Task task { other };
            *this = std::move (task);
            return *this;
        }
 
        Task (Task&& other) noexcept
        {
            std::swap (Handle_, other.Handle_);
        }
 
        Task& operator= (Task&& other) noexcept
        {
            std::swap (Handle_, other.Handle_);
            return *this;
        }
 
        auto operator co_await () const noexcept
        {
            return detail::TaskAwaiter<promise_type> { Handle_ };
        }
    };
 
    namespace detail
    {
        template<typename R, template<typename> typename... Extensions>
        struct GetPromise
        {
            using Promise = typename Task<R, Extensions...>::promise_type;
            Promise *Promise_ = nullptr;
 
            bool await_ready () const noexcept { return false; }
 
            bool await_suspend (std::coroutine_handle<Promise> handle) noexcept
            {
                Promise_ = &handle.promise ();
                return false;
            }
 
            decltype (auto) await_resume () const noexcept
            {
                return *Promise_;
            }
        };
    }
}