subscription.h

#pragma once
 
#include "future.h"
 
namespace coro {
 
 
template<typename T, typename SubscribeFn = function<prepared_coro(promise<T>)> >
class subscription {
public:
 
    static_assert(std::is_constructible_v<function<prepared_coro(promise<T>)>, SubscribeFn>);
 
    using subscribe_fn = SubscribeFn;
 
    using canceled_awaiter = _details::has_value_awaiter<subscription,false>;
    using ret_value = typename future<T>::ret_value;
 
    subscription() = default;
    subscription(SubscribeFn &&fn): _fn(std::move(fn)) {}
    subscription(const SubscribeFn &fn): _fn(fn) {}
    subscription(subscription &&other):_fn(std::move(other._fn)) {}
    subscription &operator=(subscription &&other) {
        if (this != &other) {
            _fn = std::move(other._fn);
        }
        return *this;
    }
 
 
    decltype(auto) get() {return _fut.await_resume();}
 
 
    static constexpr bool await_ready() noexcept {return false;}
 
    std::coroutine_handle<> await_suspend(std::coroutine_handle<> h) {
        return then([h]() ->prepared_coro {return h;}).symmetric_transfer();
    }
 
    ret_value await_resume() {
        return _fut.await_resume();
    }
 
 
    canceled_awaiter operator!() {return this;}
 
    bool has_value() const {return _fut.has_value();}
 
 
 
    template<std::invocable Fn>
    prepared_coro then(Fn &&fn) {
        auto prom = _fut.get_promise();
        _fut.then(std::forward<Fn>(fn));
        return std::forward<SubscribeFn>(_fn)(std::move(prom));
    }
 
 
    template<std::invocable Fn>
    void operator>>(Fn &&fn) {
        then(std::forward<Fn>(fn));
    }
 
    struct lock_guard {
        std::atomic<bool> *unlock = nullptr;
        void operator()(auto) {
            unlock->store(true);
            unlock->notify_all();
        }
    };
 
 
    using lock_ptr = std::unique_ptr<std::remove_reference_t<ret_value>, lock_guard>;
 
    static constexpr lock_ptr init_lock() {return {};}
 
 
    static constexpr lock_ptr init_lock(std::atomic<bool> &unlock) {return lock_ptr{
        //hack: first argument is an pointer, which is ignored, but must not be nullptr
            reinterpret_cast<std::remove_reference_t<ret_value> *>(&unlock), {&unlock}
    };};
 
 
    void lock(lock_ptr &ptr) {
        std::atomic<bool> ready = {false};
        std::atomic<bool> *unlock = nullptr;
        then([&]{
            std::atomic<bool> lk;
            unlock = &lk;
            ready.store(true);
            ready.notify_all();
            lk.wait(false);
        });
        ptr.reset();
        ready.wait(false);
        if (has_value()) {
            ret_value ref = _fut.get();
            ptr = lock_ptr(&ref, {unlock});
        } else {
            unlock->store(true);
            unlock->notify_all();
            ptr = lock_ptr(nullptr, {});
        }
    }
 
protected:
    future<T> _fut;
    SubscribeFn _fn;
 
    future<T> &start() {
        if (!_fut.is_pending()) _fn(_fut.get_promise());
        return _fut;
    }
 
    void wait_internal() {
    }
    friend class _details::wait_awaiter<subscription>;
 
};
 
 
}