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 <babeltrace/prio-heap-internal.h>
  27 #include <babeltrace/babeltrace-internal.h>
  28 #include <babeltrace/assert-internal.h>
  29 #include <errno.h>
  30 #include <stdlib.h>
  31 #include <string.h>
  32
  33 #ifdef DEBUG_HEAP
  34 void check_heap(const struct ptr_heap *heap)
  35 {
  36         size_t i;
  37
  38         if (!heap->len)
  39                 return;
  40
  41         for (i = 1; i < heap->len; i++)
  42                 BT_ASSERT(!heap->gt(heap->ptrs[i], heap->ptrs[0]));
  43 }
  44 #endif
  45
  46 static
  47 size_t parent(size_t i)
  48 {
  49         return (i - 1) >> 1;
  50 }
  51
  52 static
  53 size_t left(size_t i)
  54 {
  55         return (i << 1) + 1;
  56 }
  57
  58 static
  59 size_t right(size_t i)
  60 {
  61         return (i << 1) + 2;
  62 }
  63
  64 /*
  65  * Copy of heap->ptrs pointer is invalid after heap_grow.
  66  */
  67 static
  68 int heap_grow(struct ptr_heap *heap, size_t new_len)
  69 {
  70         void **new_ptrs;
  71
  72         if (likely(heap->alloc_len >= new_len))
  73                 return 0;
  74
  75         heap->alloc_len = max_t(size_t, new_len, heap->alloc_len << 1);
  76         new_ptrs = calloc(heap->alloc_len, sizeof(void *));
  77         if (unlikely(!new_ptrs))
  78                 return -ENOMEM;
  79         if (likely(heap->ptrs))
  80                 memcpy(new_ptrs, heap->ptrs, heap->len * sizeof(void *));
  81         free(heap->ptrs);
  82         heap->ptrs = new_ptrs;
  83         return 0;
  84 }
  85
  86 static
  87 int heap_set_len(struct ptr_heap *heap, size_t new_len)
  88 {
  89         int ret;
  90
  91         ret = heap_grow(heap, new_len);
  92         if (unlikely(ret))
  93                 return ret;
  94         heap->len = new_len;
  95         return 0;
  96 }
  97
  98 int bt_heap_init(struct ptr_heap *heap, size_t alloc_len,
  99               int gt(void *a, void *b))
 100 {
 101         heap->ptrs = NULL;
 102         heap->len = 0;
 103         heap->alloc_len = 0;
 104         heap->gt = gt;
 105         /*
 106          * Minimum size allocated is 1 entry to ensure memory allocation
 107          * never fails within bt_heap_replace_max.
 108          */
 109         return heap_grow(heap, max_t(size_t, 1, alloc_len));
 110 }
 111
 112 void bt_heap_free(struct ptr_heap *heap)
 113 {
 114         free(heap->ptrs);
 115 }
 116
 117 static void heapify(struct ptr_heap *heap, size_t i)
 118 {
 119         void **ptrs = heap->ptrs;
 120         size_t l, r, largest;
 121
 122         for (;;) {
 123                 void *tmp;
 124
 125                 l = left(i);
 126                 r = right(i);
 127                 if (l < heap->len && heap->gt(ptrs[l], ptrs[i]))
 128                         largest = l;
 129                 else
 130                         largest = i;
 131                 if (r < heap->len && heap->gt(ptrs[r], ptrs[largest]))
 132                         largest = r;
 133                 if (unlikely(largest == i))
 134                         break;
 135                 tmp = ptrs[i];
 136                 ptrs[i] = ptrs[largest];
 137                 ptrs[largest] = tmp;
 138                 i = largest;
 139         }
 140         check_heap(heap);
 141 }
 142
 143 void *bt_heap_replace_max(struct ptr_heap *heap, void *p)
 144 {
 145         void *res;
 146
 147         if (unlikely(!heap->len)) {
 148                 (void) heap_set_len(heap, 1);
 149                 heap->ptrs[0] = p;
 150                 check_heap(heap);
 151                 return NULL;
 152         }
 153
 154         /* Replace the current max and heapify */
 155         res = heap->ptrs[0];
 156         heap->ptrs[0] = p;
 157         heapify(heap, 0);
 158         return res;
 159 }
 160
 161 int bt_heap_insert(struct ptr_heap *heap, void *p)
 162 {
 163         void **ptrs;
 164         size_t pos;
 165         int ret;
 166
 167         ret = heap_set_len(heap, heap->len + 1);
 168         if (unlikely(ret))
 169                 return ret;
 170         ptrs = heap->ptrs;
 171         pos = heap->len - 1;
 172         while (pos > 0 && heap->gt(p, ptrs[parent(pos)])) {
 173                 /* Move parent down until we find the right spot */
 174                 ptrs[pos] = ptrs[parent(pos)];
 175                 pos = parent(pos);
 176         }
 177         ptrs[pos] = p;
 178         check_heap(heap);
 179         return 0;
 180 }
 181
 182 void *bt_heap_remove(struct ptr_heap *heap)
 183 {
 184         switch (heap->len) {
 185         case 0:
 186                 return NULL;
 187         case 1:
 188                 (void) heap_set_len(heap, 0);
 189                 return heap->ptrs[0];
 190         }
 191         /* Shrink, replace the current max by previous last entry and heapify */
 192         heap_set_len(heap, heap->len - 1);
 193         /* len changed. previous last entry is at heap->len */
 194         return bt_heap_replace_max(heap, heap->ptrs[heap->len]);
 195 }
 196
 197 void *bt_heap_cherrypick(struct ptr_heap *heap, void *p)
 198 {
 199         size_t pos, len = heap->len;
 200
 201         for (pos = 0; pos < len; pos++)
 202                 if (unlikely(heap->ptrs[pos] == p))
 203                         goto found;
 204         return NULL;
 205 found:
 206         if (unlikely(heap->len == 1)) {
 207                 (void) heap_set_len(heap, 0);
 208                 check_heap(heap);
 209                 return heap->ptrs[0];
 210         }
 211         /* Replace p with previous last entry and heapify. */
 212         heap_set_len(heap, heap->len - 1);
 213         /* len changed. previous last entry is at heap->len */
 214         heap->ptrs[pos] = heap->ptrs[heap->len];
 215         heapify(heap, pos);
 216         return p;
 217 }
 218
 219 int bt_heap_copy(struct ptr_heap *dst, struct ptr_heap *src)
 220 {
 221         int ret;
 222
 223         ret = bt_heap_init(dst, src->alloc_len, src->gt);
 224         if (ret < 0)
 225                 goto end;
 226
 227         ret = heap_set_len(dst, src->len);
 228         if (ret < 0)
 229                 goto end;
 230
 231         memcpy(dst->ptrs, src->ptrs, src->len * sizeof(void *));
 232
 233 end:
 234         return ret;
 235 }