Root/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
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

Archive Download this file



interactive