QXmppTask.h

// SPDX-FileCopyrightText: 2022 Linus Jahn <lnj@kaidan.im>
// SPDX-FileCopyrightText: 2022 Jonah Brüchert <jbb@kaidan.im>
//
// SPDX-License-Identifier: LGPL-2.1-or-later

#ifndef QXMPPTASK_H
#define QXMPPTASK_H

#include "QXmppGlobal.h"

#include <coroutine>
#include <memory>
#include <optional>

#include <QFuture>
#include <QPointer>

#if QXMPP_DEPRECATED_SINCE(1, 11)
#include <functional>
#endif

namespace QXmpp::Private {

template<typename T>
struct TaskData {
    std::conditional_t<std::is_void_v<T>, std::monostate, std::optional<T>> result;
    std::coroutine_handle<> handle;
    QPointer<const QObject> context;
    bool finished = false;
    bool cancelled = false;
    bool hasContext = false;
    uint8_t promiseCount = 1;

    ~TaskData()
    {
        Q_ASSERT(promiseCount == 0);
    }
};

template<typename T>
struct ConstRefOrVoidHelper {
    using Type = const T &;
};
template<>
struct ConstRefOrVoidHelper<void> {
    using Type = void;
};

template<typename T>
using ConstRefOrVoid = ConstRefOrVoidHelper<T>::Type;

template<typename Continuation, typename T>
struct InvokeContinuationResultHelper {
    using Type = std::invoke_result_t<Continuation, T &&>;
};
template<typename Continuation>
struct InvokeContinuationResultHelper<Continuation, void> {
    using Type = std::invoke_result_t<Continuation>;
};

template<typename Continuation, typename T>
using InvokeContinuationResult = InvokeContinuationResultHelper<Continuation, T>::Type;

}  // namespace QXmpp::Private

template<typename T>
class QXmppTask;

template<typename T>
class QXmppPromise
{
    using Task = QXmppTask<T>;
    using SharedData = QXmpp::Private::TaskData<T>;
    using SharedDataPtr = std::shared_ptr<SharedData>;

    struct InlineData {
        Task *task = nullptr;
        QPointer<const QObject> context;
        std::coroutine_handle<> handle;
        bool cancelled = false;
        bool hasContext = false;
    };

public:
    QXmppPromise() : data(InlineData()) { }
    [[deprecated]]
    QXmppPromise(const QXmppPromise<T> &p)
    {
        p.detachData();
        data = p.data;
        sharedData().promiseCount += 1;
    }
    QXmppPromise(QXmppPromise<T> &&p)
    {
        std::swap(data, p.data);
        if (!shared()) {
            if (auto *task = inlineData().task) {
                task->setPromise(this);
            }
        }
    }
    ~QXmppPromise()
    {
        if (shared()) {
            sharedData().promiseCount -= 1;

            // cancel coroutine if any
            if (sharedData().promiseCount == 0) {
                if (auto handle = sharedData().handle) {
                    sharedData().handle = nullptr;
                    handle.destroy();
                }
            }
        } else {
            if (auto *task = inlineData().task) {
                task->setPromise(nullptr);
            }
            // cancel coroutine if any
            if (inlineData().handle) {
                inlineData().handle.destroy();
            }
        }
    }

    [[deprecated]]
    QXmppPromise<T> &operator=(const QXmppPromise<T> &p)
    {
        if (shared()) {
            sharedData().promiseCount -= 1;
        }
        p.detachData();
        data = p.data;
        if (shared()) {
            sharedData().promiseCount += 1;
        }
        return *this;
    }
    QXmppPromise<T> &operator=(QXmppPromise<T> &&p)
    {
        std::swap(data, p.data);
        if (!shared()) {
            if (auto *task = inlineData().task) {
                task->setPromise(this);
            }
        }
        return *this;
    }

    QXmppTask<T> task()
    {
        if (!shared()) {
            if (inlineData().task == nullptr) {
                return Task { this };
            } else {
                detachData();
            }
        }
        return Task { std::get<SharedDataPtr>(data) };
    }

    void finish()
        requires(std::is_void_v<T>)
    {
        if (shared()) {
            sharedData().finished = true;
        } else {
            if (auto *task = inlineData().task) {
                task->inlineData().finished = true;
            } else {
                // finish called without generating task
                detachData();
                sharedData().finished = true;
            }
        }
        invokeHandle();
    }

    template<typename U>
    void finish(U &&value)
        requires(!std::is_void_v<T>)
    {
        if (shared()) {
            sharedData().finished = true;
            sharedData().result.emplace(std::forward<U>(value));
        } else {
            if (auto *task = inlineData().task) {
                inlineData().task->inlineData().finished = true;
                inlineData().task->inlineData().result.emplace(std::forward<U>(value));
            } else {
                // finish called without generating task
                detachData();
                sharedData().finished = true;
                sharedData().result.emplace(std::forward<U>(value));
            }
        }
        invokeHandle();
    }

