collector.hpp

#pragma once
 
#include <iterator>
#include <mutex>
#include <utility>
#include <vector>
 
namespace wmtk::threading {
template <typename T>
class collector
{
    std::vector<T> m_data;
    mutable std::mutex m_mutex;
 
public:
    using value_type = T;
    using iterator = typename std::vector<T>::iterator;
    using const_iterator = typename std::vector<T>::const_iterator;
    using reference = typename std::vector<T>::reference;
    using const_reference = typename std::vector<T>::const_reference;
 
    collector() = default;
    explicit collector(std::size_t n)
        : m_data(n)
    {}
    collector(std::size_t n, const T& value)
        : m_data(n, value)
    {}
 
    collector(const collector& o) { m_data = o.m_data; }
    collector& operator=(const collector& o)
    {
        m_data = o.m_data;
        return *this;
    }
 
    void push_back(const T& v)
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_data.push_back(v);
    }
    void push_back(T&& v)
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_data.push_back(std::move(v));
    }
 
    template <typename... Args>
    void emplace_back(Args&&... args)
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_data.emplace_back(std::forward<Args>(args)...);
    }
 
    void append(const std::vector<T>& v)
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_data.insert(m_data.end(), v.begin(), v.end());
    }
    void append(std::vector<T>&& v)
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        if (m_data.empty()) {
            m_data = std::move(v);
        } else {
            m_data.insert(
                m_data.end(),
                std::make_move_iterator(v.begin()),
                std::make_move_iterator(v.end()));
        }
    }
 
    void resize(std::size_t n)
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_data.resize(n);
    }
    void resize(std::size_t n, const T& value)
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_data.resize(n, value);
    }
    void clear()
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_data.clear();
    }
    void reserve(std::size_t n)
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_data.reserve(n);
    }
 
    const std::vector<T>& data() const { return m_data; }
 
    std::vector<T> take() { return std::exchange(m_data, std::vector<T>()); }
 
    const_reference operator[](std::size_t i) const { return m_data[i]; }
    // reference operator[](std::size_t i) { return m_data[i]; }
 
    std::size_t size() const { return m_data.size(); }
    bool empty() const { return m_data.empty(); }
 
    // iterator begin() { return m_data.begin(); }
    // iterator end() { return m_data.end(); }
    const_iterator begin() const { return m_data.begin(); }
    const_iterator end() const { return m_data.end(); }
};
 
} // namespace wmtk::threading