1#include <console/panic.h>
2#include <mem/alloc.h>
3#include <string.h>
4#include <structures/minheap.h>
5#include <sync/spinlock.h>
6
7static void minheap_swap(struct minheap *heap, uint32_t a, uint32_t b) {
8 struct minheap_node *tmp = heap->nodes[a];
9 heap->nodes[a] = heap->nodes[b];
10 heap->nodes[b] = tmp;
11
12 MINHEAP_NODE_SET_INDEX(heap->nodes[a], a);
13 MINHEAP_NODE_SET_INDEX(heap->nodes[b], b);
14}
15
16static void minheap_sift_up(struct minheap *heap, uint32_t idx) {
17 while (idx > 0) {
18 uint32_t parent = (idx - 1) / 2;
19 if (heap->nodes[idx]->key >= heap->nodes[parent]->key)
20 break;
21
22 minheap_swap(heap, a: idx, b: parent);
23 idx = parent;
24 }
25}
26
27static void minheap_sift_down(struct minheap *heap, uint32_t idx) {
28 uint32_t left, right, smallest;
29
30 while (true) {
31 left = 2 * idx + 1;
32 right = 2 * idx + 2;
33 smallest = idx;
34
35 if (left < heap->size &&
36 heap->nodes[left]->key < heap->nodes[smallest]->key)
37 smallest = left;
38
39 if (right < heap->size &&
40 heap->nodes[right]->key < heap->nodes[smallest]->key)
41 smallest = right;
42
43 if (smallest == idx)
44 break;
45
46 minheap_swap(heap, a: idx, b: smallest);
47 idx = smallest;
48 }
49}
50
51struct minheap *minheap_create(void) {
52 struct minheap *heap = kmalloc(sizeof(struct minheap));
53 heap->capacity = MINHEAP_INIT_CAP;
54 heap->size = 0;
55 heap->nodes = kmalloc(sizeof(struct minheap_node *) * heap->capacity,
56 ALLOC_FLAGS_ZERO);
57 spinlock_init(lock: &heap->lock);
58 return heap;
59}
60
61void minheap_expand(struct minheap *heap, uint32_t new_size) {
62 if (new_size <= heap->capacity)
63 return;
64
65 struct minheap_node **new_nodes =
66 kmalloc(sizeof(struct minheap_node *) * new_size);
67
68 if (!new_nodes)
69 return;
70
71 memcpy(new_nodes, heap->nodes,
72 sizeof(struct minheap_node *) * heap->capacity);
73
74 kfree(heap->nodes);
75 heap->nodes = new_nodes;
76 MINHEAP_SET_CAPACITY(heap, new_size);
77}
78
79void minheap_insert(struct minheap *heap, struct minheap_node *node,
80 uint64_t key) {
81 enum irql irql = spin_lock(lock: &heap->lock);
82 if (heap->size >= heap->capacity) {
83 uint32_t new_cap = heap->capacity * 2;
84 struct minheap_node **new_nodes =
85 kmalloc(sizeof(struct minheap_node *) * new_cap);
86
87 if (!new_nodes)
88 return;
89
90 memcpy(new_nodes, heap->nodes,
91 sizeof(struct minheap_node *) * heap->capacity);
92 kfree(heap->nodes);
93 heap->nodes = new_nodes;
94 heap->capacity = new_cap;
95 }
96
97 MINHEAP_NODE_SET_INDEX(node, heap->size);
98 MINHEAP_NODE_SET_KEY(node, key);
99
100 heap->nodes[heap->size++] = node;
101
102 minheap_sift_up(heap, idx: node->index);
103 spin_unlock(lock: &heap->lock, old: irql);
104}
105
106void minheap_remove(struct minheap *heap, struct minheap_node *node) {
107 enum irql irql = spin_lock(lock: &heap->lock);
108 uint32_t idx = MINHEAP_NODE_INDEX(node);
109
110 if (idx >= MINHEAP_SIZE(heap)) {
111 spin_unlock(lock: &heap->lock, old: irql);
112 return;
113 }
114
115 heap->size--;
116 if (idx != heap->size) {
117 heap->nodes[idx] = heap->nodes[heap->size];
118 heap->nodes[idx]->index = idx;
119 minheap_sift_down(heap, idx);
120 minheap_sift_up(heap, idx);
121 }
122
123 MINHEAP_MARK_NODE_INVALID(node);
124 spin_unlock(lock: &heap->lock, old: irql);
125}
126
127struct minheap_node *minheap_pop(struct minheap *heap) {
128 if (MINHEAP_SIZE(heap) == 0)
129 return NULL;
130
131 struct minheap_node *top = heap->nodes[0];
132 heap->size--;
133 if (heap->size > 0) {
134 heap->nodes[0] = heap->nodes[heap->size];
135 heap->nodes[0]->index = 0;
136 minheap_sift_down(heap, idx: 0);
137 }
138
139 MINHEAP_MARK_NODE_INVALID(top);
140 return top;
141}
142