Commit | Line | Data |
---|---|---|
6e24628d IR |
1 | /* SPDX-License-Identifier: GPL-2.0 */ |
2 | #ifndef _LINUX_MIN_HEAP_H | |
3 | #define _LINUX_MIN_HEAP_H | |
4 | ||
5 | #include <linux/bug.h> | |
6 | #include <linux/string.h> | |
7 | #include <linux/types.h> | |
8 | ||
9 | /** | |
10 | * struct min_heap - Data structure to hold a min-heap. | |
11 | * @data: Start of array holding the heap elements. | |
12 | * @nr: Number of elements currently in the heap. | |
13 | * @size: Maximum number of elements that can be held in current storage. | |
14 | */ | |
15 | struct min_heap { | |
16 | void *data; | |
17 | int nr; | |
18 | int size; | |
19 | }; | |
20 | ||
21 | /** | |
22 | * struct min_heap_callbacks - Data/functions to customise the min_heap. | |
23 | * @elem_size: The nr of each element in bytes. | |
24 | * @less: Partial order function for this heap. | |
25 | * @swp: Swap elements function. | |
26 | */ | |
27 | struct min_heap_callbacks { | |
28 | int elem_size; | |
29 | bool (*less)(const void *lhs, const void *rhs); | |
30 | void (*swp)(void *lhs, void *rhs); | |
31 | }; | |
32 | ||
33 | /* Sift the element at pos down the heap. */ | |
34 | static __always_inline | |
35 | void min_heapify(struct min_heap *heap, int pos, | |
36 | const struct min_heap_callbacks *func) | |
37 | { | |
38 | void *left, *right, *parent, *smallest; | |
39 | void *data = heap->data; | |
40 | ||
41 | for (;;) { | |
42 | if (pos * 2 + 1 >= heap->nr) | |
43 | break; | |
44 | ||
45 | left = data + ((pos * 2 + 1) * func->elem_size); | |
46 | parent = data + (pos * func->elem_size); | |
47 | smallest = parent; | |
48 | if (func->less(left, smallest)) | |
49 | smallest = left; | |
50 | ||
51 | if (pos * 2 + 2 < heap->nr) { | |
52 | right = data + ((pos * 2 + 2) * func->elem_size); | |
53 | if (func->less(right, smallest)) | |
54 | smallest = right; | |
55 | } | |
56 | if (smallest == parent) | |
57 | break; | |
58 | func->swp(smallest, parent); | |
59 | if (smallest == left) | |
60 | pos = (pos * 2) + 1; | |
61 | else | |
62 | pos = (pos * 2) + 2; | |
63 | } | |
64 | } | |
65 | ||
66 | /* Floyd's approach to heapification that is O(nr). */ | |
67 | static __always_inline | |
68 | void min_heapify_all(struct min_heap *heap, | |
69 | const struct min_heap_callbacks *func) | |
70 | { | |
71 | int i; | |
72 | ||
73 | for (i = heap->nr / 2; i >= 0; i--) | |
74 | min_heapify(heap, i, func); | |
75 | } | |
76 | ||
77 | /* Remove minimum element from the heap, O(log2(nr)). */ | |
78 | static __always_inline | |
79 | void min_heap_pop(struct min_heap *heap, | |
80 | const struct min_heap_callbacks *func) | |
81 | { | |
82 | void *data = heap->data; | |
83 | ||
84 | if (WARN_ONCE(heap->nr <= 0, "Popping an empty heap")) | |
85 | return; | |
86 | ||
87 | /* Place last element at the root (position 0) and then sift down. */ | |
88 | heap->nr--; | |
89 | memcpy(data, data + (heap->nr * func->elem_size), func->elem_size); | |
90 | min_heapify(heap, 0, func); | |
91 | } | |
92 | ||
93 | /* | |
94 | * Remove the minimum element and then push the given element. The | |
95 | * implementation performs 1 sift (O(log2(nr))) and is therefore more | |
96 | * efficient than a pop followed by a push that does 2. | |
97 | */ | |
98 | static __always_inline | |
99 | void min_heap_pop_push(struct min_heap *heap, | |
100 | const void *element, | |
101 | const struct min_heap_callbacks *func) | |
102 | { | |
103 | memcpy(heap->data, element, func->elem_size); | |
104 | min_heapify(heap, 0, func); | |
105 | } | |
106 | ||
107 | /* Push an element on to the heap, O(log2(nr)). */ | |
108 | static __always_inline | |
109 | void min_heap_push(struct min_heap *heap, const void *element, | |
110 | const struct min_heap_callbacks *func) | |
111 | { | |
112 | void *data = heap->data; | |
113 | void *child, *parent; | |
114 | int pos; | |
115 | ||
116 | if (WARN_ONCE(heap->nr >= heap->size, "Pushing on a full heap")) | |
117 | return; | |
118 | ||
119 | /* Place at the end of data. */ | |
120 | pos = heap->nr; | |
121 | memcpy(data + (pos * func->elem_size), element, func->elem_size); | |
122 | heap->nr++; | |
123 | ||
124 | /* Sift child at pos up. */ | |
125 | for (; pos > 0; pos = (pos - 1) / 2) { | |
126 | child = data + (pos * func->elem_size); | |
127 | parent = data + ((pos - 1) / 2) * func->elem_size; | |
128 | if (func->less(parent, child)) | |
129 | break; | |
130 | func->swp(parent, child); | |
131 | } | |
132 | } | |
133 | ||
134 | #endif /* _LINUX_MIN_HEAP_H */ |