AFuture.h

/*
 * AUI Framework - Declarative UI toolkit for modern C++20
 * Copyright (C) 2020-2024 Alex2772 and Contributors
 *
 * SPDX-License-Identifier: MPL-2.0
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */
 
#pragma once
 
#include <exception>
#include <thread>
#include <utility>
#include "AUI/Traits/concepts.h"
#include "AUI/Util/ABitField.h"
#if AUI_COROUTINES
#include <coroutine>
#endif
 
#include <atomic>
#include <functional>
#include <optional>
#include "AConditionVariable.h"
#include "AMutex.h"
#include <AUI/Common/SharedPtrTypes.h>
#include <AUI/Common/AString.h>
#include <AUI/Common/AException.h>
#include <AUI/Logging/ALogger.h>
#include <AUI/Reflect/AReflect.h>
 
class AThreadPool;
 
 
class AInvocationTargetException: public AException {
 
public:
    AInvocationTargetException(const AString& message = {}, std::exception_ptr causedBy = std::current_exception()):
        AException(message, std::move(causedBy), AStacktrace::capture(3)) {}
 
 
    ~AInvocationTargetException() noexcept override = default;
};
 

AUI_ENUM_FLAG(AFutureWait) {
    JUST_WAIT = 0b00,
    ALLOW_STACKFUL_COROUTINES = 0b10,

    ALLOW_TASK_EXECUTION_IF_NOT_PICKED_UP = 0b01,
    DEFAULT = ALLOW_STACKFUL_COROUTINES | ALLOW_TASK_EXECUTION_IF_NOT_PICKED_UP,
};
 
 
namespace aui::impl::future {
    template<typename Inner>
    struct CancellationWrapper {
 
        explicit CancellationWrapper(_unique<Inner> wrapped) : wrapped(std::move(wrapped)) {}
 
        ~CancellationWrapper() {
            wrapped->cancel();
            wrapped->waitForTask();
        }
 
        [[nodiscard]]
        const _unique<Inner>& ptr() noexcept {
            return wrapped;
        }
 
        Inner* operator->() const noexcept {
            return wrapped.get();
        }
 
    private:
        _unique<Inner> wrapped;
    };
 
    template<typename T>
    struct OnSuccessCallback {
        using type = std::function<void(const T& value)>;
    };
 
    template<typename T>
    struct FutureReturnType {
        using type = T&;
    };
 
    template<>
    struct FutureReturnType<void> {
        using type = void;
    };
 
 
    template<>
    struct OnSuccessCallback<void> {
        using type = std::function<void()>;
    };
    template<typename Value = void>
    class Future
    {
    public:
        static constexpr bool isVoid = std::is_same_v<void, Value>;
        using TaskCallback = std::function<Value()>;
 
        using OnSuccessCallback = typename OnSuccessCallback<Value>::type;
 
        struct Inner {
            bool interrupted = false;
            std::conditional_t<isVoid, bool, AOptional<Value>> value;
            AOptional<AInvocationTargetException> exception;
            ASpinlockMutex mutex;
            AConditionVariable cv;
            TaskCallback task;
            OnSuccessCallback onSuccess;
            std::function<void(const AException& exception)> onError;
            _<AAbstractThread> thread;
            bool cancelled = false;
 
            explicit Inner(std::function<Value()> task) noexcept: task(std::move(task)) {
                if constexpr(isVoid) {
                    value = false;
                }
            }
 
            void waitForTask() noexcept {
                std::unique_lock lock(mutex);
                bool rethrowInterrupted = false;
                while ((thread) && !hasResult() && !cancelled) {
                    try {
                        cv.wait(lock);
                    } catch (const AThread::Interrupted&) {
                        rethrowInterrupted = true;
                    }
                }
                if (rethrowInterrupted) {
                    AThread::current()->interrupt();
                }
            }
 
            [[nodiscard]]
            bool isWaitNeeded() noexcept {
                return (thread || !cancelled) && !hasResult();
            }
 
            [[nodiscard]]
            bool hasResult() const noexcept {
                return value || exception || interrupted;
            }
 
