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 
heap_init(struct ptr_heap * heap,size_t size,gfp_t gfp_mask,int (* gt)(void *,void *))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 
heap_free(struct ptr_heap * heap)21 void heap_free(struct ptr_heap *heap)
22 {
23 	kfree(heap->ptrs);
24 }
25 
heap_insert(struct ptr_heap * heap,void * p)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 		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