zeek/src/PriorityQueue.h

// See the file "COPYING" in the main distribution directory for copyright.

#pragma once

#include "zeek/zeek-config.h"

#include <cmath>
#include <cstdint>

namespace zeek::detail {

class PriorityQueue;

class PQ_Element {
public:
    explicit PQ_Element(double t) : time(t) {}
    virtual ~PQ_Element() = default;

    double Time() const { return time; }

    int Offset() const { return offset; }
    void SetOffset(int off) { offset = off; }

    void MinimizeTime() { time = -HUGE_VAL; }

protected:
    PQ_Element() = default;
    double time = 0.0;
    int offset = -1;
};

class PriorityQueue {
public:
    explicit PriorityQueue(int initial_size = 16);
    ~PriorityQueue();

    // Returns the top of queue, or nil if the queue is empty.
    PQ_Element* Top() const {
        if ( heap_size == 0 )
            return nullptr;

        return heap[0];
    }

    // Removes (and returns) top of queue.  Returns nil if the queue
    // is empty.
    PQ_Element* Remove();

    // Removes element e.  Returns e, or nullptr if e wasn't in the queue.
    // Note that e will be modified via MinimizeTime().
    PQ_Element* Remove(PQ_Element* e);

    // Add a new element to the queue.  Returns false on failure (not enough
    // memory to add the element), true on success.
    bool Add(PQ_Element* e);

    int Size() const { return heap_size; }
    int PeakSize() const { return peak_heap_size; }
    uint64_t CumulativeNum() const { return cumulative_num; }

protected:
    bool Resize(int new_size);

    void BubbleUp(int bin);
    void BubbleDown(int bin);

    int Parent(int bin) const { return bin >> 1; }

    int LeftChild(int bin) const { return bin << 1; }

    int RightChild(int bin) const { return LeftChild(bin) + 1; }

    void SetElement(int bin, PQ_Element* e) {
        heap[bin] = e;
        e->SetOffset(bin);
    }

    void Swap(int bin1, int bin2) {
        PQ_Element* t = heap[bin1];
        SetElement(bin1, heap[bin2]);
        SetElement(bin2, t);
    }

    PQ_Element** heap = nullptr;
    int heap_size = 0;
    int peak_heap_size = 0;
    int max_heap_size = 0;
    uint64_t cumulative_num = 0;
};

} // namespace zeek::detail