intial commit

source/mijin/async/coroutine.hpp

@@ -0,0 +1,648 @@
+
+#pragma once
+
+#ifndef MIJIN_ASYNC_COROUTINE_HPP_INCLUDED
+#define MIJIN_ASYNC_COROUTINE_HPP_INCLUDED 1
+
+#include <any>
+#include <coroutine>
+#include <memory>
+#include <thread>
+#include <tuple>
+
+#include "./future.hpp"
+#include "./message_queue.hpp"
+#include "../container/optional.hpp"
+#include "../util/flag.hpp"
+#include "../util/traits.hpp"
+
+namespace mijin
+{
+
+//
+// public defines
+//
+
+#if !defined(MIJIN_COROUTINE_ENABLE_DEBUG_INFO)
+#define MIJIN_COROUTINE_ENABLE_DEBUG_INFO 0
+#endif
+
+//
+// public types
+//
+
+enum class TaskStatus
+{
+    SUSPENDED = 0,
+    RUNNING = 1,
+    WAITING = 2,
+    FINISHED = 3,
+    YIELDED = 4
+};
+
+// forward declarations
+template<typename T>
+struct TaskState;
+
+class TaskLoop;
+
+template<typename TResult = void>
+class TaskBase;
+
+namespace impl
+{
+template<typename TReturn, typename TPromise>
+struct TaskReturn
+{
+    template<typename... TArgs>
+    constexpr void return_value(TArgs&&... args) noexcept {
+        *(static_cast<TPromise&>(*this).state_) = TaskState<TReturn>(TReturn(std::forward<TArgs>(args)...), TaskStatus::FINISHED);
+    }
+
+    constexpr void return_value(TReturn value) noexcept {
+        *(static_cast<TPromise&>(*this).state_) = TaskState<TReturn>(TReturn(std::move(value)), TaskStatus::FINISHED);
+    }
+};
+
+template<typename TPromise>
+struct TaskReturn<void, TPromise>
+{
+    constexpr void return_void() noexcept {
+        static_cast<TPromise&>(*this).state_->status = TaskStatus::FINISHED;
+    }
+};
+} // namespace impl
+
+template<typename T>
+struct TaskState
+{
+    Optional<T> value;
+    TaskStatus status = TaskStatus::SUSPENDED;
+
+    TaskState() = default;
+    TaskState(const TaskState&) = default;
+    TaskState(TaskState&&) noexcept = default;
+    constexpr TaskState(T _value, TaskStatus _status) noexcept : value(std::move(_value)), status(_status) {}
+    TaskState& operator=(const TaskState&) = default;
+    TaskState& operator=(TaskState&&) noexcept = default;
+};
+template<>
+struct TaskState<void>
+{
+    TaskStatus status = TaskStatus::SUSPENDED;
+
+    TaskState() = default;
+    TaskState(const TaskState&) = default;
+    TaskState(TaskState&&) noexcept = default;
+    constexpr TaskState(TaskStatus _status) noexcept : status(_status) {}
+    TaskState& operator=(const TaskState&) = default;
+    TaskState& operator=(TaskState&&) noexcept = default;
+};
+
+template<typename TValue>
+struct TaskAwaitableFuture
+{
+    FuturePtr<TValue> future;
+
+    [[nodiscard]] constexpr bool await_ready() const noexcept { return future->ready(); }
+    constexpr void await_suspend(std::coroutine_handle<>) const noexcept {}
+    constexpr TValue await_resume() const noexcept {
+        if constexpr (std::is_same_v<TValue, void>) {
+            return;
+        }
+        else {
+            return std::move(future->get());
+        }
+    }
+};
+
+template<typename... TArgs>
+struct TaskAwaitableSignal
+{
+    std::shared_ptr<std::tuple<TArgs...>> data;
+
+    [[nodiscard]] constexpr bool await_ready() const noexcept { return false; }
+    constexpr void await_suspend(std::coroutine_handle<>) const noexcept {}
+    constexpr auto& await_resume() const noexcept {
+        return *data;
+    }
+};
+
+template<typename TSingleArg>
+struct TaskAwaitableSignal<TSingleArg>
+{
+    std::shared_ptr<TSingleArg> data;
+
+    [[nodiscard]] constexpr bool await_ready() const noexcept { return false; }
+    constexpr void await_suspend(std::coroutine_handle<>) const noexcept {}
+    constexpr auto& await_resume() const noexcept {
+        return *data;
+    }
+};
+
+template<>
+struct TaskAwaitableSignal<>
+{
+    [[nodiscard]] constexpr bool await_ready() const noexcept { return false; }
+    constexpr void await_suspend(std::coroutine_handle<>) const noexcept {}
+    constexpr void await_resume() const noexcept {}
+};
+
+template<typename TTraits>
+struct TaskPromise : impl::TaskReturn<typename TTraits::result_t, TaskPromise<TTraits>>
+{
+    using handle_t = std::coroutine_handle<TaskPromise>;
+    using task_t = typename TTraits::task_t;
+    using result_t = typename TTraits::result_t;
+
+    std::shared_ptr<TaskState<result_t>> state_ = std::make_shared<TaskState<result_t>>();
+    TaskLoop* loop_ = nullptr;
+
+    constexpr task_t get_return_object() noexcept { return task_t(handle_t::from_promise(*this)); }
+    constexpr std::suspend_always initial_suspend() noexcept { return {}; }
+    constexpr std::suspend_always final_suspend() noexcept { return {}; }
+    
+    // template<typename TValue>
+    // constexpr std::suspend_always yield_value(TValue value) noexcept {
+    //     *state_ = TaskState<result_t>(std::move(value), TaskStatus::YIELDED);
+    //     return {};
+    // }
+
+    // TODO: implement yielding (can't use futures for this)
+
+    constexpr void unhandled_exception() noexcept {}
+
+    template<typename TValue>
+    auto await_transform(FuturePtr<TValue> future) noexcept
+    {
+        MIJIN_ASSERT(loop_ != nullptr, "Cannot await future outside of a loop!");
+        TaskAwaitableFuture<TValue> awaitable{future};
+        if (!awaitable.await_ready())
+        {
+            state_->status = TaskStatus::WAITING;
+            future->sigSet.connect([this, future]() mutable
+            {
+                state_->status = TaskStatus::SUSPENDED;
+            }, Oneshot::YES);
+        }
+        return awaitable;
+    }
+
+    template<typename TResultOther>
+    auto await_transform(TaskBase<TResultOther> task) noexcept
+    {
+        MIJIN_ASSERT(loop_ != nullptr, "Cannot await another task outside of a loop!");
+        auto future = delayEvaluation<TResultOther>(loop_)->addTask(std::move(task)); // hackidyhack: delay evaluation of the type of loop_ as it is only forward-declared here
+        return await_transform(future);
+    }
+
+    template<typename TFirstArg, typename TSecondArg, typename... TArgs>
+    auto await_transform(Signal<TFirstArg, TSecondArg, TArgs...>& signal) noexcept
+    {
+        auto data = std::make_shared<std::tuple<TFirstArg, TSecondArg, TArgs...>>();
+        signal.connect([this, data](TFirstArg arg0, TSecondArg arg1, TArgs... args) mutable
+        {
+            *data = std::make_tuple(std::move(arg0), std::move(arg1), std::move(args)...);
+            state_->status = TaskStatus::SUSPENDED;
+        }, Oneshot::YES);
+        TaskAwaitableSignal<TFirstArg, TSecondArg, TArgs...> awaitable{data};
+        state_->status = TaskStatus::WAITING;
+        return awaitable;
+    }
+
+    template<typename TFirstArg>
+    auto await_transform(Signal<TFirstArg>& signal) noexcept
+    {
+        auto data = std::make_shared<TFirstArg>();
+        signal.connect([this, data](TFirstArg arg0) mutable
+        {
+            *data = std::move(arg0);
+            state_->status = TaskStatus::SUSPENDED;
+        }, Oneshot::YES);
+        TaskAwaitableSignal<TFirstArg> awaitable{data};
+        state_->status = TaskStatus::WAITING;
+        return awaitable;
+    }
+
+    auto await_transform(Signal<>& signal) noexcept
+    {
+        signal.connect([this]()
+        {
+            state_->status = TaskStatus::SUSPENDED;
+        }, Oneshot::YES);
+        TaskAwaitableSignal<> awaitable{};
+        state_->status = TaskStatus::WAITING;
+        return awaitable;
+    }
+
+    std::suspend_always await_transform(std::suspend_always) noexcept
+    {
+        state_->status = TaskStatus::SUSPENDED;
+        return std::suspend_always();
+    }
+
+    std::suspend_never await_transform(std::suspend_never) noexcept {
+        return std::suspend_never();
+    }
+};
+
+template<typename TResult>
+class TaskBase
+{
+public:
+    using task_t = TaskBase;
+    using result_t = TResult;
+    struct Traits
+    {
+        using task_t = TaskBase;
+        using result_t = TResult;
+    };
+public:
+    using promise_type = TaskPromise<Traits>;
+    using handle_t = typename promise_type::handle_t;
+private:
+    handle_t handle_;
+    std::shared_ptr<TaskState<result_t>> state_;
+public:
+    constexpr explicit TaskBase(handle_t handle) noexcept : handle_(handle), state_(handle.promise().state_) {}
+    TaskBase(const TaskBase&) = default;
+    TaskBase(TaskBase&& other) noexcept = default;
+    ~TaskBase() noexcept;
+public:
+    TaskBase& operator=(const TaskBase&) = default;
+    TaskBase& operator=(TaskBase&& other) noexcept = default;
+
+    [[nodiscard]]
+    constexpr bool operator==(const TaskBase& other) const noexcept { return handle_ == other.handle_; }
+
+    [[nodiscard]]
+    constexpr bool operator!=(const TaskBase& other) const noexcept { return handle_ != other.handle_; }
+public:
+    constexpr TaskState<TResult>& resume() noexcept
+    {
+        state_->status = TaskStatus::RUNNING;
+        handle_.resume();
+        return *state_;
+    }
+
+    [[nodiscard]]
+    constexpr TaskState<TResult>& state() noexcept
+    {
+        return *state_;
+    }
+private:
+    [[nodiscard]]
+    constexpr handle_t handle() const noexcept { return handle_; }
+    [[nodiscard]]
+    constexpr TaskLoop* getLoop() noexcept
+    {
+        return handle_.promise().loop_;
+    }
+    constexpr void setLoop(TaskLoop* loop) noexcept
+    {
+        // MIJIN_ASSERT(handle_.promise().loop_ == nullptr
+        //     || handle_.promise().loop_ == loop
+        //     || loop == nullptr, "Task already has a loop assigned!");
+        handle_.promise().loop_ = loop;
+    }
+
+    friend class TaskLoop;
+
+    template<typename TTask>
+    friend class WrappedTask;
+};
+
+class WrappedTaskBase
+{
+public:
+    virtual ~WrappedTaskBase() = default;
+public:
+    virtual TaskStatus status() noexcept = 0;
+    // virtual std::any result() noexcept = 0;
+    virtual void resume() noexcept = 0;
+    virtual void* raw() noexcept = 0;
+    virtual std::coroutine_handle<> handle() noexcept = 0;
+    virtual void setLoop(TaskLoop* loop) noexcept = 0;
+
+    [[nodiscard]] inline bool canResume() {
+        const TaskStatus stat = status();
+        return (stat == TaskStatus::SUSPENDED || stat == TaskStatus::YIELDED);
+    }
+};
+
+template<typename TTask>
+class WrappedTask : public WrappedTaskBase
+{
+private:
+    TTask task_;
+public:
+    constexpr explicit WrappedTask(TTask&& task) noexcept : task_(std::move(task)) {}
+    WrappedTask(const WrappedTask&) = delete;
+    WrappedTask(WrappedTask&&) noexcept = default;
+public:
+    WrappedTask& operator=(const WrappedTask&) = delete;
+    WrappedTask& operator=(WrappedTask&&) noexcept = default;
+public:
+    TaskStatus status() noexcept override { return task_.state().status; }
+    // std::any result() noexcept override
+    // {
+    //     if constexpr (std::is_same_v<typename TTask::result_t, void>) {
+    //         return {};
+    //     }
+    //     else {
+    //         return std::any(task_.state().value);
+    //     }
+    // }
+    void resume() noexcept override { task_.resume(); }
+    void* raw() noexcept override { return &task_; }
+    std::coroutine_handle<> handle() noexcept override { return task_.handle(); }
+    void setLoop(TaskLoop* loop) noexcept override { task_.setLoop(loop); }
+};
+
+template<typename TTask>
+std::unique_ptr<WrappedTask<TTask>> wrapTask(TTask&& task) noexcept
+{
+    return std::make_unique<WrappedTask<TTask>>(std::forward<TTask>(task));
+}
+
+class TaskLoop
+{
+public:
+    MIJIN_DEFINE_FLAG(CanContinue);
+    MIJIN_DEFINE_FLAG(IgnoreWaiting);
+protected:
+    using wrapped_task_t = WrappedTaskBase;
+    using wrapped_task_base_ptr_t = std::unique_ptr<wrapped_task_t>;
+
+    struct StoredTask
+    {
+        wrapped_task_base_ptr_t task;
+        std::function<void(StoredTask&)> setFuture;
+        std::any resultData;
+    };
+    using task_vector_t = std::vector<StoredTask>;
+
+    template<typename TTask>
+    using wrapped_task_ptr_t = std::unique_ptr<WrappedTask<TTask>>;
+public:
+    TaskLoop() = default;
+    TaskLoop(const TaskLoop&) = delete;
+    TaskLoop(TaskLoop&&) = delete;
+    virtual ~TaskLoop() = default;
+
+    TaskLoop& operator=(const TaskLoop&) = delete;
+    TaskLoop& operator=(TaskLoop&&) = delete;
+
+    template<typename TResult>
+    inline FuturePtr<TResult> addTask(TaskBase<TResult> task) noexcept;
+
+    virtual void transferCurrentTask(TaskLoop& otherLoop) noexcept = 0;
+    virtual void addStoredTask(StoredTask&& storedTask) noexcept = 0;
+
+    [[nodiscard]] static TaskLoop& current() noexcept;
+protected:
+    inline TaskStatus tickTask(StoredTask& task) noexcept;
+protected:
+    static inline TaskLoop*& currentLoopStorage() noexcept;
+    template<typename TResult>
+    static inline void setFutureHelper(StoredTask& storedTask) noexcept;
+};
+
+template<typename TResult = void>
+using Task = TaskBase<TResult>;
+
+class SimpleTaskLoop : public TaskLoop
+{
+private:
+    task_vector_t tasks_;
+    task_vector_t newTasks_;
+    task_vector_t::iterator currentTask_;
+    MessageQueue<StoredTask> queuedTasks_;
+    std::thread::id threadId_;
+
+public: // TaskLoop implementation
+    void transferCurrentTask(TaskLoop& otherLoop) noexcept override;
+    void addStoredTask(StoredTask&& storedTask) noexcept override;
+
+public: // public interface
+    [[nodiscard]] constexpr bool empty() const noexcept { return tasks_.empty() && newTasks_.empty(); }
+    inline CanContinue tick() noexcept;
+    inline void runUntilDone(IgnoreWaiting ignoreWaiting = IgnoreWaiting::NO) noexcept;
+private:
+    inline void assertCorrectThread() { MIJIN_ASSERT(threadId_ == std::thread::id() || threadId_ == std::this_thread::get_id(), "Unsafe to TaskLoop from different thread!"); }
+};
+
+class MultiThreadedTaskLoop : public TaskLoop
+{
+private:
+    task_vector_t parkedTasks_; // buffer for tasks that don't fit into readyTasks_
+    MessageQueue<StoredTask> queuedTasks_; // tasks that should be appended to parked tasks
+    MessageQueue<StoredTask> readyTasks_; // task queue to send tasks to a worker thread
+    MessageQueue<StoredTask> returningTasks_; // task that have executed on a worker thread and return for further processing
+    std::jthread managerThread_;
+    std::vector<std::jthread> workerThreads_;
+
+public: // TaskLoop implementation
+    void transferCurrentTask(TaskLoop& otherLoop) noexcept override;
+    void addStoredTask(StoredTask&& storedTask) noexcept override;
+
+public: // public interface
+    void start(std::size_t numWorkerThreads);
+    void stop();
+private: // private stuff
+    void managerThread(std::stop_token stopToken);
+    void workerThread(std::stop_token stopToken, std::size_t workerId);
+
+    static thread_local StoredTask* currentTask_;
+};
+
+//
+// public functions
+//
+
+template<typename TResult>
+TaskBase<TResult>::~TaskBase() noexcept
+{
+    if (handle_)
+    {
+        // handle_.destroy();
+    }
+}
+
+template<typename TResult>
+inline FuturePtr<TResult> TaskLoop::addTask(TaskBase<TResult> task) noexcept
+{
+    MIJIN_ASSERT(!task.getLoop(), "Attempting to add task that already has a loop!");
+    task.setLoop(this);
+
+    auto future = std::make_shared<Future<TResult>>();
+
+    auto setFuture = &setFutureHelper<TResult>;
+
+    // add tasks to a seperate vector first as we might be running another task right now
+    addStoredTask(StoredTask{
+        .task = wrapTask(std::move(task)),
+        .setFuture = setFuture,
+        .resultData = future
+    });
+
+    return future;
+}
+
+inline TaskStatus TaskLoop::tickTask(StoredTask& task) noexcept
+{
+    TaskStatus status = {};
+    do
+    {
+        task.task->resume();
+        status = task.task ? task.task->status() : TaskStatus::WAITING; // no inner task -> task switch context (and will be removed later)
+    }
+    while (status == TaskStatus::RUNNING);
+
+    if (status == TaskStatus::YIELDED || status == TaskStatus::FINISHED)
+    {
+        task.setFuture(task);
+    }
+    return status;
+}
+
+/* static */ inline auto TaskLoop::current() noexcept -> TaskLoop&
+{
+    MIJIN_ASSERT(currentLoopStorage() != nullptr, "Attempting to fetch current loop while no coroutine is running!");
+    return *currentLoopStorage();
+}
+
+/* static */ auto TaskLoop::currentLoopStorage() noexcept -> TaskLoop*&
+{
+    static thread_local TaskLoop* storage = nullptr;
+    return storage;
+}
+
+template<typename TResult>
+/* static */ inline void TaskLoop::setFutureHelper(StoredTask& storedTask) noexcept
+{
+    TaskBase<TResult>& task = *static_cast<TaskBase<TResult>*>(storedTask.task->raw());
+    auto future = std::any_cast<FuturePtr<TResult>>(storedTask.resultData);
+
+    if constexpr (!std::is_same_v<TResult, void>)
+    {
+        MIJIN_ASSERT(!task.state().value.empty(), "Task did not produce a value?");
+        future->set(std::move(task.state().value.get()));
+    }
+    else {
+        future->set();
+    }
+}
+
+inline std::suspend_always switchContext(TaskLoop& taskLoop)
+{
+    TaskLoop& currentTaskLoop = TaskLoop::current();
+    if (&currentTaskLoop == &taskLoop) {
+        return {};
+    }
+    currentTaskLoop.transferCurrentTask(taskLoop);
+    return {};
+}
+
+inline auto SimpleTaskLoop::tick() noexcept -> CanContinue
+{
+    // set current taskloop
+    MIJIN_ASSERT(currentLoopStorage() == nullptr, "Trying to tick a loop from a coroutine, this is not supported.");
+    currentLoopStorage() = this;
+    threadId_ = std::this_thread::get_id();
+
+    // move over all tasks from newTasks
+    for (StoredTask& task : newTasks_)
+    {
+        tasks_.push_back(std::move(task));
+    }
+    newTasks_.clear();
+
+    // also pick up tasks from other threads
+    while(true)
+    {
+        std::optional<StoredTask> task = queuedTasks_.tryPop();
+        if (!task.has_value()) {
+            break;
+        }
+        tasks_.push_back(std::move(*task));
+    }
+
+    // remove any tasks that are finished executing
+    auto it = std::remove_if(tasks_.begin(), tasks_.end(), [](StoredTask& task) {
+        return task.task->status() == TaskStatus::FINISHED;
+    });
+    tasks_.erase(it, tasks_.end());
+
+    CanContinue canContinue = CanContinue::NO;
+    // then execute all tasks that can be executed
+    for (currentTask_ = tasks_.begin(); currentTask_ != tasks_.end(); ++currentTask_)
+    {
+        StoredTask& task = *currentTask_;
+        TaskStatus status = task.task->status();
+        if (status != TaskStatus::SUSPENDED && status != TaskStatus::YIELDED)
+        {
+            MIJIN_ASSERT(status == TaskStatus::WAITING, "Task with invalid status in task list!");
+            continue;
+        }
+
+        status = tickTask(task);
+
+        if (status == TaskStatus::SUSPENDED || status == TaskStatus::YIELDED)
+        {
+            canContinue = CanContinue::YES;
+        }
+    }
+    // reset current loop
+    currentLoopStorage() = nullptr;
+
+    // remove any tasks that have been transferred to another queue
+    it = std::remove_if(tasks_.begin(), tasks_.end(), [](const StoredTask& task) {
+        return task.task == nullptr;
+    });
+    tasks_.erase(it, tasks_.end());
+
+    return canContinue;
+}
+
+inline void SimpleTaskLoop::runUntilDone(IgnoreWaiting ignoreWaiting) noexcept
+{
+    while (!tasks_.empty() || !newTasks_.empty())
+    {
+        const CanContinue canContinue = tick();
+        if (ignoreWaiting && !canContinue)
+        {
+            break;
+        }
+    }
+}
+
+// utility stuff
+
+inline std::suspend_always c_suspend() {
+    return std::suspend_always();
+}
+
+template<template<typename...> typename TCollection, typename TType, typename... TTemplateArgs>
+Task<> c_allDone(const TCollection<FuturePtr<TType>, TTemplateArgs...>& futures)
+{
+    bool allDone = true;
+    do
+    {
+        allDone = true;
+        for (const FuturePtr<TType>& future : futures)
+        {
+            if (future && !future->ready()) {
+                allDone = false;
+                break;
+            }
+        }
+        co_await c_suspend();
+    } while (!allDone);
+}
+
+#if MIJIN_COROUTINE_ENABLE_DEBUG_INFO
+#endif
+}
+
+#endif // MIJIN_ASYNC_COROUTINE_HPP_INCLUDED