PriorityStableQueue.hpp

// SPDX-FileCopyrightText: 2025 Vincent Leroy
// SPDX-License-Identifier: MIT
//
// This file is part of dfx.
//
// Licensed under the MIT License. See the LICENSE file in the project root
// for full license information.
 
#pragma once
 
// Standard includes
#include <algorithm>
#include <concepts>
#include <condition_variable>
#include <mutex>
#include <optional>
#include <span>
#include <vector>
 
// Project includes
#include "../CopyMoveControl.hpp"
#include "../MonotonicIdAllocator.hpp"
 
namespace dfx::Utils
{
template<typename Mapper, typename T>
concept PriorityMapper = requires(Mapper m, T const & val)
{
    { m(val) } -> std::convertible_to<int>;
};

template<typename T>
class PriorityStableQueue
{
private:
    struct QueueItem
    {
        T item;
        int priority;
        uint64_t seq;
    };
 
    struct QueueItemCompare
    {
        bool operator()(QueueItem const & a, QueueItem const & b) const noexcept
        {
            // higher priority first
            auto const ap = a.priority;
            auto const bp = b.priority;
            if (ap != bp)
                return ap < bp;
 
            // FIFO among equal priorities: smaller seq first
            return a.seq > b.seq;
        }
    };
 
public:
    using value_type      = T;
    using size_type       = std::size_t;
    using difference_type = std::ptrdiff_t;
    using reference       = value_type &;
    using const_reference = value_type const &;
    using pointer         = value_type *;
    using const_pointer   = value_type const *;
 
public:
    PriorityStableQueue() = default;
 
    DFX_DISABLE_COPY_AND_MOVE(PriorityStableQueue);
 
public:
    void push(T item, int priority = 0);

    void pushRange(std::span<T> items, int priority = 0);

    template<typename Mapper>
    requires PriorityMapper<Mapper, T>
    void pushRange(std::span<T> items, Mapper && priorityMapper);

    std::optional<T> pop(std::stop_token stopToken = {});

    std::optional<T> tryPop();

    std::vector<T> drain();

    void clear();

    bool empty() const;

    size_type size() const;
 
private:
    mutable std::mutex                    _mutex;
    std::vector<QueueItem>                _stableQueue;
    Utils::MonotonicIdAllocator<uint64_t> _seqAllocator;
    std::condition_variable_any           _cvQueue;
};
 
template<typename T>
inline void PriorityStableQueue<T>::push(T item, int priority)
{
    {
        std::lock_guard lock(_mutex);
        _stableQueue.emplace_back(std::move(item), priority, _seqAllocator.next());
        std::push_heap(_stableQueue.begin(), _stableQueue.end(), QueueItemCompare());
    }
 
    _cvQueue.notify_one();
}
 
template<typename T>
inline void PriorityStableQueue<T>::pushRange(std::span<T> items, int priority)
{
    if (items.empty())
        return ;
 
    {
        std::lock_guard lock(_mutex);
 
        for (auto & item : items)
            _stableQueue.emplace_back(std::move(item), priority, _seqAllocator.next());
 
        std::make_heap(_stableQueue.begin(), _stableQueue.end(), QueueItemCompare{});
    }
 
    _cvQueue.notify_all();
}
 
template<typename T>
template<typename Mapper>
requires PriorityMapper<Mapper, T>
inline void PriorityStableQueue<T>::pushRange(std::span<T> items, Mapper && priorityMapper)
{
    if (items.empty())
        return ;
 
    {
        std::lock_guard lock(_mutex);
 
        for (auto & item : items)
            _stableQueue.emplace_back(std::move(item), priorityMapper(item), _seqAllocator.next());
 
        std::make_heap(_stableQueue.begin(), _stableQueue.end(), QueueItemCompare{});
    }
 
    _cvQueue.notify_all();
}
 
template<typename T>
inline std::optional<T> PriorityStableQueue<T>::pop(std::stop_token stopToken)
{
    std::unique_lock lock(_mutex);
    if (!_cvQueue.wait(lock, stopToken, [this] noexcept { return !_stableQueue.empty(); }))
        return std::nullopt;
 
    // std::conditional_variable_any::wait return value, in case a stop is requested,
    // is not always false but the value returned by the predicate at this moment
    // So here we ensure that if a stop was requested, it will be honored asap
    if (stopToken.stop_requested())
        return std::nullopt;
 
    std::pop_heap(_stableQueue.begin(), _stableQueue.end(), QueueItemCompare());
    auto item = std::move(_stableQueue.back().item);
    _stableQueue.pop_back();
 
    return item;
}
 
template<typename T>
inline std::optional<T> PriorityStableQueue<T>::tryPop()
{
    std::lock_guard lock(_mutex);
    if (_stableQueue.empty())
        return std::nullopt;
 
    std::pop_heap(_stableQueue.begin(), _stableQueue.end(), QueueItemCompare());
    auto item = std::move(_stableQueue.back().item);
    _stableQueue.pop_back();
 
    return item;
}
 
template<typename T>
inline std::vector<T> PriorityStableQueue<T>::drain()
{
    std::vector<QueueItem> items;
 
    {
        std::lock_guard lock(_mutex);
        if (_stableQueue.empty())
            return {};
 
        items = std::move(_stableQueue);
        _stableQueue.clear();
    }
 
    // std::sort is faster than std::sort_heap
    std::sort(items.begin(), items.end(), QueueItemCompare());
    // We want the highest priority first
    std::reverse(items.begin(), items.end());
 
    std::vector<T> result;
    result.reserve(items.size());
    for (auto & qi : items)
        result.push_back(std::move(qi.item));
 
    return result;
}
 
template<typename T>
inline void PriorityStableQueue<T>::clear()
{
    std::lock_guard lock(_mutex);
    _stableQueue.clear();
}
 
template<typename T>
inline bool PriorityStableQueue<T>::empty() const
{
    std::lock_guard lock(_mutex);
    return _stableQueue.empty();
}
 
template<typename T>
inline PriorityStableQueue<T>::size_type PriorityStableQueue<T>::size() const
{
    std::lock_guard lock(_mutex);
    return _stableQueue.size();
}
} // !namespace dfx::Utils