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; |
39 | void *data = heap->data; |
40 | void *root = data + pos * func->elem_size; |
41 | int i = pos, j; |
42 | |
43 | /* Find the sift-down path all the way to the leaves. */ |
44 | for (;;) { |
45 | if (i * 2 + 2 >= heap->nr) |
46 | break; |
47 | left = data + (i * 2 + 1) * func->elem_size; |
48 | right = data + (i * 2 + 2) * func->elem_size; |
49 | i = func->less(left, right) ? i * 2 + 1 : i * 2 + 2; |
50 | } |
51 | |
52 | /* Special case for the last leaf with no sibling. */ |
53 | if (i * 2 + 2 == heap->nr) |
54 | i = i * 2 + 1; |
55 | |
56 | /* Backtrack to the correct location. */ |
57 | while (i != pos && func->less(root, data + i * func->elem_size)) |
58 | i = (i - 1) / 2; |
59 | |
60 | /* Shift the element into its correct place. */ |
61 | j = i; |
62 | while (i != pos) { |
63 | i = (i - 1) / 2; |
64 | func->swp(data + i * func->elem_size, data + j * func->elem_size); |
65 | } |
66 | } |
67 | |
68 | /* Floyd's approach to heapification that is O(nr). */ |
69 | static __always_inline |
70 | void min_heapify_all(struct min_heap *heap, |
71 | const struct min_heap_callbacks *func) |
72 | { |
73 | int i; |
74 | |
75 | for (i = heap->nr / 2 - 1; i >= 0; i--) |
76 | min_heapify(heap, pos: i, func); |
77 | } |
78 | |
79 | /* Remove minimum element from the heap, O(log2(nr)). */ |
80 | static __always_inline |
81 | void min_heap_pop(struct min_heap *heap, |
82 | const struct min_heap_callbacks *func) |
83 | { |
84 | void *data = heap->data; |
85 | |
86 | if (WARN_ONCE(heap->nr <= 0, "Popping an empty heap" )) |
87 | return; |
88 | |
89 | /* Place last element at the root (position 0) and then sift down. */ |
90 | heap->nr--; |
91 | memcpy(data, data + (heap->nr * func->elem_size), func->elem_size); |
92 | min_heapify(heap, pos: 0, func); |
93 | } |
94 | |
95 | /* |
96 | * Remove the minimum element and then push the given element. The |
97 | * implementation performs 1 sift (O(log2(nr))) and is therefore more |
98 | * efficient than a pop followed by a push that does 2. |
99 | */ |
100 | static __always_inline |
101 | void min_heap_pop_push(struct min_heap *heap, |
102 | const void *element, |
103 | const struct min_heap_callbacks *func) |
104 | { |
105 | memcpy(heap->data, element, func->elem_size); |
106 | min_heapify(heap, pos: 0, func); |
107 | } |
108 | |
109 | /* Push an element on to the heap, O(log2(nr)). */ |
110 | static __always_inline |
111 | void min_heap_push(struct min_heap *heap, const void *element, |
112 | const struct min_heap_callbacks *func) |
113 | { |
114 | void *data = heap->data; |
115 | void *child, *parent; |
116 | int pos; |
117 | |
118 | if (WARN_ONCE(heap->nr >= heap->size, "Pushing on a full heap" )) |
119 | return; |
120 | |
121 | /* Place at the end of data. */ |
122 | pos = heap->nr; |
123 | memcpy(data + (pos * func->elem_size), element, func->elem_size); |
124 | heap->nr++; |
125 | |
126 | /* Sift child at pos up. */ |
127 | for (; pos > 0; pos = (pos - 1) / 2) { |
128 | child = data + (pos * func->elem_size); |
129 | parent = data + ((pos - 1) / 2) * func->elem_size; |
130 | if (func->less(parent, child)) |
131 | break; |
132 | func->swp(parent, child); |
133 | } |
134 | } |
135 | |
136 | #endif /* _LINUX_MIN_HEAP_H */ |
137 | |