            [[nodiscard]]
            bool hasValue() const noexcept {
                return bool(value);
            }
 
            bool setThread(_<AAbstractThread> thr) noexcept {
                std::unique_lock lock(mutex);
                if (cancelled) return true;
                if (thread) return true;
                thread = std::move(thr);
                return false;
            }
 
            void wait(const _weak<CancellationWrapper<Inner>>& innerWeak, ABitField<AFutureWait> flags = AFutureWait::DEFAULT) noexcept;
 
            void cancel() noexcept {
                std::unique_lock lock(mutex);
                if (!cancelled) {
                    cancelled = true;
                    if (thread && !hasResult()) {
                        thread->interrupt();
                    }
                }
            }

            bool tryExecute(const _weak<CancellationWrapper<Inner>>& innerWeak) {
                /*
                 * We should assume that this == nullptr or invalid.
                 * We can assert that this pointer is safe only if we hold at least one shared_ptr.
                 * It allows callers to pass a weak_ptr in any state.
                 */
                if (auto innerCancellation = innerWeak.lock()) {
                    if (value) return false;
                    auto& inner = innerCancellation->ptr();
                    if (inner->setThread(AThread::current())) return false;
                    try {
                        if (task == nullptr) { // task is executed in wait() function
                            return false;
                        }
                        std::unique_lock lock(mutex);
                        if (task == nullptr) { // task is executed in wait() function
                            return false;
                        }
                        auto func = std::move(task);
                        AUI_ASSERT(bool(func));
                        lock.unlock();
                        innerCancellation = nullptr;
                        if constexpr(isVoid) {
                            func();
                            if (auto sharedPtrLock = innerWeak.lock()) {
                                lock.lock();
                                value = true;
                                cv.notify_all();
                                notifyOnSuccessCallback(lock);
 
                                (void)sharedPtrLock; // sharedPtrLock is *used*
                                if (lock.owns_lock()) lock.unlock(); // unlock earlier because destruction of shared_ptr may cause deadlock
                            }
                        } else {
                            auto result = func();
                            if (auto sharedPtrLock = innerWeak.lock()) {
                                lock.lock();
                                value = std::move(result);
                                cv.notify_all();
                                notifyOnSuccessCallback(lock);
 
                                (void)sharedPtrLock; // sharedPtrLock is *used*
                                if (lock.owns_lock()) lock.unlock(); // unlock earlier because destruction of shared_ptr may cause deadlock
                            }
                        }
                    } catch (const AThread::Interrupted&) {
                        if (auto sharedPtrLock = innerWeak.lock()) {
                            inner->reportInterrupted();
                        }
                        throw;
                    } catch (...) {
                        if (auto sharedPtrLock = innerWeak.lock()) {
                            inner->reportException();
                        }
                        return false;
                    } 
                }
                return true;
            }
 
            void reportInterrupted() noexcept {
                std::unique_lock lock(mutex);
                interrupted = true;
                cv.notify_all();
            }
 
            void reportException(std::exception_ptr causedBy = std::current_exception()) noexcept {
                if (cancelled) {
                    return;
                }
                std::unique_lock lock(mutex);
                exception.emplace("exception reported", std::move(causedBy));
                cv.notify_all();
                if (!onError) {
                    return;
                }
                auto localOnError = std::move(onError);
                lock.unlock();
                localOnError(*exception);
            }
 

            void notifyOnSuccessCallback(std::unique_lock<decltype(mutex)>& lock) noexcept {
                AUI_ASSERT(lock.owns_lock());
                if (cancelled) {
                    return;
                }
                if (value && onSuccess) {
                    auto localOnSuccess = std::move(onSuccess);
                    onSuccess = nullptr;
                    lock.unlock();
                    invokeOnSuccessCallback(localOnSuccess);
                }
            }
 
            template<typename F>
            void invokeOnSuccessCallback(F&& f) {
                try {
                    if constexpr (isVoid) {
                        f();
                    } else {
                        f(*value);
                    }
                } catch (const AException& e) {
                    ALogger::err("AFuture") << "AFuture onSuccess thrown an exception: " << e;
                }
            }
 
