Commit | Line | Data |
---|---|---|
f3bc08c5 MD |
1 | /* |
2 | * LICENSING: this file is copied from the Linux kernel. We should therefore | |
3 | * assume a GPLv2 license for the code that comes from the Linux mainline. | |
4 | */ | |
5 | ||
6 | /* | |
7 | * Static-sized priority heap containing pointers. Based on CLR, chapter 7. | |
8 | */ | |
9 | ||
10 | #include <linux/slab.h> | |
11 | #include <linux/prio_heap.h> | |
12 | ||
13 | int heap_init(struct ptr_heap *heap, size_t size, gfp_t gfp_mask, | |
14 | int (*gt)(void *, void *)) | |
15 | { | |
16 | heap->ptrs = kmalloc(size, gfp_mask); | |
17 | if (!heap->ptrs) | |
18 | return -ENOMEM; | |
19 | heap->size = 0; | |
20 | heap->max = size / sizeof(void *); | |
21 | heap->gt = gt; | |
22 | return 0; | |
23 | } | |
24 | ||
25 | void heap_free(struct ptr_heap *heap) | |
26 | { | |
27 | kfree(heap->ptrs); | |
28 | } | |
29 | ||
30 | static void heapify(struct ptr_heap *heap, int pos) | |
31 | { | |
32 | void **ptrs = heap->ptrs; | |
33 | void *p = ptrs[pos]; | |
34 | ||
35 | while (1) { | |
36 | int left = 2 * pos + 1; | |
37 | int right = 2 * pos + 2; | |
38 | int largest = pos; | |
39 | if (left < heap->size && heap->gt(ptrs[left], p)) | |
40 | largest = left; | |
41 | if (right < heap->size && heap->gt(ptrs[right], ptrs[largest])) | |
42 | largest = right; | |
43 | if (largest == pos) | |
44 | break; | |
45 | /* Push p down the heap one level and bump one up */ | |
46 | ptrs[pos] = ptrs[largest]; | |
47 | ptrs[largest] = p; | |
48 | pos = largest; | |
49 | } | |
50 | } | |
51 | ||
52 | void *heap_replace_max(struct ptr_heap *heap, void *p) | |
53 | { | |
54 | void *res; | |
55 | void **ptrs = heap->ptrs; | |
56 | ||
57 | if (!heap->size) { | |
58 | ptrs[0] = p; | |
59 | heap->size = 1; | |
60 | return NULL; | |
61 | } | |
62 | ||
63 | /* Replace the current max and heapify */ | |
64 | res = ptrs[0]; | |
65 | ptrs[0] = p; | |
66 | heapify(heap, 0); | |
67 | return res; | |
68 | } | |
69 | ||
70 | void *heap_insert(struct ptr_heap *heap, void *p) | |
71 | { | |
72 | void **ptrs = heap->ptrs; | |
73 | int pos; | |
74 | ||
75 | if (heap->size < heap->max) { | |
76 | /* Heap insertion */ | |
77 | pos = heap->size++; | |
78 | while (pos > 0 && heap->gt(p, ptrs[(pos-1)/2])) { | |
79 | ptrs[pos] = ptrs[(pos-1)/2]; | |
80 | pos = (pos-1)/2; | |
81 | } | |
82 | ptrs[pos] = p; | |
83 | return NULL; | |
84 | } | |
85 | ||
86 | /* The heap is full, so something will have to be dropped */ | |
87 | ||
88 | /* If the new pointer is greater than the current max, drop it */ | |
89 | if (heap->gt(p, ptrs[0])) | |
90 | return p; | |
91 | ||
92 | /* Replace the current max and heapify */ | |
93 | return heap_replace_max(heap, p); | |
94 | } | |
95 | ||
96 | void *heap_remove(struct ptr_heap *heap) | |
97 | { | |
98 | void **ptrs = heap->ptrs; | |
99 | ||
100 | switch (heap->size) { | |
101 | case 0: | |
102 | return NULL; | |
103 | case 1: | |
104 | heap->size = 0; | |
105 | return ptrs[0]; | |
106 | } | |
107 | ||
108 | /* Shrink, replace the current max by previous last entry and heapify */ | |
109 | return heap_replace_max(heap, ptrs[--heap->size]); | |
110 | } | |
111 | ||
112 | void *heap_cherrypick(struct ptr_heap *heap, void *p) | |
113 | { | |
114 | void **ptrs = heap->ptrs; | |
115 | size_t pos, size = heap->size; | |
116 | ||
117 | for (pos = 0; pos < size; pos++) | |
118 | if (ptrs[pos] == p) | |
119 | goto found; | |
120 | return NULL; | |
121 | found: | |
122 | if (heap->size == 1) { | |
123 | heap->size = 0; | |
124 | return ptrs[0]; | |
125 | } | |
126 | /* | |
127 | * Replace p with previous last entry and heapify. | |
128 | */ | |
129 | ptrs[pos] = ptrs[--heap->size]; | |
130 | heapify(heap, pos); | |
131 | return p; | |
132 | } |