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 "alloc-util.h"
  33 #include "prioq.h"
  34 #include "util.h"
  35
  36 struct prioq_item {
  37         void *data;
  38         unsigned *idx;
  39 };
  40
  41 struct Prioq {
  42         compare_func_t compare_func;
  43         unsigned n_items, n_allocated;
  44
  45         struct prioq_item *items;
  46 };
  47
  48 Prioq *prioq_new(compare_func_t compare_func) {
  49         Prioq *q;
  50
  51         q = new0(Prioq, 1);
  52         if (!q)
  53                 return q;
  54
  55         q->compare_func = compare_func;
  56         return q;
  57 }
  58
  59 Prioq* prioq_free(Prioq *q) {
  60         if (!q)
  61                 return NULL;
  62
  63         free(q->items);
  64         free(q);
  65
  66         return NULL;
  67 }
  68
  69 int prioq_ensure_allocated(Prioq **q, compare_func_t compare_func) {
  70         assert(q);
  71
  72         if (*q)
  73                 return 0;
  74
  75         *q = prioq_new(compare_func);
  76         if (!*q)
  77                 return -ENOMEM;
  78
  79         return 0;
  80 }
  81
  82 static void swap(Prioq *q, unsigned j, unsigned k) {
  83         void *saved_data;
  84         unsigned *saved_idx;
  85
  86         assert(q);
  87         assert(j < q->n_items);
  88         assert(k < q->n_items);
  89
  90         assert(!q->items[j].idx || *(q->items[j].idx) == j);
  91         assert(!q->items[k].idx || *(q->items[k].idx) == k);
  92
  93         saved_data = q->items[j].data;
  94         saved_idx = q->items[j].idx;
  95         q->items[j].data = q->items[k].data;
  96         q->items[j].idx = q->items[k].idx;
  97         q->items[k].data = saved_data;
  98         q->items[k].idx = saved_idx;
  99
 100         if (q->items[j].idx)
 101                 *q->items[j].idx = j;
 102
 103         if (q->items[k].idx)
 104                 *q->items[k].idx = k;
 105 }
 106
 107 static unsigned shuffle_up(Prioq *q, unsigned idx) {
 108         assert(q);
 109
 110         while (idx > 0) {
 111                 unsigned k;
 112
 113                 k = (idx-1)/2;
 114
 115                 if (q->compare_func(q->items[k].data, q->items[idx].data) <= 0)
 116                         break;
 117
 118                 swap(q, idx, k);
 119                 idx = k;
 120         }
 121
 122         return idx;
 123 }
 124
 125 static unsigned shuffle_down(Prioq *q, unsigned idx) {
 126         assert(q);
 127
 128         for (;;) {
 129                 unsigned j, k, s;
 130
 131                 k = (idx+1)*2; /* right child */
 132                 j = k-1;       /* left child */
 133
 134                 if (j >= q->n_items)
 135                         break;
 136
 137                 if (q->compare_func(q->items[j].data, q->items[idx].data) < 0)
 138
 139                         /* So our left child is smaller than we are, let's
 140                          * remember this fact */
 141                         s = j;
 142                 else
 143                         s = idx;
 144
 145                 if (k < q->n_items &&
 146                     q->compare_func(q->items[k].data, q->items[s].data) < 0)
 147
 148                         /* So our right child is smaller than we are, let's
 149                          * remember this fact */
 150                         s = k;
 151
 152                 /* s now points to the smallest of the three items */
 153
 154                 if (s == idx)
 155                         /* No swap necessary, we're done */
 156                         break;
 157
 158                 swap(q, idx, s);
 159                 idx = s;
 160         }
 161
 162         return idx;
 163 }
 164
 165 int prioq_put(Prioq *q, void *data, unsigned *idx) {
 166         struct prioq_item *i;
 167         unsigned k;
 168
 169         assert(q);
 170
 171         if (q->n_items >= q->n_allocated) {
 172                 unsigned n;
 173                 struct prioq_item *j;
 174
 175                 n = MAX((q->n_items+1) * 2, 16u);
 176                 j = realloc(q->items, sizeof(struct prioq_item) * n);
 177                 if (!j)
 178                         return -ENOMEM;
 179
 180                 q->items = j;
 181                 q->n_allocated = n;
 182         }
 183
 184         k = q->n_items++;
 185         i = q->items + k;
 186         i->data = data;
 187         i->idx = idx;
 188
 189         if (idx)
 190                 *idx = k;
 191
 192         shuffle_up(q, k);
 193
 194         return 0;
 195 }
 196
 197 static void remove_item(Prioq *q, struct prioq_item *i) {
 198         struct prioq_item *l;
 199
 200         assert(q);
 201         assert(i);
 202
 203         l = q->items + q->n_items - 1;
 204
 205         if (i == l)
 206                 /* Last entry, let's just remove it */
 207                 q->n_items--;
 208         else {
 209                 unsigned k;
 210
 211                 /* Not last entry, let's replace the last entry with
 212                  * this one, and reshuffle */
 213
 214                 k = i - q->items;
 215
 216                 i->data = l->data;
 217                 i->idx = l->idx;
 218                 if (i->idx)
 219                         *i->idx = k;
 220                 q->n_items--;
 221
 222                 k = shuffle_down(q, k);
 223                 shuffle_up(q, k);
 224         }
 225 }
 226
 227 _pure_ static struct prioq_item* find_item(Prioq *q, void *data, unsigned *idx) {
 228         struct prioq_item *i;
 229
 230         assert(q);
 231
 232         if (idx) {
 233                 if (*idx == PRIOQ_IDX_NULL ||
 234                     *idx > q->n_items)
 235                         return NULL;
 236
 237                 i = q->items + *idx;
 238                 if (i->data != data)
 239                         return NULL;
 240
 241                 return i;
 242         } else {
 243                 for (i = q->items; i < q->items + q->n_items; i++)
 244                         if (i->data == data)
 245                                 return i;
 246                 return NULL;
 247         }
 248 }
 249
 250 int prioq_remove(Prioq *q, void *data, unsigned *idx) {
 251         struct prioq_item *i;
 252
 253         if (!q)
 254                 return 0;
 255
 256         i = find_item(q, data, idx);
 257         if (!i)
 258                 return 0;
 259
 260         remove_item(q, i);
 261         return 1;
 262 }
 263
 264 int prioq_reshuffle(Prioq *q, void *data, unsigned *idx) {
 265         struct prioq_item *i;
 266         unsigned k;
 267
 268         assert(q);
 269
 270         i = find_item(q, data, idx);
 271         if (!i)
 272                 return 0;
 273
 274         k = i - q->items;
 275         k = shuffle_down(q, k);
 276         shuffle_up(q, k);
 277         return 1;
 278 }
 279
 280 void *prioq_peek(Prioq *q) {
 281
 282         if (!q)
 283                 return NULL;
 284
 285         if (q->n_items <= 0)
 286                 return NULL;
 287
 288         return q->items[0].data;
 289 }
 290
 291 void *prioq_pop(Prioq *q) {
 292         void *data;
 293
 294         if (!q)
 295                 return NULL;
 296
 297         if (q->n_items <= 0)
 298                 return NULL;
 299
 300         data = q->items[0].data;
 301         remove_item(q, q->items);
 302         return data;
 303 }
 304
 305 unsigned prioq_size(Prioq *q) {
 306
 307         if (!q)
 308                 return 0;
 309
 310         return q->n_items;
 311 }
 312
 313 bool prioq_isempty(Prioq *q) {
 314
 315         if (!q)
 316                 return true;
 317
 318         return q->n_items <= 0;
 319 }