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