feat: Implement min-heap data structure

3 files changed, 158 insertions, 0 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 750d7b6..68a9b98 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -8,6 +8,7 @@ project(mpu
 include(cmake/Diagnostics.cmake)
 
 add_executable(mpu
+  "src/heap.c"
   "src/p2p.c"
   "src/main.c")
 
diff --git a/src/heap.c b/src/heap.c
new file mode 100644
index 0000000..d0aed61
--- /dev/null
+++ b/src/heap.c
@@ -0,0 +1,104 @@
+/* SPDX-License-Identifier: GPL-3.0-or-later */
+
+#include "heap.h"
+
+#include <stdint.h>
+#include <stdlib.h>
+
+int heap_init(heap_t* heap, size_t capacity, size_t user_sizeof) {
+    heap->nodes = malloc(capacity * (sizeof(heap_node_t) + user_sizeof));
+    if (!heap->nodes) {
+        return -1;
+    }
+
+    void* ptr = heap->nodes + capacity;
+    for (size_t i = 0; i < capacity; i++) {
+        heap->nodes[i].value = 0;
+        heap->nodes[i].user = ptr;
+        ptr += user_sizeof;
+    }
+
+    heap->capacity = capacity;
+    heap->count = 0;
+    return 0;
+}
+
+#define PARENT(i) (((i) - 1) / 2)
+#define SWAP(a, b) do { \
+    heap_node_t tmp = *(a); \
+    *(a) = *(b); \
+    *(b) = tmp; \
+} while (0)
+
+heap_node_t* heap_next_insert(heap_t* heap) {
+    if (heap->count >= heap->capacity) {
+        return nullptr;
+    }
+
+    return &heap->nodes[heap->count];
+}
+
+void heap_insert(heap_t* heap, uint64_t value) {
+    if (heap->count >= heap->capacity) {
+        return;
+    }
+
+    heap->count++;
+
+    size_t i = heap->count - 1;
+    heap->nodes[i].value = value;
+
+    while (i != 0 && heap->nodes[PARENT(i)].value > heap->nodes[i].value) {
+        SWAP(&heap->nodes[i], &heap->nodes[PARENT(i)]);
+        i = PARENT(i);
+    }
+}
+
+#define LEFT(i) (2 * (i) + 1)
+#define RIGHT(i) (2 * (i) + 2)
+
+heap_node_t* heap_next_extract(heap_t* heap) {
+    if (heap->count == 0) {
+        return nullptr;
+    }
+
+    return &heap->nodes[0];
+}
+
+void heap_extract(heap_t* heap) {
+    if (heap->count <= 1) {
+        heap->count = 0;
+        return;
+    }
+
+    SWAP(&heap->nodes[0], &heap->nodes[heap->count - 1]);
+    heap->count--;
+
+    size_t i = 0;
+    while (true) {
+        size_t smallest = i;
+        size_t left = LEFT(i);
+        size_t right = RIGHT(i);
+
+        if (left < heap->count && heap->nodes[left].value < heap->nodes[smallest].value) {
+            smallest = left;
+        }
+
+        if (right < heap->count && heap->nodes[right].value < heap->nodes[smallest].value) {
+            smallest = right;
+        }
+
+        if (smallest == i) {
+            break;
+        }
+
+        SWAP(&heap->nodes[i], &heap->nodes[smallest]);
+        i = smallest;
+    }
+}
+
+void heap_free(heap_t* heap) {
+    free(heap->nodes);
+    heap->capacity = 0;
+    heap->count = 0;
+}
diff --git a/src/heap.h b/src/heap.h
new file mode 100644
index 0000000..6ae6507
--- /dev/null
+++ b/src/heap.h
@@ -0,0 +1,53 @@
+/* SPDX-License-Identifier: GPL-3.0-or-later */
+
+#pragma once
+
+#include <stddef.h>
+#include <stdint.h>
+
+//
+// This module implements a binary heap, specifically a min-heap. Each node contains a 64-bit
+// value and a pointer to user data (user data is reused when nodes are extracted). All memory
+// is a single contiguous chunk.
+//
+// In a min-heap, the root node always has the smallest value. The root node can be
+// acquired with heap_next_extract() and extracted (removed) with heap_extract().
+// The runtime of heap_extract() is O(log n).
+//
+// Prior to insertion, the node that will be inserted can already be obtained with
+// heap_next_insert(), which MUST be succeeded by heap_insert() to actually insert the node.
+// The runtime of heap_insert() is O(log n).
+//
+
+typedef struct {
+    uint64_t value;
+    void* user;
+} heap_node_t; //!< Node in a heap
+
+typedef struct {
+    heap_node_t* nodes;
+    size_t capacity;
+    size_t count;
+} heap_t; //!< Min-heap data structure
+
+/// Initialize a min-heap with the given capacity. The user_sizeof parameter specifies the size
+/// of user data stored in each node.
+/// Returns 0 on success, -1 on allocation failure.
+int heap_init(heap_t* heap, size_t capacity, size_t user_sizeof);
+
+/// Obtain the heap node that will be inserted next.
+/// Returns heap node on success, or nullptr if the heap is full.
+heap_node_t* heap_next_insert(heap_t* heap);
+
+/// Insert a value into the heap.
+void heap_insert(heap_t* heap, uint64_t value);
+
+/// Peek the minimum value in the heap without removing it.
+/// Returns heap node on success, or nullptr if the heap is empty.
+heap_node_t* heap_next_extract(heap_t* heap);
+
+/// Extract a node from the heap.
+void heap_extract(heap_t* heap);
+
+/// Free all resources associated with the heap.
+void heap_free(heap_t* heap);