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