min_heap.h (3377B)
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, *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). */ 67static __always_inline 68void 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)). */ 78static __always_inline 79void 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 */ 98static __always_inline 99void 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)). */ 108static __always_inline 109void 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 */