From 36eefefc9a9597fe9671aa084903a699d16a1144 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Thu, 20 Jul 2023 09:52:20 -0700
Subject: [PATCH v5 2/3] expand binaryheap api

---
 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 a6be4b42ae..d24a61fd2b 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 ae0af68f0d..84c08d7ebb 100644
--- a/src/include/lib/binaryheap.h
+++ b/src/include/lib/binaryheap.h
@@ -56,8 +56,11 @@ extern void binaryheap_build(binaryheap *heap);
 extern void binaryheap_add(binaryheap *heap, bh_elem_type d);
 extern bh_elem_type binaryheap_first(binaryheap *heap);
 extern bh_elem_type binaryheap_remove_first(binaryheap *heap);
+extern void binaryheap_remove_node(binaryheap *heap, int n);
 extern void binaryheap_replace_first(binaryheap *heap, bh_elem_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