/*-------------------------------------------------------------------------

 *

 * binaryheap.c

 *    A simple binary heap implementation

 *

 * Portions Copyright (c) 2012-2025, PostgreSQL Global Development Group

 *

 * IDENTIFICATION

 *    src/common/binaryheap.c

 *

 *-------------------------------------------------------------------------

 */

#ifdef FRONTEND

#include "postgres_fe.h"

#else

#include "postgres.h"

#endif

#include <math.h>

#ifdef FRONTEND

#include "common/logging.h"

#endif

#include "lib/binaryheap.h"

static void sift_down(binaryheap *heap, int node_off);

static void sift_up(binaryheap *heap, int node_off);

/*

 * binaryheap_allocate

 *

 * Returns a pointer to a newly-allocated heap that has the capacity to

 * store the given number of nodes, with the heap property defined by

 * the given comparator function, which will be invoked with the additional

 * argument specified by 'arg'.

 */

binaryheap *

binaryheap_allocate(int capacity, binaryheap_comparator compare, void *arg)

{

  int     sz;

  binaryheap *heap;

  sz = offsetof(binaryheap, bh_nodes) + sizeof(bh_node_type) * capacity;

  heap = (binaryheap *) palloc(sz);

  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

 *

 * Resets the heap to an empty state, losing its data content but not the

 * parameters passed at allocation.

 */

void

binaryheap_reset(binaryheap *heap)

{

  heap->bh_size = 0;

  heap->bh_has_heap_property = true;

}

/*

 * binaryheap_free

 *

 * Releases memory used by the given binaryheap.

 */

void

binaryheap_free(binaryheap *heap)

{

  pfree(heap);

}

/*

 * These utility functions return the offset of the left child, right

 * child, and parent of the node at the given index, respectively.

 *

 * The heap is represented as an array of nodes, with the root node

 * stored at index 0. The left child of node i is at index 2*i+1, and

 * the right child at 2*i+2. The parent of node i is at index (i-1)/2.

 */

static inline int

left_offset(int i)

{

  return 2 * i + 1;

}

static inline int

right_offset(int i)

{

  return 2 * i + 2;

}

static inline int

parent_offset(int i)

{

  return (i - 1) / 2;

}

/*

 * binaryheap_add_unordered

 *

 * Adds the given datum to the end of the heap's list of nodes in O(1) without

 * preserving the heap property. This is a convenience to add elements quickly

 * to a new heap. To obtain a valid heap, one must call binaryheap_build()

 * afterwards.

 */

void

binaryheap_add_unordered(binaryheap *heap, bh_node_type d)

{

  if (heap->bh_size >= heap->bh_space)

  {

#ifdef FRONTEND

    pg_fatal("out of binary heap slots");

#else

    elog(ERROR, "out of binary heap slots");

#endif

  }

  heap->bh_has_heap_property = false;

  heap->bh_nodes[heap->bh_size] = d;

  heap->bh_size++;

}

/*

 * binaryheap_build

 *

 * Assembles a valid heap in O(n) from the nodes added by

 * binaryheap_add_unordered(). Not needed otherwise.

 */

void

binaryheap_build(binaryheap *heap)

{

  int     i;

  for (i = parent_offset(heap->bh_size - 1); i >= 0; i--)

    sift_down(heap, i);

  heap->bh_has_heap_property = true;

}

/*

 * binaryheap_add

 *

 * Adds the given datum to the heap in O(log n) time, while preserving

 * the heap property.

 */

void

binaryheap_add(binaryheap *heap, bh_node_type d)

{

  if (heap->bh_size >= heap->bh_space)

  {

#ifdef FRONTEND

    pg_fatal("out of binary heap slots");

#else

    elog(ERROR, "out of binary heap slots");

#endif

  }

  heap->bh_nodes[heap->bh_size] = d;

  heap->bh_size++;

  sift_up(heap, heap->bh_size - 1);

}

/*

 * binaryheap_first

 *

 * Returns a pointer to the first (root, topmost) node in the heap

 * without modifying the heap. The caller must ensure that this

 * routine is not used on an empty heap. Always O(1).

 */

bh_node_type

binaryheap_first(binaryheap *heap)

{

  Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);

  return heap->bh_nodes[0];

}

