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

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