min_heap.h

  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)	\
 28struct _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
 37typedef 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 */
 47struct 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
 66static 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 */
 91static __always_inline
 92void 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 */
117static __always_inline
118void 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 */
146static __always_inline
147void 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 */
168static __always_inline
169void *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
179static __always_inline
180void 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
212static 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. */
220static __always_inline
221void __min_heap_init_inline(min_heap_char *heap, void *data, size_t 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. */
235static __always_inline
236void *__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. */
246static __always_inline
247bool __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. */
256static __always_inline
257void __min_heap_sift_down_inline(min_heap_char *heap, size_t 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)). */
296static __always_inline
297void __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). */
323static __always_inline
324void __min_heapify_all_inline(min_heap_char *heap, size_t elem_size,
325			      const struct min_heap_callbacks *func, void *args)
326{
327	ssize_t 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)). */
338static __always_inline
339bool __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 */
364static __always_inline
365void __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)). */
377static __always_inline
378bool __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	size_t 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)). */
403static __always_inline
404bool __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
431void __min_heap_init(min_heap_char *heap, void *data, size_t size);
432void *__min_heap_peek(struct min_heap_char *heap);
433bool __min_heap_full(min_heap_char *heap);
434void __min_heap_sift_down(min_heap_char *heap, size_t pos, size_t elem_size,
435			  const struct min_heap_callbacks *func, void *args);
436void __min_heap_sift_up(min_heap_char *heap, size_t elem_size, size_t idx,
437			const struct min_heap_callbacks *func, void *args);
438void __min_heapify_all(min_heap_char *heap, size_t elem_size,
439		       const struct min_heap_callbacks *func, void *args);
440bool __min_heap_pop(min_heap_char *heap, size_t elem_size,
441		    const struct min_heap_callbacks *func, void *args);
442void __min_heap_pop_push(min_heap_char *heap, const void *element, size_t elem_size,
443			 const struct min_heap_callbacks *func, void *args);
444bool __min_heap_push(min_heap_char *heap, const void *element, size_t elem_size,
445		     const struct min_heap_callbacks *func, void *args);
446bool __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 */