    bool cancelled() const
    {
        return shared() ? sharedData().cancelled : inlineData().cancelled;
    }

private:
    friend class QXmppTask<T>;

    bool shared() const { return std::holds_alternative<SharedDataPtr>(data); }
    InlineData &inlineData()
    {
        Q_ASSERT(!shared());
        return std::get<InlineData>(data);
    }
    const InlineData &inlineData() const
    {
        Q_ASSERT(!shared());
        return std::get<InlineData>(data);
    }
    SharedData &sharedData()
    {
        Q_ASSERT(shared());
        return *std::get<SharedDataPtr>(data);
    }
    const SharedData &sharedData() const
    {
        Q_ASSERT(shared());
        return *std::get<SharedDataPtr>(data);
    }

    bool contextAlive() const
    {
        if (shared()) {
            return sharedData().context != nullptr || !sharedData().hasContext;
        } else {
            return inlineData().context != nullptr || !inlineData().hasContext;
        }
    }

    void invokeHandle()
    {
        auto &handleRef = shared() ? sharedData().handle : inlineData().handle;
        if (auto handle = handleRef) {
            handleRef = nullptr;
            if (contextAlive()) {
                handle.resume();
            } else {
                handle.destroy();
            }
        }
    }

    // Moves data from QXmppTask object to shared_ptr that can be accessed by multiple tasks and
    // multiple promises.
    // Historically required because task and promise must be copyable.
    void detachData() const
    {
        if (shared()) {
            return;
        }

        if (inlineData().task != nullptr) {
            auto &taskData = inlineData().task->inlineData();

            auto sharedData = std::make_shared<SharedData>(
                std::move(taskData.result),
                inlineData().handle,
                inlineData().context,
                taskData.finished,
                inlineData().cancelled,
                inlineData().hasContext,
                1);
            inlineData().task->data = sharedData;
            data = std::move(sharedData);
        } else {
            data = std::make_shared<SharedData>();
        }
    }

    mutable std::variant<InlineData, SharedDataPtr> data;
};

template<typename T>
class QXmppTask
{
    using Task = QXmppTask<T>;
    using SharedData = QXmpp::Private::TaskData<T>;
    using SharedDataPtr = std::shared_ptr<SharedData>;

    struct InlineData {
        QXmppPromise<T> *promise = nullptr;
        std::conditional_t<std::is_void_v<T>, std::monostate, std::optional<T>> result;
        bool finished = false;
    };

public:
    QXmppTask(QXmppTask &&t) : data(InlineData {})
    {
        std::swap(data, t.data);
        if (!shared()) {
            if (auto *p = inlineData().promise) {
                p->inlineData().task = this;
            }
        }
    }
    QXmppTask(const QXmppTask &) = delete;
    ~QXmppTask()
    {
        if (!shared()) {
            if (auto *p = inlineData().promise) {
                p->inlineData().task = nullptr;
            }
        }
    }

    QXmppTask &operator=(QXmppTask &&t) noexcept
    {
        // unregister with old promise
        if (!shared()) {
            if (auto *p = inlineData().promise) {
                p->inlineData().task = nullptr;
                // to swap nullptr into `t` in next step
                inlineData().promise = nullptr;
            }
        }
        std::swap(data, t.data);
        // register with new promise
        if (!shared()) {
            if (auto *p = inlineData().promise) {
                p->inlineData().task = this;
            }
        }
        return *this;
    }
    QXmppTask &operator=(const QXmppTask &) = delete;

    bool await_ready() const noexcept { return isFinished(); }
    void await_suspend(std::coroutine_handle<> handle)
    {
        auto replace = [](auto &var, auto newValue) {
            if (var) {
                var.destroy();
            }
            var = newValue;
        };

        if (shared()) {
            if (sharedData().promiseCount > 0 && !sharedData().cancelled) {
                replace(sharedData().handle, handle);
            } else {
                handle.destroy();
            }
        } else {
            if (auto *p = inlineData().promise; p && !p->cancelled()) {
                replace(p->inlineData().handle, handle);
            } else {
                handle.destroy();
            }
        }
    }
    auto await_resume()
    {
        if constexpr (!std::is_void_v<T>) {
            return takeResult();
        }
    }

    template<typename Continuation>
    auto then(const QObject *context, Continuation continuation)
        -> QXmppTask<QXmpp::Private::InvokeContinuationResult<Continuation, T>>
    {
        using Result = QXmpp::Private::InvokeContinuationResult<Continuation, T>;
        QXmppTask<T> task = std::move(*this);
        if constexpr (std::is_void_v<T>) {
            co_await task.withContext(context);
            if constexpr (std::is_void_v<Result>) {
                continuation();
            } else {
                co_return continuation();
            }
        } else {
            if constexpr (std::is_void_v<Result>) {
                continuation(co_await task.withContext(context));
            } else {
                co_return continuation(co_await task.withContext(context));
            }
        }
    }

