Commit | Line | Data |
---|---|---|
ef5ebbac PP |
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 */ |