1#include "pq.h"
2#include <stddef.h>
3#include <stdint.h>
4#include <stdio.h>
5#include <stdlib.h>
6#include <string.h>
7
8struct pq_entry {
9    uint64_t priority;
10    unsigned char data[];
11};
12
13struct pq {
14    size_t data_size;
15    size_t capacity;
16    size_t size;
17    unsigned char * entries; // Array of `struct pq_entry`
18};
19
20#define PQ_PARENT(I)    (I == 0 ? 0 : ((I - 1) / 2))
21#define PQ_LEFT(I)      (I * 2 + 1)
22#define PQ_RIGHT(I)     (I * 2 + 2)
23
24#define PQ_ENTRY(PQ, I)  ((struct pq_entry *) &(PQ)->entries[I * PQ_ENTRY_SIZE(PQ)])
25#define PQ_ENTRY_SIZE(PQ) (sizeof(uint64_t) + (PQ)->data_size)
26
27struct pq *
28pq_create(size_t data_size) {
29    struct pq * pq = calloc(1, sizeof *pq);
30    if (pq == NULL) {
31        return NULL;
32    }
33
34    pq->data_size = data_size;
35    pq->capacity = 1024;
36    pq->size = 0;
37    pq->entries = calloc(pq->capacity, PQ_ENTRY_SIZE(pq));
38    if (pq->entries == NULL) {
39        return free(pq), NULL;
40    }
41
42    return pq;
43}
44
45void
46pq_destroy(struct pq * pq) {
47    free(pq->entries);
48    free(pq);
49}
50
51void
52pq_clear(struct pq * pq) {
53    pq->size = 0;
54
55#ifdef PQ_DEBUG
56    memset(pq->entries, 0, PQ_ENTRY_SIZE(pq) * pq->capacity);
57#endif
58}
59
60int
61pq_push(struct pq * pq, uint64_t priority, const void * data) {
62    size_t idx = pq->size++;
63
64    if (pq->size > pq->capacity) {
65        size_t new_capacity = pq->capacity * 2;
66        void * new_entries = realloc(pq->entries, new_capacity * PQ_ENTRY_SIZE(pq));
67        if (new_entries == NULL) {
68            return -1;
69        }
70        pq->entries = new_entries;
71        memset(PQ_ENTRY(pq, pq->capacity), 0, PQ_ENTRY_SIZE(pq) * pq->capacity);
72        pq->capacity = new_capacity;
73    }
74
75    struct pq_entry * entry = PQ_ENTRY(pq, idx);
76    while (idx > 0) {
77        size_t parent_idx = PQ_PARENT(idx);
78        struct pq_entry * parent = PQ_ENTRY(pq, parent_idx);
79        if (parent->priority <= priority) {
80            break;
81        }
82
83        memcpy(entry, parent, PQ_ENTRY_SIZE(pq));
84        idx = parent_idx;
85        entry = parent;
86    }
87
88    entry->priority = priority;
89    if (data != NULL) {
90        memcpy(entry->data, data, pq->data_size);
91    }
92
93    return 0;
94}
95
96int
97pq_pop(struct pq * pq, uint64_t * priority_out, void * data_out) {
98    if (pq->size == 0)
99        return -1;
100
101    size_t idx = 0;
102    struct pq_entry * entry = PQ_ENTRY(pq, idx);
103    if (priority_out != NULL) {
104        *priority_out = entry->priority;
105    }
106    if (data_out != NULL) {
107        memcpy(data_out, entry->data, pq->data_size);
108    }
109
110    size_t child_idx = --pq->size;
111    uint64_t priority = PQ_ENTRY(pq, child_idx)->priority;
112    while (idx < pq->size) {
113        size_t left_idx = PQ_LEFT(idx); 
114        size_t right_idx = PQ_RIGHT(idx); 
115
116        size_t smallest_idx = idx;
117        size_t smallest_priority = priority;
118        if (left_idx < pq->size) {
119            struct pq_entry * left_entry = PQ_ENTRY(pq, left_idx);
120            if (left_entry->priority < smallest_priority) {
121                smallest_priority = left_entry->priority;
122                smallest_idx = left_idx;
123            }
124        }
125        if (right_idx < pq->size) {
126            struct pq_entry * right_entry = PQ_ENTRY(pq, right_idx);
127            if (right_entry->priority < smallest_priority) {
128                smallest_priority = right_entry->priority;
129                smallest_idx = right_idx;
130            }
131        }
132        if (smallest_idx == idx) {
133            break;
134        }
135        memcpy(PQ_ENTRY(pq, idx), PQ_ENTRY(pq, smallest_idx), PQ_ENTRY_SIZE(pq));
136        idx = smallest_idx;
137    }
138    if (child_idx != idx) {
139        memcpy(PQ_ENTRY(pq, idx), PQ_ENTRY(pq, child_idx), PQ_ENTRY_SIZE(pq));
140    }
141
142#ifdef PQ_DEBUG
143    for (size_t i = 0; i < pq->size; i++) {
144        entry = PQ_ENTRY(pq, i);
145        if (entry->priority < *priority_out)
146            printf("Found entry with priority %zu < %zu\n", entry->priority, *priority_out);
147    }
148#endif
149
150    return 0;
151}
152
153size_t
154pq_size(struct pq * pq) {
155    return pq->size;
156}
157
158#ifdef PQ_DEBUG
159void
160pq_print(struct pq * pq) {
161    for (size_t i = 0; i < pq->size; i++) {
162        struct pq_entry * entry = PQ_ENTRY(pq, i);
163        struct pq_entry * parent = PQ_ENTRY(pq, PQ_PARENT(i));
164        if (parent->priority > entry->priority) {
165            printf("!");
166        }
167        printf("%lu ", entry->priority);
168    }
169    printf("\n");
170}
171#endif