src/cpp-common/bt2c/prio-heap.hpp

   1 /*
   2  * Copyright (c) 2011 Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
   3  * Copyright (c) 2023 Philippe Proulx <pproulx@efficios.com>
   4  *
   5  * SPDX-License-Identifier: MIT
   6  */
   7
   8 #ifndef BABELTRACE_CPP_COMMON_BT2C_PRIO_HEAP_HPP
   9 #define BABELTRACE_CPP_COMMON_BT2C_PRIO_HEAP_HPP
  10
  11 #include <functional>
  12 #include <type_traits>
  13 #include <utility>
  14 #include <vector>
  15
  16 #include "common/assert.h"
  17
  18 namespace bt2c {
  19
  20 /*
  21  * A templated C++ version of what used to be the `bt_heap_` C API,
  22  * written by Mathieu Desnoyers, which implements an efficient heap data
  23  * structure.
  24  *
  25  * This implements a static-sized priority heap based on CLRS,
  26  * chapter 6.
  27  *
  28  * This version copies instances of `T` during its operations, so it's
  29  * best to use with small objects such as pointers, integers, and small
  30  * PODs.
  31  *
  32  * `T` must be default-constructible, copy-constructible, and
  33  * copy-assignable.
  34  *
  35  * `CompT` is the type of the callable comparator. It must be possible
  36  * to call an instance `comp` of `CompT` as such:
  37  *
  38  *     comp(a, b)
  39  *
  40  * `comp` accepts two different `const T&` values and returns a value
  41  * contextually convertible to `bool` which must be true if `a` appears
  42  * _after_ `b`.
  43  *
  44  * The benefit of this version over `std::priority_queue` is the
  45  * replaceTop() method which you can call to remove the top (greatest)
  46  * element and then insert a new one immediately afterwards with a
  47  * single heap rebalance.
  48  */
  49 template <typename T, typename CompT = std::greater<T>>
  50 class PrioHeap final
  51 {
  52     static_assert(std::is_default_constructible<T>::value, "`T` is default-constructible.");
  53     static_assert(std::is_copy_constructible<T>::value, "`T` is copy-constructible.");
  54     static_assert(std::is_copy_assignable<T>::value, "`T` is copy-assignable.");
  55
  56 public:
  57     /*
  58      * Builds a priority heap using the comparator `comp` and with an
  59      * initial capacity of `cap` elements.
  60      */
  61     explicit PrioHeap(CompT comp, const std::size_t cap) : _mComp {std::move(comp)}
  62     {
  63         _mElems.reserve(cap);
  64     }
  65
  66     /*
  67      * Builds a priority heap using the comparator `comp` and with an
  68      * initial capacity of zero.
  69      */
  70     explicit PrioHeap(CompT comp) : PrioHeap {std::move(comp), 0}
  71     {
  72     }
  73
  74     /*
  75      * Builds a priority heap using a default comparator and with an
  76      * initial capacity of zero.
  77      */
  78     explicit PrioHeap() : PrioHeap {CompT {}, 0}
  79     {
  80     }
  81
  82     /*
  83      * Number of contained elements.
  84      */
  85     std::size_t len() const noexcept
  86     {
  87         return _mElems.size();
  88     }
  89
  90     /*
  91      * Whether or not this heap is empty.
  92      */
  93     bool isEmpty() const noexcept
  94     {
  95         return _mElems.empty();
  96     }
  97
  98     /*
  99      * Removes all the elements.
 100      */
 101     void clear()
 102     {
 103         _mElems.clear();
 104     }
 105
 106     /*
 107      * Current top (greatest) element (`const` version).
 108      */
 109     const T& top() const noexcept
 110     {
 111         BT_ASSERT_DBG(!this->isEmpty());
 112         this->_validate();
 113         return _mElems.front();
 114     }
 115
 116     /*
 117      * Current top (greatest) element.
 118      */
 119     T& top() noexcept
 120     {
 121         BT_ASSERT_DBG(!this->isEmpty());
 122         this->_validate();
 123         return _mElems.front();
 124     }
 125
 126     /*
 127      * Inserts a copy of the element `elem`.
 128      */
 129     void insert(const T& elem)
 130     {
 131         /* Default-construct the new one */
 132         _mElems.resize(_mElems.size() + 1);
 133
 134         auto pos = this->len() - 1;
 135
 136         while (pos > 0 && this->_gt(elem, this->_parentElem(pos))) {
 137             /* Move parent down until we find the right spot */
 138             _mElems[pos] = this->_parentElem(pos);
 139             pos = this->_parentPos(pos);
 140         }
 141
 142         _mElems[pos] = elem;
 143         this->_validate();
 144     }
 145
 146     /*
 147      * Removes the top (greatest) element.
 148      *
 149      * This heap must not be empty.
 150      */
 151     void removeTop()
 152     {
 153         BT_ASSERT_DBG(!this->isEmpty());
 154
 155         if (_mElems.size() == 1) {
 156             /* Fast path for a single element */
 157             _mElems.clear();
 158             return;
 159         }
 160
 161         /*
 162          * Shrink, replace the current top by the previous last element,
 163          * and heapify.
 164          */
 165         const auto lastElem = _mElems.back();
 166
 167         _mElems.resize(_mElems.size() - 1);
 168         return this->replaceTop(lastElem);
 169     }
 170
 171     /*
 172      * Removes the top (greatest) element, and inserts a copy of `elem`.
 173      *
 174      * Equivalent to using removeTop() and then insert(), but more
 175      * efficient (single rebalance).
 176      *
 177      * This heap must not be empty.
 178      */
 179     void replaceTop(const T& elem)
 180     {
 181         BT_ASSERT_DBG(!this->isEmpty());
 182
 183         /* Replace the current top and heapify */
 184         _mElems[0] = elem;
 185         this->_heapify(0);
 186     }
 187
 188 private:
 189     static std::size_t _parentPos(const std::size_t pos) noexcept
 190     {
 191         return (pos - 1) >> 1;
 192     }
 193
 194     void _heapify(std::size_t pos)
 195     {
 196         while (true) {
 197             std::size_t largestPos;
 198             const auto leftPos = (pos << 1) + 1;
 199
 200             if (leftPos < this->len() && this->_gt(_mElems[leftPos], _mElems[pos])) {
 201                 largestPos = leftPos;
 202             } else {
 203                 largestPos = pos;
 204             }
 205
 206             const auto rightPos = (pos << 1) + 2;
 207
 208             if (rightPos < this->len() && this->_gt(_mElems[rightPos], _mElems[largestPos])) {
 209                 largestPos = rightPos;
 210             }
 211
 212             if (G_UNLIKELY(largestPos == pos)) {
 213                 break;
 214             }
 215
 216             const auto tmpElem = _mElems[pos];
 217
 218             _mElems[pos] = _mElems[largestPos];
 219             _mElems[largestPos] = tmpElem;
 220             pos = largestPos;
 221         }
 222
 223         this->_validate();
 224     }
 225
 226     T& _parentElem(const std::size_t pos) noexcept
 227     {
 228         return _mElems[this->_parentPos(pos)];
 229     }
 230
 231     bool _gt(const T& a, const T& b) const
 232     {
 233         /* Forward to user comparator */
 234         return _mComp(a, b);
 235     }
 236
 237     void _validate() const noexcept
 238     {
 239 #ifdef BT_DEBUG_MODE
 240         if (_mElems.empty()) {
 241             return;
 242         }
 243
 244         for (std::size_t i = 1; i < _mElems.size(); ++i) {
 245             BT_ASSERT_DBG(!this->_gt(_mElems[i], _mElems.front()));
 246         }
 247 #endif /* BT_DEBUG_MODE */
 248     }
 249
 250     CompT _mComp;
 251     std::vector<T> _mElems;
 252 };
 253
 254 } /* namespace bt2c */
 255
 256 #endif /* BABELTRACE_CPP_COMMON_BT2C_PRIO_HEAP_HPP */