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 * The Min Heap API provides utilities for managing min-heaps, a binary tree
11 * structure where each node's value is less than or equal to its children's
12 * values, ensuring the smallest element is at the root.
13 *
14 * Users should avoid directly calling functions prefixed with __min_heap_*().
15 * Instead, use the provided macro wrappers.
16 *
17 * For further details and examples, refer to Documentation/core-api/min_heap.rst.
18 */
19
20 /**
21 * Data structure to hold a min-heap.
22 * @nr: Number of elements currently in the heap.
23 * @size: Maximum number of elements that can be held in current storage.
24 * @data: Pointer to the start of array holding the heap elements.
25 * @preallocated: Start of the static preallocated array holding the heap elements.
26 */
27 #define MIN_HEAP_PREALLOCATED(_type, _name, _nr) \
28 struct _name { \
29 size_t nr; \
30 size_t size; \
31 _type *data; \
32 _type preallocated[_nr]; \
33 }
34
35 #define DEFINE_MIN_HEAP(_type, _name) MIN_HEAP_PREALLOCATED(_type, _name, 0)
36
37 typedef DEFINE_MIN_HEAP(char, min_heap_char) min_heap_char;
38
39 #define __minheap_cast(_heap) (typeof((_heap)->data[0]) *)
40 #define __minheap_obj_size(_heap) sizeof((_heap)->data[0])
41
42 /**
43 * struct min_heap_callbacks - Data/functions to customise the min_heap.
44 * @less: Partial order function for this heap.
45 * @swp: Swap elements function.
46 */
47 struct min_heap_callbacks {
48 bool (*less)(const void *lhs, const void *rhs, void *args);
49 void (*swp)(void *lhs, void *rhs, void *args);
50 };
51
52 /**
53 * is_aligned - is this pointer & size okay for word-wide copying?
54 * @base: pointer to data
55 * @size: size of each element
56 * @align: required alignment (typically 4 or 8)
57 *
58 * Returns true if elements can be copied using word loads and stores.
59 * The size must be a multiple of the alignment, and the base address must
60 * be if we do not have CONFIG_HAVE_EFFICIENT_UNALIGNED_ACCESS.
61 *
62 * For some reason, gcc doesn't know to optimize "if (a & mask || b & mask)"
63 * to "if ((a | b) & mask)", so we do that by hand.
64 */
65 __attribute_const__ __always_inline
is_aligned(const void * base,size_t size,unsigned char align)66 static bool is_aligned(const void *base, size_t size, unsigned char align)
67 {
68 unsigned char lsbits = (unsigned char)size;
69
70 (void)base;
71 #ifndef CONFIG_HAVE_EFFICIENT_UNALIGNED_ACCESS
72 lsbits |= (unsigned char)(uintptr_t)base;
73 #endif
74 return (lsbits & (align - 1)) == 0;
75 }
76
77 /**
78 * swap_words_32 - swap two elements in 32-bit chunks
79 * @a: pointer to the first element to swap
80 * @b: pointer to the second element to swap
81 * @n: element size (must be a multiple of 4)
82 *
83 * Exchange the two objects in memory. This exploits base+index addressing,
84 * which basically all CPUs have, to minimize loop overhead computations.
85 *
86 * For some reason, on x86 gcc 7.3.0 adds a redundant test of n at the
87 * bottom of the loop, even though the zero flag is still valid from the
88 * subtract (since the intervening mov instructions don't alter the flags).
89 * Gcc 8.1.0 doesn't have that problem.
90 */
91 static __always_inline
swap_words_32(void * a,void * b,size_t n)92 void swap_words_32(void *a, void *b, size_t n)
93 {
94 do {
95 u32 t = *(u32 *)(a + (n -= 4));
96 *(u32 *)(a + n) = *(u32 *)(b + n);
97 *(u32 *)(b + n) = t;
98 } while (n);
99 }
100
101 /**
102 * swap_words_64 - swap two elements in 64-bit chunks
103 * @a: pointer to the first element to swap
104 * @b: pointer to the second element to swap
105 * @n: element size (must be a multiple of 8)
106 *
107 * Exchange the two objects in memory. This exploits base+index
108 * addressing, which basically all CPUs have, to minimize loop overhead
109 * computations.
110 *
111 * We'd like to use 64-bit loads if possible. If they're not, emulating
112 * one requires base+index+4 addressing which x86 has but most other
113 * processors do not. If CONFIG_64BIT, we definitely have 64-bit loads,
114 * but it's possible to have 64-bit loads without 64-bit pointers (e.g.
115 * x32 ABI). Are there any cases the kernel needs to worry about?
116 */
117 static __always_inline
swap_words_64(void * a,void * b,size_t n)118 void swap_words_64(void *a, void *b, size_t n)
119 {
120 do {
121 #ifdef CONFIG_64BIT
122 u64 t = *(u64 *)(a + (n -= 8));
123 *(u64 *)(a + n) = *(u64 *)(b + n);
124 *(u64 *)(b + n) = t;
125 #else
126 /* Use two 32-bit transfers to avoid base+index+4 addressing */
127 u32 t = *(u32 *)(a + (n -= 4));
128 *(u32 *)(a + n) = *(u32 *)(b + n);
129 *(u32 *)(b + n) = t;
130
131 t = *(u32 *)(a + (n -= 4));
132 *(u32 *)(a + n) = *(u32 *)(b + n);
133 *(u32 *)(b + n) = t;
134 #endif
135 } while (n);
136 }
137
138 /**
139 * swap_bytes - swap two elements a byte at a time
140 * @a: pointer to the first element to swap
141 * @b: pointer to the second element to swap
142 * @n: element size
143 *
144 * This is the fallback if alignment doesn't allow using larger chunks.
145 */
146 static __always_inline
swap_bytes(void * a,void * b,size_t n)147 void swap_bytes(void *a, void *b, size_t n)
148 {
149 do {
150 char t = ((char *)a)[--n];
151 ((char *)a)[n] = ((char *)b)[n];
152 ((char *)b)[n] = t;
153 } while (n);
154 }
155
156 /*
157 * The values are arbitrary as long as they can't be confused with
158 * a pointer, but small integers make for the smallest compare
159 * instructions.
160 */
161 #define SWAP_WORDS_64 ((void (*)(void *, void *, void *))0)
162 #define SWAP_WORDS_32 ((void (*)(void *, void *, void *))1)
163 #define SWAP_BYTES ((void (*)(void *, void *, void *))2)
164
165 /*
166 * Selects the appropriate swap function based on the element size.
167 */
168 static __always_inline
select_swap_func(const void * base,size_t size)169 void *select_swap_func(const void *base, size_t size)
170 {
171 if (is_aligned(base, size, 8))
172 return SWAP_WORDS_64;
173 else if (is_aligned(base, size, 4))
174 return SWAP_WORDS_32;
175 else
176 return SWAP_BYTES;
177 }
178
179 static __always_inline
do_swap(void * a,void * b,size_t size,void (* swap_func)(void * lhs,void * rhs,void * args),void * priv)180 void do_swap(void *a, void *b, size_t size, void (*swap_func)(void *lhs, void *rhs, void *args),
181 void *priv)
182 {
183 if (swap_func == SWAP_WORDS_64)
184 swap_words_64(a, b, size);
185 else if (swap_func == SWAP_WORDS_32)
186 swap_words_32(a, b, size);
187 else if (swap_func == SWAP_BYTES)
188 swap_bytes(a, b, size);
189 else
190 swap_func(a, b, priv);
191 }
192
193 /**
194 * parent - given the offset of the child, find the offset of the parent.
195 * @i: the offset of the heap element whose parent is sought. Non-zero.
196 * @lsbit: a precomputed 1-bit mask, equal to "size & -size"
197 * @size: size of each element
198 *
199 * In terms of array indexes, the parent of element j = @i/@size is simply
200 * (j-1)/2. But when working in byte offsets, we can't use implicit
201 * truncation of integer divides.
202 *
203 * Fortunately, we only need one bit of the quotient, not the full divide.
204 * @size has a least significant bit. That bit will be clear if @i is
205 * an even multiple of @size, and set if it's an odd multiple.
206 *
207 * Logically, we're doing "if (i & lsbit) i -= size;", but since the
208 * branch is unpredictable, it's done with a bit of clever branch-free
209 * code instead.
210 */
211 __attribute_const__ __always_inline
parent(size_t i,unsigned int lsbit,size_t size)212 static size_t parent(size_t i, unsigned int lsbit, size_t size)
213 {
214 i -= size;
215 i -= size & -(i & lsbit);
216 return i / 2;
217 }
218
219 /* Initialize a min-heap. */
220 static __always_inline
__min_heap_init_inline(min_heap_char * heap,void * data,int size)221 void __min_heap_init_inline(min_heap_char *heap, void *data, int size)
222 {
223 heap->nr = 0;
224 heap->size = size;
225 if (data)
226 heap->data = data;
227 else
228 heap->data = heap->preallocated;
229 }
230
231 #define min_heap_init_inline(_heap, _data, _size) \
232 __min_heap_init_inline(container_of(&(_heap)->nr, min_heap_char, nr), _data, _size)
233
234 /* Get the minimum element from the heap. */
235 static __always_inline
__min_heap_peek_inline(struct min_heap_char * heap)236 void *__min_heap_peek_inline(struct min_heap_char *heap)
237 {
238 return heap->nr ? heap->data : NULL;
239 }
240
241 #define min_heap_peek_inline(_heap) \
242 (__minheap_cast(_heap) \
243 __min_heap_peek_inline(container_of(&(_heap)->nr, min_heap_char, nr)))
244
245 /* Check if the heap is full. */
246 static __always_inline
__min_heap_full_inline(min_heap_char * heap)247 bool __min_heap_full_inline(min_heap_char *heap)
248 {
249 return heap->nr == heap->size;
250 }
251
252 #define min_heap_full_inline(_heap) \
253 __min_heap_full_inline(container_of(&(_heap)->nr, min_heap_char, nr))
254
255 /* Sift the element at pos down the heap. */
256 static __always_inline
__min_heap_sift_down_inline(min_heap_char * heap,int pos,size_t elem_size,const struct min_heap_callbacks * func,void * args)257 void __min_heap_sift_down_inline(min_heap_char *heap, int pos, size_t elem_size,
258 const struct min_heap_callbacks *func, void *args)
259 {
260 const unsigned long lsbit = elem_size & -elem_size;
261 void *data = heap->data;
262 void (*swp)(void *lhs, void *rhs, void *args) = func->swp;
263 /* pre-scale counters for performance */
264 size_t a = pos * elem_size;
265 size_t b, c, d;
266 size_t n = heap->nr * elem_size;
267
268 if (!swp)
269 swp = select_swap_func(data, elem_size);
270
271 /* Find the sift-down path all the way to the leaves. */
272 for (b = a; c = 2 * b + elem_size, (d = c + elem_size) < n;)
273 b = func->less(data + c, data + d, args) ? c : d;
274
275 /* Special case for the last leaf with no sibling. */
276 if (d == n)
277 b = c;
278
279 /* Backtrack to the correct location. */
280 while (b != a && func->less(data + a, data + b, args))
281 b = parent(b, lsbit, elem_size);
282
283 /* Shift the element into its correct place. */
284 c = b;
285 while (b != a) {
286 b = parent(b, lsbit, elem_size);
287 do_swap(data + b, data + c, elem_size, swp, args);
288 }
289 }
290
291 #define min_heap_sift_down_inline(_heap, _pos, _func, _args) \
292 __min_heap_sift_down_inline(container_of(&(_heap)->nr, min_heap_char, nr), _pos, \
293 __minheap_obj_size(_heap), _func, _args)
294
295 /* Sift up ith element from the heap, O(log2(nr)). */
296 static __always_inline
__min_heap_sift_up_inline(min_heap_char * heap,size_t elem_size,size_t idx,const struct min_heap_callbacks * func,void * args)297 void __min_heap_sift_up_inline(min_heap_char *heap, size_t elem_size, size_t idx,
298 const struct min_heap_callbacks *func, void *args)
299 {
300 const unsigned long lsbit = elem_size & -elem_size;
301 void *data = heap->data;
302 void (*swp)(void *lhs, void *rhs, void *args) = func->swp;
303 /* pre-scale counters for performance */
304 size_t a = idx * elem_size, b;
305
306 if (!swp)
307 swp = select_swap_func(data, elem_size);
308
309 while (a) {
310 b = parent(a, lsbit, elem_size);
311 if (func->less(data + b, data + a, args))
312 break;
313 do_swap(data + a, data + b, elem_size, swp, args);
314 a = b;
315 }
316 }
317
318 #define min_heap_sift_up_inline(_heap, _idx, _func, _args) \
319 __min_heap_sift_up_inline(container_of(&(_heap)->nr, min_heap_char, nr), \
320 __minheap_obj_size(_heap), _idx, _func, _args)
321
322 /* Floyd's approach to heapification that is O(nr). */
323 static __always_inline
__min_heapify_all_inline(min_heap_char * heap,size_t elem_size,const struct min_heap_callbacks * func,void * args)324 void __min_heapify_all_inline(min_heap_char *heap, size_t elem_size,
325 const struct min_heap_callbacks *func, void *args)
326 {
327 int i;
328
329 for (i = heap->nr / 2 - 1; i >= 0; i--)
330 __min_heap_sift_down_inline(heap, i, elem_size, func, args);
331 }
332
333 #define min_heapify_all_inline(_heap, _func, _args) \
334 __min_heapify_all_inline(container_of(&(_heap)->nr, min_heap_char, nr), \
335 __minheap_obj_size(_heap), _func, _args)
336
337 /* Remove minimum element from the heap, O(log2(nr)). */
338 static __always_inline
__min_heap_pop_inline(min_heap_char * heap,size_t elem_size,const struct min_heap_callbacks * func,void * args)339 bool __min_heap_pop_inline(min_heap_char *heap, size_t elem_size,
340 const struct min_heap_callbacks *func, void *args)
341 {
342 void *data = heap->data;
343
344 if (WARN_ONCE(heap->nr <= 0, "Popping an empty heap"))
345 return false;
346
347 /* Place last element at the root (position 0) and then sift down. */
348 heap->nr--;
349 memcpy(data, data + (heap->nr * elem_size), elem_size);
350 __min_heap_sift_down_inline(heap, 0, elem_size, func, args);
351
352 return true;
353 }
354
355 #define min_heap_pop_inline(_heap, _func, _args) \
356 __min_heap_pop_inline(container_of(&(_heap)->nr, min_heap_char, nr), \
357 __minheap_obj_size(_heap), _func, _args)
358
359 /*
360 * Remove the minimum element and then push the given element. The
361 * implementation performs 1 sift (O(log2(nr))) and is therefore more
362 * efficient than a pop followed by a push that does 2.
363 */
364 static __always_inline
__min_heap_pop_push_inline(min_heap_char * heap,const void * element,size_t elem_size,const struct min_heap_callbacks * func,void * args)365 void __min_heap_pop_push_inline(min_heap_char *heap, const void *element, size_t elem_size,
366 const struct min_heap_callbacks *func, void *args)
367 {
368 memcpy(heap->data, element, elem_size);
369 __min_heap_sift_down_inline(heap, 0, elem_size, func, args);
370 }
371
372 #define min_heap_pop_push_inline(_heap, _element, _func, _args) \
373 __min_heap_pop_push_inline(container_of(&(_heap)->nr, min_heap_char, nr), _element, \
374 __minheap_obj_size(_heap), _func, _args)
375
376 /* Push an element on to the heap, O(log2(nr)). */
377 static __always_inline
__min_heap_push_inline(min_heap_char * heap,const void * element,size_t elem_size,const struct min_heap_callbacks * func,void * args)378 bool __min_heap_push_inline(min_heap_char *heap, const void *element, size_t elem_size,
379 const struct min_heap_callbacks *func, void *args)
380 {
381 void *data = heap->data;
382 int pos;
383
384 if (WARN_ONCE(heap->nr >= heap->size, "Pushing on a full heap"))
385 return false;
386
387 /* Place at the end of data. */
388 pos = heap->nr;
389 memcpy(data + (pos * elem_size), element, elem_size);
390 heap->nr++;
391
392 /* Sift child at pos up. */
393 __min_heap_sift_up_inline(heap, elem_size, pos, func, args);
394
395 return true;
396 }
397
398 #define min_heap_push_inline(_heap, _element, _func, _args) \
399 __min_heap_push_inline(container_of(&(_heap)->nr, min_heap_char, nr), _element, \
400 __minheap_obj_size(_heap), _func, _args)
401
402 /* Remove ith element from the heap, O(log2(nr)). */
403 static __always_inline
__min_heap_del_inline(min_heap_char * heap,size_t elem_size,size_t idx,const struct min_heap_callbacks * func,void * args)404 bool __min_heap_del_inline(min_heap_char *heap, size_t elem_size, size_t idx,
405 const struct min_heap_callbacks *func, void *args)
406 {
407 void *data = heap->data;
408 void (*swp)(void *lhs, void *rhs, void *args) = func->swp;
409
410 if (WARN_ONCE(heap->nr <= 0, "Popping an empty heap"))
411 return false;
412
413 if (!swp)
414 swp = select_swap_func(data, elem_size);
415
416 /* Place last element at the root (position 0) and then sift down. */
417 heap->nr--;
418 if (idx == heap->nr)
419 return true;
420 do_swap(data + (idx * elem_size), data + (heap->nr * elem_size), elem_size, swp, args);
421 __min_heap_sift_up_inline(heap, elem_size, idx, func, args);
422 __min_heap_sift_down_inline(heap, idx, elem_size, func, args);
423
424 return true;
425 }
426
427 #define min_heap_del_inline(_heap, _idx, _func, _args) \
428 __min_heap_del_inline(container_of(&(_heap)->nr, min_heap_char, nr), \
429 __minheap_obj_size(_heap), _idx, _func, _args)
430
431 void __min_heap_init(min_heap_char *heap, void *data, int size);
432 void *__min_heap_peek(struct min_heap_char *heap);
433 bool __min_heap_full(min_heap_char *heap);
434 void __min_heap_sift_down(min_heap_char *heap, int pos, size_t elem_size,
435 const struct min_heap_callbacks *func, void *args);
436 void __min_heap_sift_up(min_heap_char *heap, size_t elem_size, size_t idx,
437 const struct min_heap_callbacks *func, void *args);
438 void __min_heapify_all(min_heap_char *heap, size_t elem_size,
439 const struct min_heap_callbacks *func, void *args);
440 bool __min_heap_pop(min_heap_char *heap, size_t elem_size,
441 const struct min_heap_callbacks *func, void *args);
442 void __min_heap_pop_push(min_heap_char *heap, const void *element, size_t elem_size,
443 const struct min_heap_callbacks *func, void *args);
444 bool __min_heap_push(min_heap_char *heap, const void *element, size_t elem_size,
445 const struct min_heap_callbacks *func, void *args);
446 bool __min_heap_del(min_heap_char *heap, size_t elem_size, size_t idx,
447 const struct min_heap_callbacks *func, void *args);
448
449 #define min_heap_init(_heap, _data, _size) \
450 __min_heap_init(container_of(&(_heap)->nr, min_heap_char, nr), _data, _size)
451 #define min_heap_peek(_heap) \
452 (__minheap_cast(_heap) __min_heap_peek(container_of(&(_heap)->nr, min_heap_char, nr)))
453 #define min_heap_full(_heap) \
454 __min_heap_full(container_of(&(_heap)->nr, min_heap_char, nr))
455 #define min_heap_sift_down(_heap, _pos, _func, _args) \
456 __min_heap_sift_down(container_of(&(_heap)->nr, min_heap_char, nr), _pos, \
457 __minheap_obj_size(_heap), _func, _args)
458 #define min_heap_sift_up(_heap, _idx, _func, _args) \
459 __min_heap_sift_up(container_of(&(_heap)->nr, min_heap_char, nr), \
460 __minheap_obj_size(_heap), _idx, _func, _args)
461 #define min_heapify_all(_heap, _func, _args) \
462 __min_heapify_all(container_of(&(_heap)->nr, min_heap_char, nr), \
463 __minheap_obj_size(_heap), _func, _args)
464 #define min_heap_pop(_heap, _func, _args) \
465 __min_heap_pop(container_of(&(_heap)->nr, min_heap_char, nr), \
466 __minheap_obj_size(_heap), _func, _args)
467 #define min_heap_pop_push(_heap, _element, _func, _args) \
468 __min_heap_pop_push(container_of(&(_heap)->nr, min_heap_char, nr), _element, \
469 __minheap_obj_size(_heap), _func, _args)
470 #define min_heap_push(_heap, _element, _func, _args) \
471 __min_heap_push(container_of(&(_heap)->nr, min_heap_char, nr), _element, \
472 __minheap_obj_size(_heap), _func, _args)
473 #define min_heap_del(_heap, _idx, _func, _args) \
474 __min_heap_del(container_of(&(_heap)->nr, min_heap_char, nr), \
475 __minheap_obj_size(_heap), _idx, _func, _args)
476
477 #endif /* _LINUX_MIN_HEAP_H */
478