            template<typename Callback>
            void addOnSuccessCallback(Callback&& callback) {
                if constexpr (isVoid) {
                    if (onSuccess) {
                        onSuccess = [prev = std::move(onSuccess),
                                     callback = std::forward<Callback>(callback)]() mutable {
                            prev();
                            callback();
                        };
                    } else {
                        onSuccess = [callback = std::forward<Callback>(callback)]() mutable { callback(); };
                    }
                } else {
                    if (onSuccess) {
                        onSuccess = [prev = std::move(onSuccess),
                                     callback = std::forward<Callback>(callback)](const Value& v) mutable {
                            prev(v);
                            callback(v);
                        };
                    } else {
                        onSuccess = [callback = std::forward<Callback>(callback)](const Value& v) mutable {
                            callback(v);
                        };
                    }
                }
            }
 
            template<typename Callback>
            void addOnErrorCallback(Callback&& callback) {
                if (onError) {
                    onError = [prev = std::move(onError),
                               callback = std::forward<Callback>(callback)](const AException& v) {
                        prev(v);
                        callback(v);
                    };
                } else {
                    onError = [callback = std::forward<Callback>(callback)](const AException& v) { callback(v); };
                }
            }
        };
 
#if AUI_COROUTINES
        struct CoPromiseType;
#endif
 
    protected:
        _<CancellationWrapper<Inner>> mInner;
 
 
    public:
        Future(TaskCallback task = nullptr): mInner(_new<CancellationWrapper<Inner>>((_unique<Inner>)(new Inner(std::move(task))))) {}
 
        [[nodiscard]]
        const _<CancellationWrapper<Inner>>& inner() const noexcept {
            return mInner;
        }

        [[nodiscard]]
        bool isWaitNeeded() const noexcept {
            return (*mInner)->isWaitNeeded();
        }

        [[nodiscard]]
        bool hasResult() const noexcept {
            return (*mInner)->hasResult();
        }

        [[nodiscard]]
        bool hasValue() const noexcept {
            return (*mInner)->hasValue();
        }
 
        void reportException() const noexcept {
            (*mInner)->reportException();
        }
 
        template<typename Callback>
        void onSuccess(Callback&& callback) const {
            if (hasValue()) { // cheap lookahead
                (*mInner)->invokeOnSuccessCallback(std::forward<Callback>(callback));
                return;
            }
            std::unique_lock lock((*mInner)->mutex);
            if (hasValue()) { // not so cheap, but cheapier than adding the callback to inner
                (*mInner)->invokeOnSuccessCallback(std::forward<Callback>(callback));
                return;
            }
            (*mInner)->addOnSuccessCallback(std::forward<Callback>(callback));
        }
 
        template<aui::invocable<const AException&> Callback>
        void onError(Callback&& callback) const {
            std::unique_lock lock((*mInner)->mutex);
            (*mInner)->addOnErrorCallback(std::forward<Callback>(callback));
        }

        template<aui::invocable Callback>
        void onFinally(Callback&& callback) const {
            std::unique_lock lock((*mInner)->mutex);
            (*mInner)->addOnSuccessCallback([callback](const auto&...) { callback(); });
            (*mInner)->addOnErrorCallback([callback = std::move(callback)](const auto&...) { callback(); });
            (*mInner)->notifyOnSuccessCallback(lock);
        }

        void cancel() const noexcept {
            (*mInner)->cancel();
        }
 
        void reportInterrupted() const {
            (*mInner)->reportInterrupted();
        }

        void wait(AFutureWait flags = AFutureWait::DEFAULT) const {
            (*mInner)->wait(mInner, flags);
            checkForSelfWait();
        }

        typename FutureReturnType<Value>::type get(AFutureWait flags = AFutureWait::DEFAULT) const {
            AThread::interruptionPoint();
 
            (*mInner)->wait(mInner, flags);
 
            AThread::interruptionPoint();
            if ((*mInner)->exception) {
                throw *(*mInner)->exception;
            }
            if ((*mInner)->interrupted) {
                throw AInvocationTargetException("Future execution interrupted");
            }
            checkForSelfWait();
            if constexpr(!isVoid) {
                return *(*mInner)->value;
            }
        }
 

