lib/prio_heap/prio_heap.c

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