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, *parent, *smallest; 39 void *data = heap->data; 40 41 for (;;) { 42 if (pos * 2 + 1 >= heap->nr) 43 break; 44 45 left = data + ((pos * 2 + 1) * func->elem_size); 46 parent = data + (pos * func->elem_size); 47 smallest = parent; 48 if (func->less(left, smallest)) 49 smallest = left; 50 51 if (pos * 2 + 2 < heap->nr) { 52 right = data + ((pos * 2 + 2) * func->elem_size); 53 if (func->less(right, smallest)) 54 smallest = right; 55 } 56 if (smallest == parent) 57 break; 58 func->swp(smallest, parent); 59 if (smallest == left) 60 pos = (pos * 2) + 1; 61 else 62 pos = (pos * 2) + 2; 63 } 64 } 65 66 /* Floyd's approach to heapification that is O(nr). */ 67 static __always_inline 68 void min_heapify_all(struct min_heap *heap, 69 const struct min_heap_callbacks *func) 70 { 71 int i; 72 73 for (i = heap->nr / 2; i >= 0; i--) 74 min_heapify(heap, i, func); 75 } 76 77 /* Remove minimum element from the heap, O(log2(nr)). */ 78 static __always_inline 79 void min_heap_pop(struct min_heap *heap, 80 const struct min_heap_callbacks *func) 81 { 82 void *data = heap->data; 83 84 if (WARN_ONCE(heap->nr <= 0, "Popping an empty heap")) 85 return; 86 87 /* Place last element at the root (position 0) and then sift down. */ 88 heap->nr--; 89 memcpy(data, data + (heap->nr * func->elem_size), func->elem_size); 90 min_heapify(heap, 0, func); 91 } 92 93 /* 94 * Remove the minimum element and then push the given element. The 95 * implementation performs 1 sift (O(log2(nr))) and is therefore more 96 * efficient than a pop followed by a push that does 2. 97 */ 98 static __always_inline 99 void min_heap_pop_push(struct min_heap *heap, 100 const void *element, 101 const struct min_heap_callbacks *func) 102 { 103 memcpy(heap->data, element, func->elem_size); 104 min_heapify(heap, 0, func); 105 } 106 107 /* Push an element on to the heap, O(log2(nr)). */ 108 static __always_inline 109 void min_heap_push(struct min_heap *heap, const void *element, 110 const struct min_heap_callbacks *func) 111 { 112 void *data = heap->data; 113 void *child, *parent; 114 int pos; 115 116 if (WARN_ONCE(heap->nr >= heap->size, "Pushing on a full heap")) 117 return; 118 119 /* Place at the end of data. */ 120 pos = heap->nr; 121 memcpy(data + (pos * func->elem_size), element, func->elem_size); 122 heap->nr++; 123 124 /* Sift child at pos up. */ 125 for (; pos > 0; pos = (pos - 1) / 2) { 126 child = data + (pos * func->elem_size); 127 parent = data + ((pos - 1) / 2) * func->elem_size; 128 if (func->less(parent, child)) 129 break; 130 func->swp(parent, child); 131 } 132 } 133 134 #endif /* _LINUX_MIN_HEAP_H */ 135