From 8e01f7140bd0c70cceb586b5fb4d199b38cf1956 Mon Sep 17 00:00:00 2001 From: Nathan Bossart Date: Thu, 20 Jul 2023 09:52:20 -0700 Subject: [PATCH v6 2/4] Add function for removing arbitrary nodes in binaryheap. This commit introduces binaryheap_remove_node(), which can be used to remove any node from a binary heap. The implementation is straightforward. The target node is replaced with the last node in the heap, and then we sift as needed to preserve the heap property. This new function is intended for use in a follow-up commit that will improve the performance of pg_restore. Discussion: https://postgr.es/m/3612876.1689443232%40sss.pgh.pa.us --- src/common/binaryheap.c | 29 +++++++++++++++++++++++++++++ src/include/lib/binaryheap.h | 3 +++ 2 files changed, 32 insertions(+) diff --git a/src/common/binaryheap.c b/src/common/binaryheap.c index 4836cd1c6d..8c6cc6e5e2 100644 --- a/src/common/binaryheap.c +++ b/src/common/binaryheap.c @@ -211,6 +211,35 @@ binaryheap_remove_first(binaryheap *heap) return result; } +/* + * binaryheap_remove_node + * + * Removes the nth node from the heap. The caller must ensure that there are + * at least (n - 1) nodes in the heap. O(log n) worst case. + */ +void +binaryheap_remove_node(binaryheap *heap, int n) +{ + int cmp; + + Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property); + Assert(n >= 0 && n < heap->bh_size); + + /* compare last node to the one that is being removed */ + cmp = heap->bh_compare(heap->bh_nodes[--heap->bh_size], + heap->bh_nodes[n], + heap->bh_arg); + + /* remove the last node, placing it in the vacated entry */ + heap->bh_nodes[n] = heap->bh_nodes[heap->bh_size]; + + /* sift as needed to preserve the heap property */ + if (cmp > 0) + sift_up(heap, n); + else if (cmp < 0) + sift_down(heap, n); +} + /* * binaryheap_replace_first * diff --git a/src/include/lib/binaryheap.h b/src/include/lib/binaryheap.h index 4756f43a6d..7060c77a52 100644 --- a/src/include/lib/binaryheap.h +++ b/src/include/lib/binaryheap.h @@ -59,8 +59,11 @@ extern void binaryheap_build(binaryheap *heap); extern void binaryheap_add(binaryheap *heap, bh_node_type d); extern bh_node_type binaryheap_first(binaryheap *heap); extern bh_node_type binaryheap_remove_first(binaryheap *heap); +extern void binaryheap_remove_node(binaryheap *heap, int n); extern void binaryheap_replace_first(binaryheap *heap, bh_node_type d); #define binaryheap_empty(h) ((h)->bh_size == 0) +#define binaryheap_size(h) ((h)->bh_size) +#define binaryheap_get_node(h, n) ((h)->bh_nodes[n]) #endif /* BINARYHEAP_H */ -- 2.25.1