src/lib/prio-heap/prio-heap.c

   1 /*
   2  * Static-sized priority heap containing pointers. Based on CLRS,
   3  * chapter 6.
   4  *
   5  * Copyright 2011 - Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
   6  *
   7  * Permission is hereby granted, free of charge, to any person obtaining a copy
   8  * of this software and associated documentation files (the "Software"), to deal
   9  * in the Software without restriction, including without limitation the rights
  10  * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  11  * copies of the Software, and to permit persons to whom the Software is
  12  * furnished to do so, subject to the following conditions:
  13  *
  14  * The above copyright notice and this permission notice shall be included in
  15  * all copies or substantial portions of the Software.
  16  *
  17  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  18  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  19  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  20  * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  21  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  22  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  23  * SOFTWARE.
  24  */
  25
  26 #include "common/macros.h"
  27 #include "common/assert.h"
  28 #include <errno.h>
  29 #include <stdlib.h>
  30 #include <string.h>
  31
  32 #include "prio-heap.h"
  33
  34 #ifdef DEBUG_HEAP
  35 void check_heap(const struct ptr_heap *heap)
  36 {
  37         size_t i;
  38
  39         if (!heap->len)
  40                 return;
  41
  42         for (i = 1; i < heap->len; i++)
  43                 BT_ASSERT_DBG(!heap->gt(heap->ptrs[i], heap->ptrs[0]));
  44 }
  45 #endif
  46
  47 static
  48 size_t parent(size_t i)
  49 {
  50         return (i - 1) >> 1;
  51 }
  52
  53 static
  54 size_t left(size_t i)
  55 {
  56         return (i << 1) + 1;
  57 }
  58
  59 static
  60 size_t right(size_t i)
  61 {
  62         return (i << 1) + 2;
  63 }
  64
  65 /*
  66  * Copy of heap->ptrs pointer is invalid after heap_grow.
  67  */
  68 static
  69 int heap_grow(struct ptr_heap *heap, size_t new_len)
  70 {
  71         void **new_ptrs;
  72
  73         if (G_LIKELY(heap->alloc_len >= new_len))
  74                 return 0;
  75
  76         heap->alloc_len = bt_max_t(size_t, new_len, heap->alloc_len << 1);
  77         new_ptrs = calloc(heap->alloc_len, sizeof(void *));
  78         if (G_UNLIKELY(!new_ptrs))
  79                 return -ENOMEM;
  80         if (G_LIKELY(heap->ptrs))
  81                 memcpy(new_ptrs, heap->ptrs, heap->len * sizeof(void *));
  82         free(heap->ptrs);
  83         heap->ptrs = new_ptrs;
  84         return 0;
  85 }
  86
  87 static
  88 int heap_set_len(struct ptr_heap *heap, size_t new_len)
  89 {
  90         int ret;
  91
  92         ret = heap_grow(heap, new_len);
  93         if (G_UNLIKELY(ret))
  94                 return ret;
  95         heap->len = new_len;
  96         return 0;
  97 }
  98
  99 int bt_heap_init(struct ptr_heap *heap, size_t alloc_len,
 100               int gt(void *a, void *b))
 101 {
 102         heap->ptrs = NULL;
 103         heap->len = 0;
 104         heap->alloc_len = 0;
 105         heap->gt = gt;
 106         /*
 107          * Minimum size allocated is 1 entry to ensure memory allocation
 108          * never fails within bt_heap_replace_max.
 109          */
 110         return heap_grow(heap, bt_max_t(size_t, 1, alloc_len));
 111 }
 112
 113 void bt_heap_free(struct ptr_heap *heap)
 114 {
 115         free(heap->ptrs);
 116 }
 117
 118 static void heapify(struct ptr_heap *heap, size_t i)
 119 {
 120         void **ptrs = heap->ptrs;
 121         size_t l, r, largest;
 122
 123         for (;;) {
 124                 void *tmp;
 125
 126                 l = left(i);
 127                 r = right(i);
 128                 if (l < heap->len && heap->gt(ptrs[l], ptrs[i]))
 129                         largest = l;
 130                 else
 131                         largest = i;
 132                 if (r < heap->len && heap->gt(ptrs[r], ptrs[largest]))
 133                         largest = r;
 134                 if (G_UNLIKELY(largest == i))
 135                         break;
 136                 tmp = ptrs[i];
 137                 ptrs[i] = ptrs[largest];
 138                 ptrs[largest] = tmp;
 139                 i = largest;
 140         }
 141         check_heap(heap);
 142 }
 143
 144 void *bt_heap_replace_max(struct ptr_heap *heap, void *p)
 145 {
 146         void *res;
 147
 148         if (G_UNLIKELY(!heap->len)) {
 149                 (void) heap_set_len(heap, 1);
 150                 heap->ptrs[0] = p;
 151                 check_heap(heap);
 152                 return NULL;
 153         }
 154
 155         /* Replace the current max and heapify */
 156         res = heap->ptrs[0];
 157         heap->ptrs[0] = p;
 158         heapify(heap, 0);
 159         return res;
 160 }
 161
 162 int bt_heap_insert(struct ptr_heap *heap, void *p)
 163 {
 164         void **ptrs;
 165         size_t pos;
 166         int ret;
 167
 168         ret = heap_set_len(heap, heap->len + 1);
 169         if (G_UNLIKELY(ret))
 170                 return ret;
 171         ptrs = heap->ptrs;
 172         pos = heap->len - 1;
 173         while (pos > 0 && heap->gt(p, ptrs[parent(pos)])) {
 174                 /* Move parent down until we find the right spot */
 175                 ptrs[pos] = ptrs[parent(pos)];
 176                 pos = parent(pos);
 177         }
 178         ptrs[pos] = p;
 179         check_heap(heap);
 180         return 0;
 181 }
 182
 183 void *bt_heap_remove(struct ptr_heap *heap)
 184 {
 185         switch (heap->len) {
 186         case 0:
 187                 return NULL;
 188         case 1:
 189                 (void) heap_set_len(heap, 0);
 190                 return heap->ptrs[0];
 191         }
 192         /* Shrink, replace the current max by previous last entry and heapify */
 193         heap_set_len(heap, heap->len - 1);
 194         /* len changed. previous last entry is at heap->len */
 195         return bt_heap_replace_max(heap, heap->ptrs[heap->len]);
 196 }
 197
 198 void *bt_heap_cherrypick(struct ptr_heap *heap, void *p)
 199 {
 200         size_t pos, len = heap->len;
 201
 202         for (pos = 0; pos < len; pos++)
 203                 if (G_UNLIKELY(heap->ptrs[pos] == p))
 204                         goto found;
 205         return NULL;
 206 found:
 207         if (G_UNLIKELY(heap->len == 1)) {
 208                 (void) heap_set_len(heap, 0);
 209                 check_heap(heap);
 210                 return heap->ptrs[0];
 211         }
 212         /* Replace p with previous last entry and heapify. */
 213         heap_set_len(heap, heap->len - 1);
 214         /* len changed. previous last entry is at heap->len */
 215         heap->ptrs[pos] = heap->ptrs[heap->len];
 216         heapify(heap, pos);
 217         return p;
 218 }
 219
 220 int bt_heap_copy(struct ptr_heap *dst, struct ptr_heap *src)
 221 {
 222         int ret;
 223
 224         ret = bt_heap_init(dst, src->alloc_len, src->gt);
 225         if (ret < 0)
 226                 goto end;
 227
 228         ret = heap_set_len(dst, src->len);
 229         if (ret < 0)
 230                 goto end;
 231
 232         memcpy(dst->ptrs, src->ptrs, src->len * sizeof(void *));
 233
 234 end:
 235         return ret;
 236 }