From 658f4c1899b4721275986e50f388dd5be88e8c20 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Thu, 20 Jul 2023 09:52:20 -0700
Subject: [PATCH v7 3/5] 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.
---
 src/common/binaryheap.c      | 44 ++++++++++++++++++++++++++++++++++++
 src/include/lib/binaryheap.h |  3 +++
 2 files changed, 47 insertions(+)

diff --git a/src/common/binaryheap.c b/src/common/binaryheap.c
index 07021d10b8..f1ed0e157e 100644
--- a/src/common/binaryheap.c
+++ b/src/common/binaryheap.c
@@ -232,6 +232,50 @@ binaryheap_remove_first(binaryheap *heap, void *result)
 	sift_down(heap, 0);
 }
 
+/*
+ * binaryheap_nth
+ *
+ * Returns the nth node from the heap.  The caller must ensure that there are
+ * at least (n - 1) nodes in the heap.  Always O(1).
+ */
+void
+binaryheap_nth(binaryheap *heap, int n, void *result)
+{
+	Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);
+	Assert(n >= 0 && n < heap->bh_size);
+
+	memmove(result, bh_node(heap, n), heap->bh_elem_size);
+}
+
+/*
+ * binaryheap_remove_nth
+ *
+ * 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_nth(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(bh_node(heap, --heap->bh_size),
+						   bh_node(heap, n),
+						   heap->bh_arg);
+
+	/* remove the last node, placing it in the vacated entry */
+	bh_set(heap, n, bh_node(heap, 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 b8c7ef81ef..083821ddbd 100644
--- a/src/include/lib/binaryheap.h
+++ b/src/include/lib/binaryheap.h
@@ -50,8 +50,11 @@ extern void binaryheap_build(binaryheap *heap);
 extern void binaryheap_add(binaryheap *heap, void *d);
 extern void binaryheap_first(binaryheap *heap, void *result);
 extern void binaryheap_remove_first(binaryheap *heap, void *result);
+extern void binaryheap_nth(binaryheap *heap, int n, void *result);
+extern void binaryheap_remove_nth(binaryheap *heap, int n);
 extern void binaryheap_replace_first(binaryheap *heap, void *d);
 
 #define binaryheap_empty(h)			((h)->bh_size == 0)
+#define binaryheap_size(h)			((h)->bh_size)
 
 #endif							/* BINARYHEAP_H */
-- 
2.25.1