    QXmppTask<T> &withContext(const QObject *c)
    {
        if (shared()) {
            if (!sharedData().finished) {
                sharedData().context = c;
                sharedData().hasContext = true;
            }
        } else {
            if (auto *p = inlineData().promise) {
                p->inlineData().context = c;
                p->inlineData().hasContext = true;
            }
        }
        return *this;
    }

    void cancel()
    {
        if (shared()) {
            sharedData().cancelled = true;
            if (auto handle = sharedData().handle) {
                sharedData().handle = nullptr;
                handle.destroy();
            }
        } else {
            if (auto *p = inlineData().promise) {
                p->inlineData().cancelled = true;
                if (auto handle = p->inlineData().handle) {
                    p->inlineData().handle = nullptr;
                    handle.destroy();
                }
            }
        }
    }

    [[nodiscard]]
    bool isFinished() const
    {
        return shared() ? sharedData().finished : inlineData().finished;
    }

    [[nodiscard]]
    bool hasResult() const
        requires(!std::is_void_v<T>)
    {
        return shared() ? sharedData().result.has_value() : inlineData().result.has_value();
    }

    [[nodiscard]]
    QXmpp::Private::ConstRefOrVoid<T> result() const
        requires(!std::is_void_v<T>)
    {
        Q_ASSERT(isFinished());
        Q_ASSERT(hasResult());
        return shared() ? sharedData().result.value() : inlineData().result.value();
    }

    [[nodiscard]]
    T takeResult()
        requires(!std::is_void_v<T>)
    {
        Q_ASSERT(isFinished());
        Q_ASSERT(hasResult());
        auto &result = shared() ? sharedData().result : inlineData().result;

        auto value = std::move(*result);
        result.reset();
        return value;
    }

    [[nodiscard]]
    QFuture<T> toFuture(const QObject *context)
    {
        QFutureInterface<T> interface;

        if constexpr (std::is_same_v<T, void>) {
            then(context, [interface]() mutable {
                interface.reportFinished();
            });
        } else {
            then(context, [interface](T &&val) mutable {
                interface.reportResult(val);
                interface.reportFinished();
            });
        }

        return interface.future();
    }

private:
    friend class QXmppPromise<T>;

    explicit QXmppTask(QXmppPromise<T> *p) : data(InlineData {})
    {
        inlineData().promise = p;

        Q_ASSERT(p->inlineData().task == nullptr);
        p->inlineData().task = this;
    }
    explicit QXmppTask(SharedDataPtr data) : data(std::move(data)) { }

    bool shared() const { return std::holds_alternative<SharedDataPtr>(data); }
    InlineData &inlineData()
    {
        Q_ASSERT(!shared());
        return std::get<InlineData>(data);
    }
    const InlineData &inlineData() const
    {
        Q_ASSERT(!shared());
        return std::get<InlineData>(data);
    }
    SharedData &sharedData()
    {
        Q_ASSERT(shared());
        return *std::get<SharedDataPtr>(data);
    }
    const SharedData &sharedData() const
    {
        Q_ASSERT(shared());
        return *std::get<SharedDataPtr>(data);
    }

    void setPromise(QXmppPromise<T> *p)
    {
        inlineData().promise = p;
    }

    std::variant<InlineData, SharedDataPtr> data;
};

namespace std {

template<typename T, typename... Args>
struct coroutine_traits<QXmppTask<T>, Args...> {
    struct promise_type {
        QXmppPromise<T> p;

        QXmppTask<T> get_return_object() { return p.task(); }
        std::suspend_never initial_suspend() noexcept { return {}; }
        std::suspend_never final_suspend() noexcept { return {}; }

        void unhandled_exception()
        {
            // exception handling currently not supported
            throw std::current_exception();
        }

        void return_value(T value) { p.finish(std::move(value)); }
    };
};

template<typename... Args>
struct coroutine_traits<QXmppTask<void>, Args...> {
    struct promise_type {
        QXmppPromise<void> p;

        QXmppTask<void> get_return_object() { return p.task(); }
        std::suspend_never initial_suspend() noexcept { return {}; }
        std::suspend_never final_suspend() noexcept { return {}; }

        void unhandled_exception()
        {
            // exception handling currently not supported
            throw std::current_exception();
        }

        void return_void() { p.finish(); }
    };
};

}  // namespace std

namespace QXmpp {

namespace Private {

template<typename T>
struct IsTaskHelper {
    constexpr static bool Value = false;
};
template<typename T>
struct IsTaskHelper<QXmppTask<T>> {
    using Type = T;
    constexpr static bool Value = true;
};

}  // namespace Private

template<typename T>
concept IsTask = Private::IsTaskHelper<T>::Value;

template<IsTask T>
using TaskValueType = typename Private::IsTaskHelper<T>::Type;

}  // namespace QXmpp

#endif  // QXMPPTASK_H