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