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