1 /* SPDX-License-Identifier: LGPL-2.1-or-later */
2
3 /*
4 * Priority Queue
5 * The prioq object implements a priority queue. That is, it orders objects by
6 * their priority and allows O(1) access to the object with the highest
7 * priority. Insertion and removal are Θ(log n). Optionally, the caller can
8 * provide a pointer to an index which will be kept up-to-date by the prioq.
9 *
10 * The underlying algorithm used in this implementation is a Heap.
11 */
12
13 #include <errno.h>
14 #include <stdlib.h>
15
16 #include "alloc-util.h"
17 #include "hashmap.h"
18 #include "prioq.h"
19
20 struct prioq_item {
21 void *data;
22 unsigned *idx;
23 };
24
25 struct Prioq {
26 compare_func_t compare_func;
27 unsigned n_items, n_allocated;
28
29 struct prioq_item *items;
30 };
31
prioq_new(compare_func_t compare_func)32 Prioq *prioq_new(compare_func_t compare_func) {
33 Prioq *q;
34
35 q = new(Prioq, 1);
36 if (!q)
37 return q;
38
39 *q = (Prioq) {
40 .compare_func = compare_func,
41 };
42
43 return q;
44 }
45
prioq_free(Prioq * q)46 Prioq* prioq_free(Prioq *q) {
47 if (!q)
48 return NULL;
49
50 free(q->items);
51 return mfree(q);
52 }
53
prioq_ensure_allocated(Prioq ** q,compare_func_t compare_func)54 int prioq_ensure_allocated(Prioq **q, compare_func_t compare_func) {
55 assert(q);
56
57 if (*q)
58 return 0;
59
60 *q = prioq_new(compare_func);
61 if (!*q)
62 return -ENOMEM;
63
64 return 0;
65 }
66
swap(Prioq * q,unsigned j,unsigned k)67 static void swap(Prioq *q, unsigned j, unsigned k) {
68 assert(q);
69 assert(j < q->n_items);
70 assert(k < q->n_items);
71
72 assert(!q->items[j].idx || *(q->items[j].idx) == j);
73 assert(!q->items[k].idx || *(q->items[k].idx) == k);
74
75 SWAP_TWO(q->items[j].data, q->items[k].data);
76 SWAP_TWO(q->items[j].idx, q->items[k].idx);
77
78 if (q->items[j].idx)
79 *q->items[j].idx = j;
80
81 if (q->items[k].idx)
82 *q->items[k].idx = k;
83 }
84
shuffle_up(Prioq * q,unsigned idx)85 static unsigned shuffle_up(Prioq *q, unsigned idx) {
86 assert(q);
87 assert(idx < q->n_items);
88
89 while (idx > 0) {
90 unsigned k;
91
92 k = (idx-1)/2;
93
94 if (q->compare_func(q->items[k].data, q->items[idx].data) <= 0)
95 break;
96
97 swap(q, idx, k);
98 idx = k;
99 }
100
101 return idx;
102 }
103
shuffle_down(Prioq * q,unsigned idx)104 static unsigned shuffle_down(Prioq *q, unsigned idx) {
105 assert(q);
106
107 for (;;) {
108 unsigned j, k, s;
109
110 k = (idx+1)*2; /* right child */
111 j = k-1; /* left child */
112
113 if (j >= q->n_items)
114 break;
115
116 if (q->compare_func(q->items[j].data, q->items[idx].data) < 0)
117
118 /* So our left child is smaller than we are, let's
119 * remember this fact */
120 s = j;
121 else
122 s = idx;
123
124 if (k < q->n_items &&
125 q->compare_func(q->items[k].data, q->items[s].data) < 0)
126
127 /* So our right child is smaller than we are, let's
128 * remember this fact */
129 s = k;
130
131 /* s now points to the smallest of the three items */
132
133 if (s == idx)
134 /* No swap necessary, we're done */
135 break;
136
137 swap(q, idx, s);
138 idx = s;
139 }
140
141 return idx;
142 }
143
prioq_put(Prioq * q,void * data,unsigned * idx)144 int prioq_put(Prioq *q, void *data, unsigned *idx) {
145 struct prioq_item *i;
146 unsigned k;
147
148 assert(q);
149
150 if (q->n_items >= q->n_allocated) {
151 unsigned n;
152 struct prioq_item *j;
153
154 n = MAX((q->n_items+1) * 2, 16u);
155 j = reallocarray(q->items, n, sizeof(struct prioq_item));
156 if (!j)
157 return -ENOMEM;
158
159 q->items = j;
160 q->n_allocated = n;
161 }
162
163 k = q->n_items++;
164 i = q->items + k;
165 i->data = data;
166 i->idx = idx;
167
168 if (idx)
169 *idx = k;
170
171 shuffle_up(q, k);
172
173 return 0;
174 }
175
prioq_ensure_put(Prioq ** q,compare_func_t compare_func,void * data,unsigned * idx)176 int prioq_ensure_put(Prioq **q, compare_func_t compare_func, void *data, unsigned *idx) {
177 int r;
178
179 r = prioq_ensure_allocated(q, compare_func);
180 if (r < 0)
181 return r;
182
183 return prioq_put(*q, data, idx);
184 }
185
remove_item(Prioq * q,struct prioq_item * i)186 static void remove_item(Prioq *q, struct prioq_item *i) {
187 struct prioq_item *l;
188
189 assert(q);
190 assert(i);
191
192 l = q->items + q->n_items - 1;
193
194 if (i == l)
195 /* Last entry, let's just remove it */
196 q->n_items--;
197 else {
198 unsigned k;
199
200 /* Not last entry, let's replace the last entry with
201 * this one, and reshuffle */
202
203 k = i - q->items;
204
205 i->data = l->data;
206 i->idx = l->idx;
207 if (i->idx)
208 *i->idx = k;
209 q->n_items--;
210
211 k = shuffle_down(q, k);
212 shuffle_up(q, k);
213 }
214 }
215
find_item(Prioq * q,void * data,unsigned * idx)216 _pure_ static struct prioq_item* find_item(Prioq *q, void *data, unsigned *idx) {
217 struct prioq_item *i;
218
219 assert(q);
220
221 if (q->n_items <= 0)
222 return NULL;
223
224 if (idx) {
225 if (*idx == PRIOQ_IDX_NULL ||
226 *idx >= q->n_items)
227 return NULL;
228
229 i = q->items + *idx;
230 if (i->data != data)
231 return NULL;
232
233 return i;
234 } else {
235 for (i = q->items; i < q->items + q->n_items; i++)
236 if (i->data == data)
237 return i;
238 return NULL;
239 }
240 }
241
prioq_remove(Prioq * q,void * data,unsigned * idx)242 int prioq_remove(Prioq *q, void *data, unsigned *idx) {
243 struct prioq_item *i;
244
245 if (!q)
246 return 0;
247
248 i = find_item(q, data, idx);
249 if (!i)
250 return 0;
251
252 remove_item(q, i);
253 return 1;
254 }
255
prioq_reshuffle(Prioq * q,void * data,unsigned * idx)256 int prioq_reshuffle(Prioq *q, void *data, unsigned *idx) {
257 struct prioq_item *i;
258 unsigned k;
259
260 assert(q);
261
262 i = find_item(q, data, idx);
263 if (!i)
264 return 0;
265
266 k = i - q->items;
267 k = shuffle_down(q, k);
268 shuffle_up(q, k);
269 return 1;
270 }
271
prioq_peek_by_index(Prioq * q,unsigned idx)272 void *prioq_peek_by_index(Prioq *q, unsigned idx) {
273 if (!q)
274 return NULL;
275
276 if (idx >= q->n_items)
277 return NULL;
278
279 return q->items[idx].data;
280 }
281
prioq_pop(Prioq * q)282 void *prioq_pop(Prioq *q) {
283 void *data;
284
285 if (!q)
286 return NULL;
287
288 if (q->n_items <= 0)
289 return NULL;
290
291 data = q->items[0].data;
292 remove_item(q, q->items);
293 return data;
294 }
295
prioq_size(Prioq * q)296 unsigned prioq_size(Prioq *q) {
297
298 if (!q)
299 return 0;
300
301 return q->n_items;
302 }
303
prioq_isempty(Prioq * q)304 bool prioq_isempty(Prioq *q) {
305
306 if (!q)
307 return true;
308
309 return q->n_items <= 0;
310 }
311