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         heap->gfpmask = gfpmask;
  95         /*
  96          * Minimum size allocated is 1 entry to ensure memory allocation
  97          * never fails within heap_replace_max.
  98          */
  99         return heap_grow(heap, max_t(size_t, 1, alloc_len));
 100 }
 101
 102 void lttng_heap_free(struct lttng_ptr_heap *heap)
 103 {
 104         kfree(heap->ptrs);
 105 }
 106
 107 static void heapify(struct lttng_ptr_heap *heap, size_t i)
 108 {
 109         void **ptrs = heap->ptrs;
 110         size_t l, r, largest;
 111
 112         for (;;) {
 113                 void *tmp;
 114
 115                 l = left(i);
 116                 r = right(i);
 117                 if (l < heap->len && heap->gt(ptrs[l], ptrs[i]))
 118                         largest = l;
 119                 else
 120                         largest = i;
 121                 if (r < heap->len && heap->gt(ptrs[r], ptrs[largest]))
 122                         largest = r;
 123                 if (largest == i)
 124                         break;
 125                 tmp = ptrs[i];
 126                 ptrs[i] = ptrs[largest];
 127                 ptrs[largest] = tmp;
 128                 i = largest;
 129         }
 130         lttng_check_heap(heap);
 131 }
 132
 133 void *lttng_heap_replace_max(struct lttng_ptr_heap *heap, void *p)
 134 {
 135         void *res;
 136
 137         if (!heap->len) {
 138                 (void) heap_set_len(heap, 1);
 139                 heap->ptrs[0] = p;
 140                 lttng_check_heap(heap);
 141                 return NULL;
 142         }
 143
 144         /* Replace the current max and heapify */
 145         res = heap->ptrs[0];
 146         heap->ptrs[0] = p;
 147         heapify(heap, 0);
 148         return res;
 149 }
 150
 151 int lttng_heap_insert(struct lttng_ptr_heap *heap, void *p)
 152 {
 153         void **ptrs;
 154         size_t pos;
 155         int ret;
 156
 157         ret = heap_set_len(heap, heap->len + 1);
 158         if (ret)
 159                 return ret;
 160         ptrs = heap->ptrs;
 161         pos = heap->len - 1;
 162         while (pos > 0 && heap->gt(p, ptrs[parent(pos)])) {
 163                 /* Move parent down until we find the right spot */
 164                 ptrs[pos] = ptrs[parent(pos)];
 165                 pos = parent(pos);
 166         }
 167         ptrs[pos] = p;
 168         lttng_check_heap(heap);
 169         return 0;
 170 }
 171
 172 void *lttng_heap_remove(struct lttng_ptr_heap *heap)
 173 {
 174         switch (heap->len) {
 175         case 0:
 176                 return NULL;
 177         case 1:
 178                 (void) heap_set_len(heap, 0);
 179                 return heap->ptrs[0];
 180         }
 181         /* Shrink, replace the current max by previous last entry and heapify */
 182         heap_set_len(heap, heap->len - 1);
 183         /* len changed. previous last entry is at heap->len */
 184         return lttng_heap_replace_max(heap, heap->ptrs[heap->len]);
 185 }
 186
 187 void *lttng_heap_cherrypick(struct lttng_ptr_heap *heap, void *p)
 188 {
 189         size_t pos, len = heap->len;
 190
 191         for (pos = 0; pos < len; pos++)
 192                 if (heap->ptrs[pos] == p)
 193                         goto found;
 194         return NULL;
 195 found:
 196         if (heap->len == 1) {
 197                 (void) heap_set_len(heap, 0);
 198                 lttng_check_heap(heap);
 199                 return heap->ptrs[0];
 200         }
 201         /* Replace p with previous last entry and heapify. */
 202         heap_set_len(heap, heap->len - 1);
 203         /* len changed. previous last entry is at heap->len */
 204         heap->ptrs[pos] = heap->ptrs[heap->len];
 205         heapify(heap, pos);
 206         return p;
 207 }