diff --git a/prio-queue.c b/prio-queue.c index ec33ac27db..9748528ce6 100644 --- a/prio-queue.c +++ b/prio-queue.c @@ -58,22 +58,10 @@ void prio_queue_put(struct prio_queue *queue, void *thing) } } -void *prio_queue_get(struct prio_queue *queue) +static void sift_down_root(struct prio_queue *queue) { - void *result; size_t ix, child; - if (!queue->nr) - return NULL; - if (!queue->compare) - return queue->array[--queue->nr].data; /* LIFO */ - - result = queue->array[0].data; - if (!--queue->nr) - return result; - - queue->array[0] = queue->array[queue->nr]; - /* Push down the one at the root */ for (ix = 0; ix * 2 + 1 < queue->nr; ix = child) { child = ix * 2 + 1; /* left */ @@ -86,6 +74,23 @@ void *prio_queue_get(struct prio_queue *queue) swap(queue, child, ix); } +} + +void *prio_queue_get(struct prio_queue *queue) +{ + void *result; + + if (!queue->nr) + return NULL; + if (!queue->compare) + return queue->array[--queue->nr].data; /* LIFO */ + + result = queue->array[0].data; + if (!--queue->nr) + return result; + + queue->array[0] = queue->array[queue->nr]; + sift_down_root(queue); return result; } @@ -97,3 +102,17 @@ void *prio_queue_peek(struct prio_queue *queue) return queue->array[queue->nr - 1].data; return queue->array[0].data; } + +void prio_queue_replace(struct prio_queue *queue, void *thing) +{ + if (!queue->nr) { + prio_queue_put(queue, thing); + } else if (!queue->compare) { + queue->array[queue->nr - 1].ctr = queue->insertion_ctr++; + queue->array[queue->nr - 1].data = thing; + } else { + queue->array[0].ctr = queue->insertion_ctr++; + queue->array[0].data = thing; + sift_down_root(queue); + } +} diff --git a/prio-queue.h b/prio-queue.h index 38d032636d..da7fad2f1f 100644 --- a/prio-queue.h +++ b/prio-queue.h @@ -52,6 +52,14 @@ void *prio_queue_get(struct prio_queue *); */ void *prio_queue_peek(struct prio_queue *); +/* + * Replace the "thing" that compares the smallest with a new "thing", + * like prio_queue_get()+prio_queue_put() would do, but in a more + * efficient way. Does the same as prio_queue_put() if the queue is + * empty. + */ +void prio_queue_replace(struct prio_queue *queue, void *thing); + void clear_prio_queue(struct prio_queue *); /* Reverse the LIFO elements */ diff --git a/t/unit-tests/u-prio-queue.c b/t/unit-tests/u-prio-queue.c index 145e689c9c..63e58114ae 100644 --- a/t/unit-tests/u-prio-queue.c +++ b/t/unit-tests/u-prio-queue.c @@ -13,6 +13,7 @@ static int intcmp(const void *va, const void *vb, void *data UNUSED) #define STACK -3 #define GET -4 #define REVERSE -5 +#define REPLACE -6 static int show(int *v) { @@ -51,6 +52,15 @@ static void test_prio_queue(int *input, size_t input_size, case REVERSE: prio_queue_reverse(&pq); break; + case REPLACE: + peek = prio_queue_peek(&pq); + cl_assert(i + 1 < input_size); + cl_assert(input[i + 1] >= 0); + cl_assert(j < result_size); + cl_assert_equal_i(result[j], show(peek)); + j++; + prio_queue_replace(&pq, &input[++i]); + break; default: prio_queue_put(&pq, &input[i]); break; @@ -81,6 +91,13 @@ void test_prio_queue__empty(void) ((int []){ 1, 2, MISSING, 1, 2, MISSING })); } +void test_prio_queue__replace(void) +{ + TEST_INPUT(((int []){ REPLACE, 6, 2, 4, REPLACE, 5, 7, GET, + REPLACE, 1, DUMP }), + ((int []){ MISSING, 2, 4, 5, 1, 6, 7 })); +} + void test_prio_queue__stack(void) { TEST_INPUT(((int []){ STACK, 8, 1, 5, 4, 6, 2, 3, DUMP }), @@ -92,3 +109,9 @@ void test_prio_queue__reverse_stack(void) TEST_INPUT(((int []){ STACK, 1, 2, 3, 4, 5, 6, REVERSE, DUMP }), ((int []){ 1, 2, 3, 4, 5, 6 })); } + +void test_prio_queue__replace_stack(void) +{ + TEST_INPUT(((int []){ STACK, 8, 1, 5, REPLACE, 4, 6, 2, 3, DUMP }), + ((int []){ 5, 3, 2, 6, 4, 1, 8 })); +}