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 | { | |
c641722e | 38 | void *left, *right; |
6e24628d | 39 | void *data = heap->data; |
c641722e KWC |
40 | void *root = data + pos * func->elem_size; |
41 | int i = pos, j; | |
6e24628d | 42 | |
c641722e | 43 | /* Find the sift-down path all the way to the leaves. */ |
6e24628d | 44 | for (;;) { |
c641722e | 45 | if (i * 2 + 2 >= heap->nr) |
6e24628d | 46 | break; |
c641722e KWC |
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 | } | |
6e24628d | 51 | |
c641722e KWC |
52 | /* Special case for the last leaf with no sibling. */ |
53 | if (i * 2 + 2 == heap->nr) | |
54 | i = i * 2 + 1; | |
6e24628d | 55 | |
c641722e KWC |
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); | |
6e24628d IR |
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 | ||
c499c717 | 75 | for (i = heap->nr / 2 - 1; i >= 0; i--) |
6e24628d IR |
76 | min_heapify(heap, 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, 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, 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 */ |