src/basic/prioq.c

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