        typename FutureReturnType<Value>::type operator*() const {
            return **const_cast<Future*>(this);
        }

        typename FutureReturnType<Value>::type operator*() {
            return get();
        }

        Value* operator->() const {
            return &operator*();
        }
 
    private:
        void checkForSelfWait() const {
            if (!(*mInner)->hasResult() && AThread::current() == (*mInner)->thread) {
                throw AException("self wait?");
            }
        }
    };
 
}
 

template<typename T = void>
class AFuture final: public aui::impl::future::Future<T> {
private:
    using super = typename aui::impl::future::Future<T>;
 
public:
    using Task = typename super::TaskCallback;
#if AUI_COROUTINES
    using promise_type = typename super::CoPromiseType;
#endif
    using Inner = aui::impl::future::CancellationWrapper<typename super::Inner>;
 
    explicit AFuture(T immediateValue): super() {
        auto& inner = (*super::mInner);
        inner->value = std::move(immediateValue);
    }
    AFuture(Task task = nullptr) noexcept: super(std::move(task)) {}
    ~AFuture() = default;
 
    AFuture(const AFuture&) = default;
    AFuture(AFuture&&) noexcept = default;
 
    AFuture& operator=(const AFuture&) = default;
    AFuture& operator=(AFuture&&) noexcept = default;

    void supplyValue(T v) const noexcept {
        auto& inner = (*super::mInner);
        AUI_ASSERTX(inner->task == nullptr, "task is already provided");
 
        std::unique_lock lock(inner->mutex);
        inner->value = std::move(v);
        inner->cv.notify_all();
        inner->notifyOnSuccessCallback(lock);
    }

    void supplyException(std::exception_ptr causedBy = std::current_exception()) const noexcept {
        auto& inner = (*super::mInner);
        inner->reportException(std::move(causedBy));
    }
 
    AFuture& operator=(std::nullptr_t) noexcept {
        super::mInner = nullptr;
        return *this;
    }
 
    [[nodiscard]]
    bool operator==(const AFuture& r) const noexcept {
        return super::mInner == r.mInner;
    }

    template<aui::invocable<const T&> Callback>
    const AFuture& onSuccess(Callback&& callback) const noexcept {
        super::onSuccess(std::forward<Callback>(callback));
        return *this;
    }

    template<aui::invocable<const AException&> Callback>
    const AFuture& onError(Callback&& callback) const noexcept {
        super::onError(std::forward<Callback>(callback));
        return *this;
    }

    template<aui::invocable Callback>
    const AFuture& onFinally(Callback&& callback) const noexcept {
        super::onFinally(std::forward<Callback>(callback));
        return *this;
    }

    template<aui::invocable<const T&> Callback>
    auto map(Callback&& callback) -> AFuture<decltype(callback(std::declval<T>()))> const {
        AFuture<decltype(callback(std::declval<T>()))> result;
        onSuccess([result, callback = std::forward<Callback>(callback)](const T& v) {
            result.supplyValue(callback(v));
        });
        onError([result](const AException& v) {
            try {
                throw v;
            } catch (...) {
                result.reportException();
            }
        });
        return result;
    }
};
 
template<>
class AFuture<void> final: public aui::impl::future::Future<void> {
private:
    using T = void;
    using super = typename aui::impl::future::Future<T>;
 
public:
    using Task = typename super::TaskCallback;
#if AUI_COROUTINES
    using promise_type = typename super::CoPromiseType;
#endif
    using Inner = decltype(std::declval<super>().inner());
 
    AFuture(Task task = nullptr) noexcept: super(std::move(task)) {}
    ~AFuture() = default;
 
    AFuture(const AFuture&) = default;
    AFuture(AFuture&&) noexcept = default;
 
    AFuture& operator=(const AFuture&) = default;
    AFuture& operator=(AFuture&&) noexcept = default;

