lib/prio_heap/prio_heap.c

   1 /*
   2  * LICENSING: this file is copied from the Linux kernel. We should therefore
   3  * assume a GPLv2 license for the code that comes from the Linux mainline.
   4  */
   5
   6 /*
   7  * Static-sized priority heap containing pointers. Based on CLR, chapter 7.
   8  */
   9
  10 #include <linux/slab.h>
  11 #include <linux/prio_heap.h>
  12
  13 int heap_init(struct ptr_heap *heap, size_t size, gfp_t gfp_mask,
  14               int (*gt)(void *, void *))
  15 {
  16         heap->ptrs = kmalloc(size, gfp_mask);
  17         if (!heap->ptrs)
  18                 return -ENOMEM;
  19         heap->size = 0;
  20         heap->max = size / sizeof(void *);
  21         heap->gt = gt;
  22         return 0;
  23 }
  24
  25 void heap_free(struct ptr_heap *heap)
  26 {
  27         kfree(heap->ptrs);
  28 }
  29
  30 static void heapify(struct ptr_heap *heap, int pos)
  31 {
  32         void **ptrs = heap->ptrs;
  33         void *p = ptrs[pos];
  34
  35         while (1) {
  36                 int left = 2 * pos + 1;
  37                 int right = 2 * pos + 2;
  38                 int largest = pos;
  39                 if (left < heap->size && heap->gt(ptrs[left], p))
  40                         largest = left;
  41                 if (right < heap->size && heap->gt(ptrs[right], ptrs[largest]))
  42                         largest = right;
  43                 if (largest == pos)
  44                         break;
  45                 /* Push p down the heap one level and bump one up */
  46                 ptrs[pos] = ptrs[largest];
  47                 ptrs[largest] = p;
  48                 pos = largest;
  49         }
  50 }
  51
  52 void *heap_replace_max(struct ptr_heap *heap, void *p)
  53 {
  54         void *res;
  55         void **ptrs = heap->ptrs;
  56
  57         if (!heap->size) {
  58                 ptrs[0] = p;
  59                 heap->size = 1;
  60                 return NULL;
  61         }
  62
  63         /* Replace the current max and heapify */
  64         res = ptrs[0];
  65         ptrs[0] = p;
  66         heapify(heap, 0);
  67         return res;
  68 }
  69
  70 void *heap_insert(struct ptr_heap *heap, void *p)
  71 {
  72         void **ptrs = heap->ptrs;
  73         int pos;
  74
  75         if (heap->size < heap->max) {
  76                 /* Heap insertion */
  77                 pos = heap->size++;
  78                 while (pos > 0 && heap->gt(p, ptrs[(pos-1)/2])) {
  79                         ptrs[pos] = ptrs[(pos-1)/2];
  80                         pos = (pos-1)/2;
  81                 }
  82                 ptrs[pos] = p;
  83                 return NULL;
  84         }
  85
  86         /* The heap is full, so something will have to be dropped */
  87
  88         /* If the new pointer is greater than the current max, drop it */
  89         if (heap->gt(p, ptrs[0]))
  90                 return p;
  91
  92         /* Replace the current max and heapify */
  93         return heap_replace_max(heap, p);
  94 }
  95
  96 void *heap_remove(struct ptr_heap *heap)
  97 {
  98         void **ptrs = heap->ptrs;
  99
 100         switch (heap->size) {
 101         case 0:
 102                 return NULL;
 103         case 1:
 104                 heap->size = 0;
 105                 return ptrs[0];
 106         }
 107
 108         /* Shrink, replace the current max by previous last entry and heapify */
 109         return heap_replace_max(heap, ptrs[--heap->size]);
 110 }
 111
 112 void *heap_cherrypick(struct ptr_heap *heap, void *p)
 113 {
 114         void **ptrs = heap->ptrs;
 115         size_t pos, size = heap->size;
 116
 117         for (pos = 0; pos < size; pos++)
 118                 if (ptrs[pos] == p)
 119                         goto found;
 120         return NULL;
 121 found:
 122         if (heap->size == 1) {
 123                 heap->size = 0;
 124                 return ptrs[0];
 125         }
 126         /*
 127          * Replace p with previous last entry and heapify.
 128          */
 129         ptrs[pos] = ptrs[--heap->size];
 130         heapify(heap, pos);
 131         return p;
 132 }