lib/prio_heap/lttng_prio_heap.c

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