    void supplyException(std::exception_ptr causedBy = std::current_exception()) const noexcept {
        auto& inner = (*super::mInner);
        inner->reportException(std::move(causedBy));
    }

    void supplyValue() const noexcept {
        auto& inner = (*super::mInner);
        AUI_ASSERTX(inner->task == nullptr, "task is already provided");
 
        std::unique_lock lock(inner->mutex);
        inner->value = true;
        inner->cv.notify_all();
        inner->notifyOnSuccessCallback(lock);
    }
 
    AFuture& operator=(std::nullptr_t) noexcept {
        super::mInner = nullptr;
        return *this;
    }
 
    [[nodiscard]]
    bool operator==(const AFuture& r) const noexcept {
        return super::mInner == r.mInner;
    }

    template<aui::invocable Callback>
    const AFuture& onSuccess(Callback&& callback) const {
        super::onSuccess(std::forward<Callback>(callback));
        return *this;
    }

    template<aui::invocable<const AException&> Callback>
    const AFuture& onError(Callback&& callback) const {
        super::onError(std::forward<Callback>(callback));
        return *this;
    }

    template<aui::invocable Callback>
    const AFuture& onFinally(Callback&& callback) const {
        super::onFinally(std::forward<Callback>(callback));
        return *this;
    }
};
 
 
#include <AUI/Thread/AThreadPool.h>
#include <AUI/Common/AException.h>
 
template <typename Value>
void aui::impl::future::Future<Value>::Inner::wait(const _weak<CancellationWrapper<Inner>>& innerWeak,
                                                   ABitField<AFutureWait> flags) noexcept {
    if (hasResult()) return; // cheap check
    std::unique_lock lock(mutex);
    try {
        if ((thread || !cancelled) && !hasResult() && flags & AFutureWait::ALLOW_TASK_EXECUTION_IF_NOT_PICKED_UP && task) {
            // task is not have picked up by the threadpool; execute it here
            lock.unlock();
            if (tryExecute(innerWeak)) {
                return;
            }
            lock.lock();
        }
 
        if (flags & AFutureWait::ALLOW_STACKFUL_COROUTINES) {
            if (auto threadPoolWorker = _cast<AThreadPool::Worker>(AThread::current())) {
                if (hasResult()) return; // for sure
                AUI_ASSERT(lock.owns_lock());
                auto callback = [threadPoolWorker](auto&&...) {
                    threadPoolWorker->threadPool().wakeUpAll();
                };
                addOnSuccessCallback(callback);
                addOnErrorCallback(callback);
 
                threadPoolWorker->loop([&] {
                  if (lock.owns_lock())
                    lock.unlock();
                  return !hasResult();
                });
 
                return;
            }
        }
        while ((thread || !cancelled) && !hasResult()) {
            if (thread == AThread::current()) [[unlikely]] {
                // self wait?
                return;
            }
            cv.wait(lock);
        }
    } catch (const AThread::Interrupted& e) {
        e.needRethrow();
    }
}
 
#if AUI_COROUTINES
template<typename Value>
struct aui::impl::future::Future<Value>::CoPromiseType {
    AFuture<Value> future;
    auto initial_suspend() const noexcept
    {
        return std::suspend_never{};
    }
 
    auto final_suspend() const noexcept
    {
        return std::suspend_never{};
    }
    auto unhandled_exception() const noexcept {
        future.supplyException();
    }
 
    const AFuture<Value>& get_return_object() const noexcept {
        return future; 
    }
 
    void return_value(Value v) const noexcept {
        future.supplyValue(std::move(v));
    }
};
 
template<typename T>
auto operator co_await(AFuture<T> future) {
    struct Awaitable {
        AFuture<T> future;
 
        bool await_ready() const noexcept {
            return future.hasResult();
        }
 
        T await_resume() {
            return *future;
        }
 
 
        void await_suspend(std::coroutine_handle<> h)
        {
            future.onSuccess([h](const int&) {
                h.resume();
            });
            
            future.onError([h](const AException&) {
                h.resume();
            });
        }
    };
 
    return Awaitable{ std::move(future) };
}
#endif