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