LXR linux/lib/prio

   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
   9int 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
  21void heap_free(struct ptr_heap *heap)
  22{
  23        kfree(heap->ptrs);
  24}
  25
  26void *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                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}
  71