/* 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;
}