Program Listing for File TSQueue.hpp

Return to documentation for file (pennylane_lightning/core/src/utils/cuda_utils/TSQueue.hpp)

#pragma once

#include <condition_variable>
#include <memory>
#include <mutex>
#include <queue>

template <typename T> class TSQueue {
  private:
    mutable std::mutex m;
    std::queue<T> q;
    std::condition_variable cond;

  public:
    TSQueue() = default;

    void push(T data) {
        std::lock_guard<std::mutex> lk(m);
        q.push(std::move(data));
        cond.notify_one();
    }

    void wait_and_pop(T &data) {
        std::unique_lock<std::mutex> lk(m);
        cond.wait(lk, [this] { return !q.empty(); });
        data = std::move(q.front());
        q.pop();
    }
    std::shared_ptr<T> wait_and_pop() {
        std::unique_lock<std::mutex> lk(m);
        cond.wait(lk, [this] { return !q.empty(); });
        std::shared_ptr<T> res(std::make_shared<T>(std::move(q.front())));
        q.pop();
        return res;
    }
    bool try_pop(T &data) {
        std::lock_guard<std::mutex> lk(m);
        if (q.empty()) {
            return false;
        }
        data = std::move(q.front());
        q.pop();
        return true;
    }
    std::shared_ptr<T> try_pop() {
        std::lock_guard<std::mutex> lk(m);
        if (q.empty()) {
            return std::shared_ptr<T>();
        }
        std::shared_ptr<T> res(std::make_shared<T>(std::move(q.front())));
        q.pop();
        return res;
    }
    bool empty() const {
        std::lock_guard<std::mutex> lk(m);
        return q.empty();
    }
};