From 85a0f0fe2d10dc524d0d52e53969ab02c0694042 Mon Sep 17 00:00:00 2001
From: Amit Kapila <amit.kapila@enterprisedb.com>
Date: Mon, 27 May 2019 14:02:23 +0530
Subject: [PATCH] Extend binary heap functionality.

Add the routines to allocate binary heap in shared memory and to remove
nth element from binray heap.  This routines will be used by latter commit
to add support for undo workers.

Author: Kuntal Ghosh and Amit Kapila
---
 src/backend/lib/binaryheap.c | 118 +++++++++++++++++++++++++++++++++++
 src/include/lib/binaryheap.h |  12 +++-
 2 files changed, 128 insertions(+), 2 deletions(-)

diff --git a/src/backend/lib/binaryheap.c b/src/backend/lib/binaryheap.c
index a2c8967813..5f7454d05f 100644
--- a/src/backend/lib/binaryheap.c
+++ b/src/backend/lib/binaryheap.c
@@ -16,6 +16,7 @@
 #include <math.h>
 
 #include "lib/binaryheap.h"
+#include "storage/shmem.h"
 
 static void sift_down(binaryheap *heap, int node_off);
 static void sift_up(binaryheap *heap, int node_off);
@@ -47,6 +48,36 @@ binaryheap_allocate(int capacity, binaryheap_comparator compare, void *arg)
 	return heap;
 }
 
+/*
+ * binaryheap_allocate_shm
+ *
+ * It works same as binaryheap_allocate except that the heap will be created
+ * in shared memory.
+ */
+binaryheap *
+binaryheap_allocate_shm(const char *name, int capacity,
+						binaryheap_comparator compare, void *arg)
+{
+	Size		sz;
+	binaryheap *heap;
+	bool		foundBHeap;
+
+	sz = binaryheap_shmem_size(capacity);
+	heap = (binaryheap *) ShmemInitStruct(name, sz, &foundBHeap);
+
+	if (!foundBHeap)
+	{
+		heap->bh_space = capacity;
+		heap->bh_compare = compare;
+		heap->bh_arg = arg;
+
+		heap->bh_size = 0;
+		heap->bh_has_heap_property = true;
+	}
+
+	return heap;
+}
+
 /*
  * binaryheap_reset
  *
@@ -211,6 +242,79 @@ binaryheap_replace_first(binaryheap *heap, Datum d)
 		sift_down(heap, 0);
 }
 
+/*
+ * binaryheap_nth
+ *
+ * Returns a pointer to the nth (0-based) node in the heap without modifying
+ * the heap in O(1).  The caller must ensure that this routine is not used on
+ * an empty heap and is not called with n greater than or equal to the heap
+ * size.
+ */
+Datum
+binaryheap_nth(binaryheap *heap, int n)
+{
+	Assert(!binaryheap_empty(heap));
+	Assert(n < heap->bh_size);
+	return heap->bh_nodes[n];
+}
+
+/*
+ * binaryheap_remove_nth
+ *
+ * Removes the nth node (0-based) in the heap and returns a
+ * pointer to it after rebalancing the heap. The caller must ensure
+ * that this routine is not used on an empty heap.  O(log n) worst
+ * case.
+ */
+Datum
+binaryheap_remove_nth(binaryheap *heap, int n)
+{
+	Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);
+	Assert(n < heap->bh_size);
+
+
+	if (n == heap->bh_size - 1)
+	{
+		heap->bh_size--;
+		return heap->bh_nodes[heap->bh_size];
+	}
+
+	swap_nodes(heap, n, heap->bh_size - 1);
+	heap->bh_size--;
+	sift_down(heap, n);
+
+	return heap->bh_nodes[heap->bh_size];
+}
+
+/*
+ * binaryheap_remove_nth_unordered
+ *
+ * Removes the nth node (0-based) in the heap and returns a pointer to it in
+ * O(1) without preserving the heap property.  This is a convenience routine
+ * to remove elements quickly.  To obtain a valid heap, one must call
+ * binaryheap_build() afterwards.  The caller must ensure that this routine is
+ * not used on an empty heap.
+ */
+Datum
+binaryheap_remove_nth_unordered(binaryheap *heap, int n)
+{
+	Assert(!binaryheap_empty(heap));
+	Assert(n < heap->bh_size);
+
+	heap->bh_has_heap_property = false;
+
+	if (n == heap->bh_size - 1)
+	{
+		heap->bh_size--;
+		return heap->bh_nodes[heap->bh_size];
+	}
+
+	swap_nodes(heap, n, heap->bh_size - 1);
+	heap->bh_size--;
+
+	return heap->bh_nodes[heap->bh_size];
+}
+
 /*
  * Swap the contents of two nodes.
  */
@@ -305,3 +409,17 @@ sift_down(binaryheap *heap, int node_off)
 		node_off = swap_off;
 	}
 }
+
+/*
+ * Compute the size required by binary heap structure.
+ */
+Size
+binaryheap_shmem_size(int capacity)
+{
+	Size		sz;
+
+	sz = add_size(offsetof(binaryheap, bh_nodes),
+				  mul_size(sizeof(Datum), capacity));
+
+	return sz;
+}
diff --git a/src/include/lib/binaryheap.h b/src/include/lib/binaryheap.h
index 21f96b9ca6..ed9e8e8661 100644
--- a/src/include/lib/binaryheap.h
+++ b/src/include/lib/binaryheap.h
@@ -38,8 +38,11 @@ typedef struct binaryheap
 } binaryheap;
 
 extern binaryheap *binaryheap_allocate(int capacity,
-									   binaryheap_comparator compare,
-									   void *arg);
+					binaryheap_comparator compare,
+					void *arg);
+extern binaryheap *binaryheap_allocate_shm(const char *name, int capacity,
+					binaryheap_comparator compare,
+					void *arg);
 extern void binaryheap_reset(binaryheap *heap);
 extern void binaryheap_free(binaryheap *heap);
 extern void binaryheap_add_unordered(binaryheap *heap, Datum d);
@@ -48,7 +51,12 @@ extern void binaryheap_add(binaryheap *heap, Datum d);
 extern Datum binaryheap_first(binaryheap *heap);
 extern Datum binaryheap_remove_first(binaryheap *heap);
 extern void binaryheap_replace_first(binaryheap *heap, Datum d);
+extern Datum binaryheap_nth(binaryheap *heap, int n);
+extern Datum binaryheap_remove_nth(binaryheap *heap, int n);
+extern Datum binaryheap_remove_nth_unordered(binaryheap *heap, int n);
+extern Size binaryheap_shmem_size(int capacity);
 
 #define binaryheap_empty(h)			((h)->bh_size == 0)
+#define binaryheap_cur_size(h)		((h)->bh_size)
 
 #endif							/* BINARYHEAP_H */
-- 
2.30.2