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 */ 15struct 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 */ 27struct 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. */ 34static __always_inline 35void 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). */ 69static __always_inline 70void 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, i, func); 77} 78 79/* Remove minimum element from the heap, O(log2(nr)). */ 80static __always_inline 81void 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 */ 100static __always_inline 101void 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)). */ 110static __always_inline 111void 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