lib/prio_heap/prio_heap.c

   1 /*
   2  * 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 "prio_heap.h"
  21
  22 static
  23 size_t parent(size_t i)
  24 {
  25         return i >> 1;
  26 }
  27
  28 static
  29 size_t left(size_t i)
  30 {
  31         return i << 1;
  32 }
  33
  34 static
  35 size_t right(size_t i)
  36 {
  37         return (i << 1) + 1;
  38 }
  39
  40 /*
  41  * Copy of heap->ptrs pointer is invalid after heap_grow.
  42  */
  43 static
  44 int heap_grow(struct ptr_heap *heap, size_t new_len)
  45 {
  46         void **new_ptrs;
  47
  48         if (heap->alloc_len >= new_len)
  49                 return 0;
  50
  51         heap->alloc_len = max_t(size_t, new_len, heap->alloc_len << 1);
  52         new_ptrs = kmalloc(heap->alloc_len * sizeof(void *), heap->gfpmask);
  53         if (!new_ptrs)
  54                 return -ENOMEM;
  55         if (heap->ptrs)
  56                 memcpy(new_ptrs, heap->ptrs, heap->len * sizeof(void *));
  57         kfree(heap->ptrs);
  58         heap->ptrs = new_ptrs;
  59         return 0;
  60 }
  61
  62 static
  63 int heap_set_len(struct ptr_heap *heap, size_t new_len)
  64 {
  65         int ret;
  66
  67         ret = heap_grow(heap, new_len);
  68         if (ret)
  69                 return ret;
  70         heap->len = new_len;
  71         return 0;
  72 }
  73
  74 int heap_init(struct ptr_heap *heap, size_t alloc_len,
  75               gfp_t gfpmask, int gt(void *a, void *b))
  76 {
  77         heap->ptrs = NULL;
  78         heap->len = 0;
  79         heap->alloc_len = 0;
  80         heap->gt = gt;
  81         /*
  82          * Minimum size allocated is 1 entry to ensure memory allocation
  83          * never fails within heap_replace_max.
  84          */
  85         return heap_grow(heap, max_t(size_t, 1, alloc_len));
  86 }
  87
  88 void heap_free(struct ptr_heap *heap)
  89 {
  90         kfree(heap->ptrs);
  91 }
  92
  93 static void heapify(struct ptr_heap *heap, size_t i)
  94 {
  95         void **ptrs = heap->ptrs;
  96         size_t l, r, largest;
  97
  98         for (;;) {
  99                 l = left(i);
 100                 r = right(i);
 101                 if (l <= heap->len && ptrs[l] > ptrs[i])
 102                         largest = l;
 103                 else
 104                         largest = i;
 105                 if (r <= heap->len && ptrs[r] > ptrs[largest])
 106                         largest = r;
 107                 if (largest != i) {
 108                         void *tmp;
 109
 110                         tmp = ptrs[i];
 111                         ptrs[i] = ptrs[largest];
 112                         ptrs[largest] = tmp;
 113                         i = largest;
 114                         continue;
 115                 } else {
 116                         break;
 117                 }
 118         }
 119 }
 120
 121 void *heap_replace_max(struct ptr_heap *heap, void *p)
 122 {
 123         void *res;
 124
 125         if (!heap->len) {
 126                 (void) heap_set_len(heap, 1);
 127                 heap->ptrs[0] = p;
 128                 return NULL;
 129         }
 130
 131         /* Replace the current max and heapify */
 132         res = heap->ptrs[0];
 133         heap->ptrs[0] = p;
 134         heapify(heap, 0);
 135         return res;
 136 }
 137
 138 int heap_insert(struct ptr_heap *heap, void *p)
 139 {
 140         void **ptrs;
 141         size_t pos;
 142         int ret;
 143
 144         ret = heap_set_len(heap, heap->len + 1);
 145         if (ret)
 146                 return ret;
 147         ptrs = heap->ptrs;
 148         /* Add the element to the end */
 149         ptrs[heap->len - 1] = p;
 150         pos = heap->len - 1;
 151         /* Bubble it up to the appropriate position. */
 152         for (;;) {
 153                 if (pos > 0 && heap->gt(ptrs[pos], ptrs[parent(pos)])) {
 154                         void *tmp;
 155
 156                         /* Need to exchange */
 157                         tmp = ptrs[pos];
 158                         ptrs[pos] = ptrs[parent(pos)];
 159                         ptrs[parent(pos)] = tmp;
 160                         pos = parent(pos);
 161                         /*
 162                          * No need to rebalance: if we are larger than
 163                          * our parent, we are necessarily larger than
 164                          * its other child.
 165                          */
 166                 } else {
 167                         break;
 168                 }
 169         }
 170         return 0;
 171 }
 172
 173 void *heap_remove(struct ptr_heap *heap)
 174 {
 175         switch (heap->len) {
 176         case 0:
 177                 return NULL;
 178         case 1:
 179                 (void) heap_set_len(heap, 0);
 180                 return heap->ptrs[0];
 181         }
 182         /* Shrink, replace the current max by previous last entry and heapify */
 183         heap_set_len(heap, heap->len - 1);
 184         return heap_replace_max(heap, heap->ptrs[heap->len - 1]);
 185 }
 186
 187 void *heap_cherrypick(struct 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                 return heap->ptrs[0];
 199         }
 200         /* Replace p with previous last entry and heapify. */
 201         heap_set_len(heap, heap->len - 1);
 202         heap->ptrs[pos] = heap->ptrs[heap->len - 1];
 203         heapify(heap, pos);
 204         return p;
 205 }