/*

 * binaryheap_remove_first

 *

 * Removes the first (root, topmost) node 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.

 */

bh_node_type

binaryheap_remove_first(binaryheap *heap)

{

  bh_node_type result;

  Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);

  /* extract the root node, which will be the result */

  result = heap->bh_nodes[0];

  /* easy if heap contains one element */

  if (heap->bh_size == 1)

  {

    heap->bh_size--;

    return result;

  }

  /*

   * Remove the last node, placing it in the vacated root entry, and sift

   * the new root node down to its correct position.

   */

  heap->bh_nodes[0] = heap->bh_nodes[--heap->bh_size];

  sift_down(heap, 0);

  return result;

}

/*

 * binaryheap_remove_node

 *

 * Removes the nth (zero based) 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

 *

 * Replace the topmost element of a non-empty heap, preserving the heap

 * property.  O(1) in the best case, or O(log n) if it must fall back to

 * sifting the new node down.

 */

void

binaryheap_replace_first(binaryheap *heap, bh_node_type d)

{

  Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);

  heap->bh_nodes[0] = d;

  if (heap->bh_size > 1)

    sift_down(heap, 0);

}

/*

 * Sift a node up to the highest position it can hold according to the

 * comparator.

 */

static void

sift_up(binaryheap *heap, int node_off)

{

  bh_node_type node_val = heap->bh_nodes[node_off];

  /*

   * Within the loop, the node_off'th array entry is a "hole" that

   * notionally holds node_val, but we don't actually store node_val there

   * till the end, saving some unnecessary data copying steps.

   */

  while (node_off != 0)

  {

    int     cmp;

    int     parent_off;

    bh_node_type parent_val;

    /*

     * If this node is smaller than its parent, the heap condition is

     * satisfied, and we're done.

     */

    parent_off = parent_offset(node_off);

    parent_val = heap->bh_nodes[parent_off];

    cmp = heap->bh_compare(node_val,

                 parent_val,

                 heap->bh_arg);

    if (cmp <= 0)

      break;

    /*

     * Otherwise, swap the parent value with the hole, and go on to check

     * the node's new parent.

     */

    heap->bh_nodes[node_off] = parent_val;

    node_off = parent_off;

  }

  /* Re-fill the hole */

  heap->bh_nodes[node_off] = node_val;

}

/*

 * Sift a node down from its current position to satisfy the heap

 * property.

 */

static void

sift_down(binaryheap *heap, int node_off)

{

  bh_node_type node_val = heap->bh_nodes[node_off];

  /*

   * Within the loop, the node_off'th array entry is a "hole" that

   * notionally holds node_val, but we don't actually store node_val there

   * till the end, saving some unnecessary data copying steps.

   */

  while (true)

  {

    int     left_off = left_offset(node_off);

    int     right_off = right_offset(node_off);

    int     swap_off = left_off;

    /* Is the right child larger than the left child? */

    if (right_off < heap->bh_size &&

      heap->bh_compare(heap->bh_nodes[left_off],

               heap->bh_nodes[right_off],

               heap->bh_arg) < 0)

      swap_off = right_off;

    /*

     * If no children or parent is >= the larger child, heap condition is

     * satisfied, and we're done.

     */

    if (left_off >= heap->bh_size ||

      heap->bh_compare(node_val,

               heap->bh_nodes[swap_off],

               heap->bh_arg) >= 0)

      break;

    /*

     * Otherwise, swap the hole with the child that violates the heap

     * property; then go on to check its children.

     */

    heap->bh_nodes[node_off] = heap->bh_nodes[swap_off];

    node_off = swap_off;

  }

  /* Re-fill the hole */

  heap->bh_nodes[node_off] = node_val;

}