lib/prio_heap.c

   1 /*
   2  * Simple insertion-only static-sized priority heap containing
   3  * pointers, based on CLR, chapter 7
   4  */
   5
   6 #include <linux/slab.h>
   7 #include <linux/prio_heap.h>
   8
   9 int heap_init(struct ptr_heap *heap, size_t size, gfp_t gfp_mask,
  10               int (*gt)(void *, void *))
  11 {
  12         heap->ptrs = kmalloc(size, gfp_mask);
  13         if (!heap->ptrs)
  14                 return -ENOMEM;
  15         heap->size = 0;
  16         heap->max = size / sizeof(void *);
  17         heap->gt = gt;
  18         return 0;
  19 }
  20
  21 void heap_free(struct ptr_heap *heap)
  22 {
  23         kfree(heap->ptrs);
  24 }
  25
  26 void *heap_insert(struct ptr_heap *heap, void *p)
  27 {
  28         void *res;
  29         void **ptrs = heap->ptrs;
  30         int pos;
  31
  32         if (heap->size < heap->max) {
  33                 /* Heap insertion */
  34                 int pos = heap->size++;
  35                 while (pos > 0 && heap->gt(p, ptrs[(pos-1)/2])) {
  36                         ptrs[pos] = ptrs[(pos-1)/2];
  37                         pos = (pos-1)/2;
  38                 }
  39                 ptrs[pos] = p;
  40                 return NULL;
  41         }
  42
  43         /* The heap is full, so something will have to be dropped */
  44
  45         /* If the new pointer is greater than the current max, drop it */
  46         if (heap->gt(p, ptrs[0]))
  47                 return p;
  48
  49         /* Replace the current max and heapify */
  50         res = ptrs[0];
  51         ptrs[0] = p;
  52         pos = 0;
  53
  54         while (1) {
  55                 int left = 2 * pos + 1;
  56                 int right = 2 * pos + 2;
  57                 int largest = pos;
  58                 if (left < heap->size && heap->gt(ptrs[left], p))
  59                         largest = left;
  60                 if (right < heap->size && heap->gt(ptrs[right], ptrs[largest]))
  61                         largest = right;
  62                 if (largest == pos)
  63                         break;
  64                 /* Push p down the heap one level and bump one up */
  65                 ptrs[pos] = ptrs[largest];
  66                 ptrs[largest] = p;
  67                 pos = largest;
  68         }
  69         return res;
  70 }