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