lib/prio_heap/lttng_prio_heap.c

   1 /*
   2  * lttng_prio_heap.c
   3  *
   4  * Priority heap containing pointers. Based on CLRS, chapter 6.
   5  *
   6  * Copyright 2011 - Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
   7  *
   8  * Permission is hereby granted, free of charge, to any person obtaining a copy
   9  * of this software and associated documentation files (the "Software"), to deal
  10  * in the Software without restriction, including without limitation the rights
  11  * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  12  * copies of the Software, and to permit persons to whom the Software is
  13  * furnished to do so, subject to the following conditions:
  14  *
  15  * The above copyright notice and this permission notice shall be included in
  16  * all copies or substantial portions of the Software.
  17  */
  18
  19 #include <linux/slab.h>
  20 #include "lttng_prio_heap.h"
  21
  22 #ifdef DEBUG_HEAP
  23 void lttng_check_heap(const struct lttng_ptr_heap *heap)
  24 {
  25         size_t i;
  26
  27         if (!heap->len)
  28                 return;
  29
  30         for (i = 1; i < heap->len; i++)
  31                 WARN_ON_ONCE(!heap->gt(heap->ptrs[i], heap->ptrs[0]));
  32 }
  33 #endif
  34
  35 static
  36 size_t parent(size_t i)
  37 {
  38         return (i -1) >> 1;
  39 }
  40
  41 static
  42 size_t left(size_t i)
  43 {
  44         return (i << 1) + 1;
  45 }
  46
  47 static
  48 size_t right(size_t i)
  49 {
  50         return (i << 1) + 2;
  51 }
  52
  53 /*
  54  * Copy of heap->ptrs pointer is invalid after heap_grow.
  55  */
  56 static
  57 int heap_grow(struct lttng_ptr_heap *heap, size_t new_len)
  58 {
  59         void **new_ptrs;
  60
  61         if (heap->alloc_len >= new_len)
  62                 return 0;
  63
  64         heap->alloc_len = max_t(size_t, new_len, heap->alloc_len << 1);
  65         new_ptrs = kmalloc(heap->alloc_len * sizeof(void *), heap->gfpmask);
  66         if (!new_ptrs)
  67                 return -ENOMEM;
  68         if (heap->ptrs)
  69                 memcpy(new_ptrs, heap->ptrs, heap->len * sizeof(void *));
  70         kfree(heap->ptrs);
  71         heap->ptrs = new_ptrs;
  72         return 0;
  73 }
  74
  75 static
  76 int heap_set_len(struct lttng_ptr_heap *heap, size_t new_len)
  77 {
  78         int ret;
  79
  80         ret = heap_grow(heap, new_len);
  81         if (ret)
  82                 return ret;
  83         heap->len = new_len;
  84         return 0;
  85 }
  86
  87 int lttng_heap_init(struct lttng_ptr_heap *heap, size_t alloc_len,
  88               gfp_t gfpmask, int gt(void *a, void *b))
  89 {
  90         heap->ptrs = NULL;
  91         heap->len = 0;
  92         heap->alloc_len = 0;
  93         heap->gt = gt;
  94         /*
  95          * Minimum size allocated is 1 entry to ensure memory allocation
  96          * never fails within heap_replace_max.
  97          */
  98         return heap_grow(heap, max_t(size_t, 1, alloc_len));
  99 }
 100
 101 void lttng_heap_free(struct lttng_ptr_heap *heap)
 102 {
 103         kfree(heap->ptrs);
 104 }
 105
 106 static void heapify(struct lttng_ptr_heap *heap, size_t i)
 107 {
 108         void **ptrs = heap->ptrs;
 109         size_t l, r, largest;
 110
 111         for (;;) {
 112                 void *tmp;
 113
 114                 l = left(i);
 115                 r = right(i);
 116                 if (l < heap->len && heap->gt(ptrs[l], ptrs[i]))
 117                         largest = l;
 118                 else
 119                         largest = i;
 120                 if (r < heap->len && heap->gt(ptrs[r], ptrs[largest]))
 121                         largest = r;
 122                 if (largest == i)
 123                         break;
 124                 tmp = ptrs[i];
 125                 ptrs[i] = ptrs[largest];
 126                 ptrs[largest] = tmp;
 127                 i = largest;
 128         }
 129         lttng_check_heap(heap);
 130 }
 131
 132 void *lttng_heap_replace_max(struct lttng_ptr_heap *heap, void *p)
 133 {
 134         void *res;
 135
 136         if (!heap->len) {
 137                 (void) heap_set_len(heap, 1);
 138                 heap->ptrs[0] = p;
 139                 lttng_check_heap(heap);
 140                 return NULL;
 141         }
 142
 143         /* Replace the current max and heapify */
 144         res = heap->ptrs[0];
 145         heap->ptrs[0] = p;
 146         heapify(heap, 0);
 147         return res;
 148 }
 149
 150 int lttng_heap_insert(struct lttng_ptr_heap *heap, void *p)
 151 {
 152         void **ptrs;
 153         size_t pos;
 154         int ret;
 155
 156         ret = heap_set_len(heap, heap->len + 1);
 157         if (ret)
 158                 return ret;
 159         ptrs = heap->ptrs;
 160         pos = heap->len - 1;
 161         while (pos > 0 && heap->gt(p, ptrs[parent(pos)])) {
 162                 /* Move parent down until we find the right spot */
 163                 ptrs[pos] = ptrs[parent(pos)];
 164                 pos = parent(pos);
 165         }
 166         ptrs[pos] = p;
 167         lttng_check_heap(heap);
 168         return 0;
 169 }
 170
 171 void *lttng_heap_remove(struct lttng_ptr_heap *heap)
 172 {
 173         switch (heap->len) {
 174         case 0:
 175                 return NULL;
 176         case 1:
 177                 (void) heap_set_len(heap, 0);
 178                 return heap->ptrs[0];
 179         }
 180         /* Shrink, replace the current max by previous last entry and heapify */
 181         heap_set_len(heap, heap->len - 1);
 182         /* len changed. previous last entry is at heap->len */
 183         return lttng_heap_replace_max(heap, heap->ptrs[heap->len]);
 184 }
 185
 186 void *lttng_heap_cherrypick(struct lttng_ptr_heap *heap, void *p)
 187 {
 188         size_t pos, len = heap->len;
 189
 190         for (pos = 0; pos < len; pos++)
 191                 if (heap->ptrs[pos] == p)
 192                         goto found;
 193         return NULL;
 194 found:
 195         if (heap->len == 1) {
 196                 (void) heap_set_len(heap, 0);
 197                 lttng_check_heap(heap);
 198                 return heap->ptrs[0];
 199         }
 200         /* Replace p with previous last entry and heapify. */
 201         heap_set_len(heap, heap->len - 1);
 202         /* len changed. previous last entry is at heap->len */
 203         heap->ptrs[pos] = heap->ptrs[heap->len];
 204         heapify(heap, pos);
 205         return p;
 206 }