/*
 * Copyright (c) 2020 Red Hat, Inc.
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301, USA. 
 *
 * $Id: //eng/vdo-releases/aluminum/src/c++/vdo/base/heap.c#2 $
 */

#include "heap.h"

#include "errors.h"
#include "logger.h"
#include "numeric.h"

#include "statusCodes.h"

/**********************************************************************/
void initializeHeap(Heap           *heap,
                    HeapComparator *comparator,
                    HeapSwapper    *swapper,
                    void           *array,
                    size_t          capacity,
                    size_t          elementSize)
{
  *heap = (Heap) {
    .comparator  = comparator,
    .swapper     = swapper,
    .capacity    = capacity,
    .elementSize = elementSize,
  };
  if (array != NULL) {
    // Calculating child indexes is simplified by pretending the element array
    // is 1-based.
    heap->array = ((byte *) array - elementSize);
  }
}

/**********************************************************************/
static void siftHeapDown(Heap *heap, size_t topNode, size_t lastNode)
{
  // Keep sifting until the sub-heap rooted at topNode has no children.
  size_t leftChild;
  while ((leftChild = (2 * topNode)) <= lastNode) {
    // If there are two children, select the largest child to swap with.
    size_t swapNode = leftChild;
    if (leftChild < lastNode) {
      size_t rightChild = leftChild + heap->elementSize;
      if (heap->comparator(&heap->array[leftChild],
                           &heap->array[rightChild]) < 0) {
        swapNode = rightChild;
      }
    }

    // Stop sifting if topNode is at least as large as its largest child,
    // which means the heap invariant was restored by the previous swap.
    if (heap->comparator(&heap->array[topNode], &heap->array[swapNode]) >= 0) {
      return;
    }

    // Swap the element we've been sifting down with the larger child.
    heap->swapper(&heap->array[topNode], &heap->array[swapNode]);

    // Descend into the sub-heap rooted at that child, going around the loop
    // again in place of a tail-recursive call to siftHeapDown().
    topNode = swapNode;
  }

  // We sifted the element all the way to a leaf node of the heap, so the heap
  // invariant has now been restored.
}

/**********************************************************************/
void buildHeap(Heap *heap, size_t count)
{
  heap->count = minSizeT(count, heap->capacity);

  if ((heap->count < 2) || (heap->elementSize == 0)) {
    return;
  }

  /*
   * All the leaf nodes are trivially valid sub-heaps. Starting with the parent
   * of the right-most leaf node, restore the heap invariant in that sub-heap
   * by sifting the top node of the sub-heap down into one of its children's
   * valid sub-heaps (or not, if the top node is already larger than its
   * children). Continue iterating through all the interior nodes in the heap,
   * in sort of a reverse breadth-first traversal, restoring the heap
   * invariant for each (increasingly larger) sub-heap until we reach the root
   * of the heap. Once we sift the root node down into one of its two valid
   * children, the entire heap must be valid, by induction.
   *
   * Even though we operate on every node and potentially perform an O(log N)
   * traversal for each node, the combined probabilities of actually needing
   * to do a swap and the heights of the sub-heaps sum to a constant, so
   * restoring a heap from the bottom-up like this has only O(N) complexity.
   */
  size_t size       = heap->elementSize;
  size_t lastParent = size * (heap->count / 2);
  size_t lastNode   = size * heap->count;
  for (size_t topNode = lastParent; topNode > 0; topNode -= size) {
    siftHeapDown(heap, topNode, lastNode);
  }
}

/**********************************************************************/
bool popMaxHeapElement(Heap *heap, void *elementPtr)
{
  if (heap->count == 0) {
    return false;
  }

  size_t rootNode = (heap->elementSize * 1);
  size_t lastNode = (heap->elementSize * heap->count);

  // Return the maximum element (the root of the heap) if the caller wanted it.
  if (elementPtr != NULL) {
    memcpy(elementPtr, &heap->array[rootNode], heap->elementSize);
  }

  // Move the right-most leaf node to the vacated root node, reducing the
  // number of elements by one and violating the heap invariant.
  if (rootNode != lastNode) {
    memcpy(&heap->array[rootNode], &heap->array[lastNode], heap->elementSize);
  }
  heap->count -= 1;
  lastNode    -= heap->elementSize;

  // Restore the heap invariant by sifting the root back down into the heap.
  siftHeapDown(heap, rootNode, lastNode);
  return true;
}

/**********************************************************************/
static inline size_t siftAndSort(Heap *heap, size_t rootNode, size_t lastNode)
{
  /*
   * We have a valid heap, so the largest unsorted element is now at the top
   * of the heap. That element belongs at the start of the partially-sorted
   * array, preceding all the larger elements that we've already removed
   * from the heap. Swap that largest unsorted element with the the
   * right-most leaf node in the heap, moving it to its sorted position in
   * the array.
   */
  heap->swapper(&heap->array[rootNode], &heap->array[lastNode]);
  // The sorted list is now one element larger and valid. The heap is
  // one element smaller, and invalid.
  lastNode -= heap->elementSize;
  // Restore the heap invariant by sifting the swapped element back down
  // into the heap.
  siftHeapDown(heap, rootNode, lastNode);
  return lastNode;
}

/**********************************************************************/
size_t sortHeap(Heap *heap)
{
  // All zero-length records are identical and therefore already sorted, as
  // are empty or singleton arrays.
  if ((heap->count < 2) || (heap->elementSize == 0)) {
    return heap->count;
  }

  // Get the byte array offset of the root node, and the right-most leaf node
  // in the 1-based array of records that will form the heap.
  size_t rootNode = (heap->elementSize * 1);
  size_t lastNode = (heap->elementSize * heap->count);

  while (lastNode > rootNode) {
    lastNode = siftAndSort(heap, rootNode, lastNode);
  }

  size_t count = heap->count;
  heap->count = 0;
  return count;
}

/**********************************************************************/
void *sortNextHeapElement(Heap *heap)
{
  if ((heap->count == 0) || (heap->elementSize == 0)) {
    return NULL;
  }

  // Get the byte array offset of the root node, and the right-most leaf node
  // in the 1-based array of records that will form the heap.
  size_t rootNode = (heap->elementSize * 1);
  size_t lastNode = (heap->elementSize * heap->count);
  if (heap->count > 1) {
    siftAndSort(heap, rootNode, lastNode);
  }
  heap->count--;

  return &heap->array[lastNode];
}