lib/prio_heap/prio_heap.c

   1 /*
   2  * prio_heap.c
   3  *
   4  * Static-sized priority heap containing pointers. Based on CLRS,
   5  * chapter 6.
   6  *
   7  * Copyright 2011 - Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
   8  *
   9  * Permission is hereby granted, free of charge, to any person obtaining a copy
  10  * of this software and associated documentation files (the "Software"), to deal
  11  * in the Software without restriction, including without limitation the rights
  12  * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  13  * copies of the Software, and to permit persons to whom the Software is
  14  * furnished to do so, subject to the following conditions:
  15  *
  16  * The above copyright notice and this permission notice shall be included in
  17  * all copies or substantial portions of the Software.
  18  */
  19
  20 #include <linux/slab.h>
  21 #include <linux/prio_heap.h>
  22
  23 int heap_init(struct ptr_heap *heap, size_t size,
  24                gfp_t gfpmask, int gt(void *a, void *b))
  25 {
  26         WARN_ON_ONCE(size == 0);
  27         heap->ptrs = kmalloc(size * sizeof(void *), gfpmask);
  28         if (!heap->ptrs)
  29                 return -ENOMEM;
  30         heap->size = 0;
  31         heap->max = size;
  32         heap->gt = gt;
  33         return 0;
  34 }
  35
  36 void heap_free(struct ptr_heap *heap)
  37 {
  38         kfree(heap->ptrs);
  39 }
  40
  41 static void heapify(struct ptr_heap *heap, int pos)
  42 {
  43         void **ptrs = heap->ptrs;
  44         void *p = ptrs[pos];
  45
  46         while (1) {
  47                 int left = 2 * pos + 1;
  48                 int right = 2 * pos + 2;
  49                 int largest = pos;
  50                 if (left < heap->size && heap->gt(ptrs[left], p))
  51                         largest = left;
  52                 if (right < heap->size && heap->gt(ptrs[right], ptrs[largest]))
  53                         largest = right;
  54                 if (largest == pos)
  55                         break;
  56                 /* Push p down the heap one level and bump one up */
  57                 ptrs[pos] = ptrs[largest];
  58                 ptrs[largest] = p;
  59                 pos = largest;
  60         }
  61 }
  62
  63 void *heap_replace_max(struct ptr_heap *heap, void *p)
  64 {
  65         void *res;
  66         void **ptrs = heap->ptrs;
  67
  68         if (!heap->size) {
  69                 ptrs[0] = p;
  70                 heap->size = 1;
  71                 return NULL;
  72         }
  73
  74         /* Replace the current max and heapify */
  75         res = ptrs[0];
  76         ptrs[0] = p;
  77         heapify(heap, 0);
  78         return res;
  79 }
  80
  81 void *heap_insert(struct ptr_heap *heap, void *p)
  82 {
  83         void **ptrs = heap->ptrs;
  84         void *tmp = NULL;
  85
  86         if (heap->size < heap->max) {
  87                 /* Add the element to the end */
  88                 heap->ptrs[heap->size++] = p;
  89                 /* rebalance */
  90                 heapify(heap, 0);
  91                 return NULL;
  92         }
  93
  94         /*
  95          * Full. We need to replace the largest (if we are
  96          * smaller or equal to this element).
  97          */
  98         if (heap->gt(ptrs[0], p)) {
  99                 tmp = ptrs[0];
 100                 ptrs[0] = p;
 101                 /* rebalance */
 102                 heapify(heap, 0);
 103         } else {
 104                 tmp = p;
 105         }
 106         return tmp;
 107 }
 108
 109 void *heap_remove(struct ptr_heap *heap)
 110 {
 111         void **ptrs = heap->ptrs;
 112
 113         switch (heap->size) {
 114         case 0:
 115                 return NULL;
 116         case 1:
 117                 heap->size = 0;
 118                 return ptrs[0];
 119         }
 120
 121         /* Shrink, replace the current max by previous last entry and heapify */
 122         return heap_replace_max(heap, ptrs[--heap->size]);
 123 }
 124
 125 void *heap_cherrypick(struct ptr_heap *heap, void *p)
 126 {
 127         void **ptrs = heap->ptrs;
 128         size_t pos, size = heap->size;
 129
 130         for (pos = 0; pos < size; pos++)
 131                 if (ptrs[pos] == p)
 132                         goto found;
 133         return NULL;
 134 found:
 135         if (heap->size == 1) {
 136                 heap->size = 0;
 137                 return ptrs[0];
 138         }
 139         /*
 140          * Replace p with previous last entry and heapify.
 141          */
 142         ptrs[pos] = ptrs[--heap->size];
 143         heapify(heap, pos);
 144         return p;
 145 }