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, 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 */ 137