mutex.h

#pragma once
#include "prepared_coro.h"
#include "future.h"
 
#include <atomic>
#include <memory>
#include <variant>
 
 
namespace coro {
 
 
class mutex {
public:
 
    class ownership {
    public:
        ownership() = default;
        ownership(ownership &&x):_inst(x._inst) {x._inst = nullptr;}
        ownership &operator=(ownership &&x) {
            if (this != &x) {
                release();
                _inst = x._inst;
                x._inst = nullptr;
            }
            return *this;
        }
        ~ownership() {release();}
        void release() {
            if (_inst) {
                auto tmp = _inst;
                _inst = nullptr;
                tmp->unlock();
            }
        }
 
        operator bool() const {return _inst != nullptr;}
 
 
    protected:
        mutex *_inst = nullptr;
 
        friend class mutex;
 
        ownership(mutex *inst):_inst(inst) {}
    };
 
 
    ownership try_lock() {
        return ownership(try_acquire()?this:nullptr);
    }
 
 
    future<ownership> lock() {
        return [&](auto promise) {
            do_lock(promise.release());
        };
    }
 
    future<ownership> operator co_await() {
        return lock();
    }
 
    ownership lock_sync() {
        return std::move(lock().get());
    }
 
protected:
    class awaiter : public future<ownership> {
    public:
        using future<ownership>::_chain;
        static awaiter *from_future(future<ownership> *x) {
            return static_cast<awaiter *>(x);
        }
    };
 
    static awaiter *locked() {return reinterpret_cast<awaiter *>(0x1);}
 
    std::atomic<awaiter *> _req = {nullptr};
 
    awaiter * _que = nullptr;
 
 
 
    bool try_acquire() {
        awaiter *need = nullptr;
        return _req.compare_exchange_strong(need,locked(), std::memory_order_acquire);
    }
 
 
    promise<ownership>::notify do_lock(future<ownership> *fut_awt) {
        awaiter *awt = awaiter::from_future(fut_awt);
        awaiter *nx = nullptr;
        while (!_req.compare_exchange_weak(nx, awt, std::memory_order_relaxed)) {
            awt->_chain = nx;
        }
 
        if (nx == nullptr) {
            build_queue(_req.exchange(locked(), std::memory_order_acquire), awt);
            return promise<ownership>(fut_awt)(make_ownership());
        }
        return {};
    }
 
 
    void build_queue(awaiter *r, awaiter *stop) {
        while (r != stop) {
            auto n = r->_chain;
            r->_chain = _que;
            _que = r;
            r = awaiter::from_future(n);
        }
    }
 
 
    promise<ownership>::notify unlock() {
        if (!_que) {
            auto lk = locked();
            auto need = lk;
            if (_req.compare_exchange_strong(need, nullptr, std::memory_order_release)) return {};
            build_queue(_req.exchange(lk, std::memory_order_relaxed), lk);
        }
        promise<ownership> ret(_que);
        _que = static_cast<awaiter *>(_que->_chain);
        return ret(make_ownership());
    }
 
 
    ownership make_ownership() {return this;}
 
};
 
namespace _details {
    template<typename T, typename Tuple> struct tuple_add;
    template<typename T, typename ... Us> struct tuple_add<T, std::tuple<Us...> > {
        using type = std::tuple<T, Us...>;
    };
 
    template <typename Tuple> struct make_unique_types;
    template <typename T, typename... Us> struct make_unique_types<std::tuple<T, Us...> > {
        using type = std::conditional_t<
                (std::is_same_v<T, Us> || ...),
                typename make_unique_types<std::tuple<Us ...> >::type,
                typename tuple_add<T, typename make_unique_types<std::tuple<Us...> >::type>::type>;
    };
 
    template <> struct make_unique_types< std::tuple<> > {
        using type = std::tuple<>;
    };
 
    template<typename Tuple> struct make_variant_from_tuple;
    template<typename ... Ts> struct make_variant_from_tuple<std::tuple<Ts...> > {
        using type = std::variant<Ts...>;
    };
    template<typename ... Ts> using variant_of_t = typename make_variant_from_tuple<
            typename make_unique_types<std::tuple<Ts...> >::type>::type;
}
 
 
template<typename ... Mutexes>
class lock: public future<std::tuple<decltype(std::declval<Mutexes>().try_lock())...> > {
public:
 
    using ownership = std::tuple<decltype(std::declval<Mutexes>().try_lock())...>;
    using future_variant = _details::variant_of_t<std::monostate,
            decltype(std::declval<Mutexes>().lock())...>;
 
    lock(Mutexes & ... lst):_mxlist(lst...),_prom(this->get_promise()) {
        ownership ownlist;
        if (finish_lock(-1, ownlist)) {
            _prom(std::move(ownlist));
        } //if failed, ownership is released and so all held locks
    }
 
 
protected:
    std::tuple<Mutexes &...> _mxlist;
    promise<ownership> _prom;
    future_variant _fut;
 
 
    template<int idx>
    void on_lock() {
        auto &mx = std::get<idx>(_mxlist);
        using LkType = decltype(mx.lock());
        ownership ownlist;
        std::get<idx>(ownlist) = std::get<LkType>(_fut).await_resume();
        if (finish_lock(idx,ownlist)) {
            _prom(std::move(ownlist));
        }
    }
 
 
    template<int idx = 0>
    bool finish_lock(int skip, ownership &ownlist) {
        //idx out of range? success
        if constexpr(idx >= std::tuple_size_v<ownership>) {
            return true;
        } else {
            //idx == skip - skip it
            if (idx == skip) {
                return finish_lock<idx+1>(skip,ownlist);
            } else {
                //retrieve ownership
                auto &own = std::get<idx>(ownlist);
                //retrieve lock
                auto &mx = std::get<idx>(_mxlist);
                //try to lock
                own = std::get<idx>(_mxlist).try_lock();
                //we success
                if (own) {
                    //continue by next lock
                    return finish_lock<idx+1>(skip,ownlist);
                } else {
                    //failure
                    using LkType = decltype(mx.lock());
                    //construct future object
                    _fut.template emplace<LkType>();
                    LkType &fut = std::get<LkType>(_fut);
                    //acquire lock
                    fut << [&]{return mx.lock();};
                    //install callback, if failed
                    if (fut.set_callback([this]{ on_lock<idx>();}) == false) {
                        //we got ownership
                        own = fut.get();
                        //continue locking
                        return finish_lock<idx+1>(skip,ownlist);
                    }
                    //return failure
                    return false;
                }
            }
        }
    }
};
 
 
}