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 (MINHEAP_NODE_KEY(heap->nodes[idx]) >=
20 MINHEAP_NODE_KEY(heap->nodes[parent]))
21 break;
22
23 minheap_swap(heap, a: idx, b: parent);
24 idx = parent;
25 }
26}
27
28static void minheap_sift_down(struct minheap *heap, uint32_t idx) {
29 uint32_t left, right, smallest;
30
31 while (true) {
32 left = 2 * idx + 1;
33 right = 2 * idx + 2;
34 smallest = idx;
35
36 uint32_t heap_size = MINHEAP_SIZE(heap);
37 if (left < heap_size && MINHEAP_NODE_KEY(heap->nodes[left]) <
38 MINHEAP_NODE_KEY(heap->nodes[smallest]))
39 smallest = left;
40
41 if (right < heap_size && MINHEAP_NODE_KEY(heap->nodes[right]) <
42 MINHEAP_NODE_KEY(heap->nodes[smallest]))
43 smallest = right;
44
45 if (smallest == idx)
46 break;
47
48 minheap_swap(heap, a: idx, b: smallest);
49 idx = smallest;
50 }
51}
52
53struct minheap *minheap_create(void) {
54 struct minheap *heap = vnew(heap);
55 if (!heap)
56 return NULL;
57
58 MINHEAP_SET_CAPACITY(heap, MINHEAP_INIT_CAP);
59 MINHEAP_SET_SIZE(heap, 0);
60 heap->nodes = vnews(heap->nodes, MINHEAP_CAPACITY(heap), ALLOC_ZERO);
61 if (!heap->nodes) {
62 vfree(heap);
63 return NULL;
64 }
65
66 spinlock_init(&heap->lock);
67 return heap;
68}
69
70void minheap_expand(struct minheap *heap, uint32_t new_size) {
71 if (new_size <= MINHEAP_CAPACITY(heap))
72 return;
73
74 struct minheap_node **new_nodes = vnews(new_nodes, new_size);
75
76 if (!new_nodes)
77 return;
78
79 memcpy(new_nodes, heap->nodes,
80 sizeof(struct minheap_node *) * MINHEAP_CAPACITY(heap));
81
82 vfree(heap->nodes);
83 heap->nodes = new_nodes;
84 MINHEAP_SET_CAPACITY(heap, new_size);
85}
86
87void minheap_insert(struct minheap *heap, struct minheap_node *node,
88 uint64_t key) {
89 enum irql irql = spin_lock(&heap->lock);
90 uint32_t heap_cap = MINHEAP_CAPACITY(heap);
91 uint32_t heap_size = MINHEAP_SIZE(heap);
92 if (heap_size >= heap_cap) {
93 uint32_t new_cap = heap_cap * 2;
94 struct minheap_node **new_nodes = vnews(new_nodes, new_cap);
95
96 if (!new_nodes) {
97 spin_unlock(&heap->lock, irql);
98 return;
99 }
100
101 memcpy(new_nodes, heap->nodes,
102 sizeof(struct minheap_node *) * heap_cap);
103 vfree(heap->nodes);
104 heap->nodes = new_nodes;
105 MINHEAP_SET_CAPACITY(heap, new_cap);
106 }
107
108 MINHEAP_NODE_SET_INDEX(node, heap_size);
109 MINHEAP_NODE_SET_KEY(node, key);
110
111 heap->nodes[heap_size] = node;
112 MINHEAP_SET_SIZE(heap, heap_size + 1);
113
114 minheap_sift_up(heap, MINHEAP_NODE_INDEX(node));
115 spin_unlock(&heap->lock, irql);
116}
117
118void minheap_remove(struct minheap *heap, struct minheap_node *node) {
119 enum irql irql = spin_lock(&heap->lock);
120 uint32_t idx = MINHEAP_NODE_INDEX(node);
121 uint32_t heap_size = MINHEAP_SIZE(heap);
122
123 if (idx >= heap_size) {
124 spin_unlock(&heap->lock, irql);
125 return;
126 }
127
128 heap_size--;
129 MINHEAP_SET_SIZE(heap, heap_size);
130 if (idx != heap_size) {
131 heap->nodes[idx] = heap->nodes[heap_size];
132 MINHEAP_NODE_SET_INDEX(heap->nodes[idx], idx);
133 minheap_sift_down(heap, idx);
134 minheap_sift_up(heap, idx);
135 }
136
137 MINHEAP_MARK_NODE_INVALID(node);
138 spin_unlock(&heap->lock, irql);
139}
140
141struct minheap_node *minheap_pop(struct minheap *heap) {
142 uint32_t heap_size = MINHEAP_SIZE(heap);
143 if (heap_size == 0)
144 return NULL;
145
146 struct minheap_node *top = heap->nodes[0];
147 heap_size--;
148 MINHEAP_SET_SIZE(heap, heap_size);
149 if (heap_size > 0) {
150 heap->nodes[0] = heap->nodes[heap_size];
151 MINHEAP_NODE_SET_INDEX(heap->nodes[0], 0);
152 minheap_sift_down(heap, idx: 0);
153 }
154
155 MINHEAP_MARK_NODE_INVALID(top);
156 return top;
157}
158