/*
 * 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/uds-releases/jasper/src/uds/util/radixSort.c#2 $
 */

/*
 * Radix sort is implemented using an American Flag sort, an unstable,
 * in-place 8-bit radix exchange sort.
 *
 * Adapted from the algorithm in the paper by Peter M. McIlroy, Keith Bostic,
 * and M. Douglas McIlroy, "Engineering Radix Sort".
 * http://www.usenix.org/publications/compsystems/1993/win_mcilroy.pdf
 */

#include "radixSort.h"

#include "compiler.h"
#include "memoryAlloc.h"
#include "stringUtils.h"
#include "typeDefs.h"
#include "uds.h"

enum {
  // Piles smaller than this are handled with a simple insertion sort.
  INSERTION_SORT_THRESHOLD  = 12
};

// Sort keys are pointers to immutable fixed-length arrays of bytes.
typedef const uint8_t * Key;

/**
 * The keys are separated into piles based on the byte in each
 * keys at the current offset, so the number of keys with each
 * byte must be counted.
 **/
typedef struct {
  uint16_t used;       // number of non-empty bins
  uint16_t first;      // index (key byte)  of the first non-empty bin
  uint16_t last;       // index (key byte) of the last non-empty bin
  uint32_t size[256];  // size[byte] == # of occurrences of byte
} Histogram;

/**
 * Sub-tasks are manually managed on a stack, both for performance
 * and to put a logarithmic bound on the stack space needed.
 **/
typedef struct {
  Key      *firstKey;  // Pointers to first and last keys to sort, inclusive.
  Key      *lastKey;
  uint16_t  offset;    // The offset into the key at which to continue sorting.
  uint16_t  length;    // The number of bytes remaining in the sort keys.
} Task;

struct radixSorter {
  unsigned int  count;
  Histogram     bins;
  Key          *pile[256];
  Task         *endOfStack;
  Task          isList[256];
  Task          stack[];
};

/**
 * Compare a segment of two fixed-length keys starting an offset.
 *
 * @param key1    the first key
 * @param key2    the second key
 * @param offset  the offset into the keys of the first byte to compare
 * @param length  the number of bytes remaining in each key
 **/
static INLINE int compare(Key key1, Key key2, uint16_t offset, uint16_t length)
{
  return memcmp(&key1[offset], &key2[offset], length);
}

/**
 * Insert the next unsorted key into an array of sorted keys.
 *
 * @param task  the description of the keys being sorted
 * @param next  the pointer to the unsorted key to insert into
 *              the array of sorted key pointers preceding it
 **/
static INLINE void insertKey(const Task task, Key *next)
{
  // Pull the unsorted key out, freeing up the array slot.
  Key unsorted = *next;
  // Compare the key to the preceding sorted entries, shifting
  // down the ones that are larger.
  while ((--next >= task.firstKey)
         && (compare(unsorted, next[0], task.offset, task.length) < 0)) {
    next[1] = next[0];
  }
  // Insert the key into the last slot that was cleared, sorting it.
  next[1] = unsorted;
}

/**
 * Sort a range of key segments using an insertion sort. This simple sort is
 * faster than the 256-way radix sort when the number of keys to sort is
 * small.
 *
 * @param task  the description of the keys to sort
 **/
static INLINE void insertionSort(const Task task)
{
  // (firstKey .. firstKey) is trivially sorted. Repeatedly insert the next
  // key into the sorted list of keys preceding it, and voila!
  Key *next;
  for (next = task.firstKey + 1; next <= task.lastKey; next++) {
    insertKey(task, next);
  }
}

/**
 * Push a sorting task onto the task stack, increasing the stack pointer.
 **/
static INLINE void pushTask(Task    **stackPointer,
                            Key      *firstKey,
                            uint32_t  count,
                            uint16_t  offset,
                            uint16_t  length)
{
  Task *task = (*stackPointer)++;
  task->firstKey = firstKey;
  task->lastKey  = &firstKey[count - 1];
  task->offset   = offset;
  task->length   = length;
}

/**********************************************************************/
static INLINE void swapKeys(Key *a, Key *b)
{
  Key c = *a;
  *a = *b;
  *b = c;
}

/**
 * Count the number of times each byte value appears in in the arrays of keys
 * to sort at the current offset, keeping track of the number of non-empty
 * bins, and the index of the first and last non-empty bin.
 *
 * @param task  the description of the keys to sort
 * @param bins  the histogram bins receiving the counts
 **/
static INLINE void measureBins(const Task task, Histogram *bins)
{
  // Set bogus values that will will be replaced by min and max, respectively.
  bins->first = UINT8_MAX;
  bins->last = 0;

  // Subtle invariant: bins->used and bins->size[] are zero because the
  // sorting code clears it all out as it goes. Even though this structure is
  // re-used, we don't need to pay to zero it before starting a new tally.

  Key *keyPtr;
  for (keyPtr = task.firstKey; keyPtr <= task.lastKey; keyPtr++) {
    // Increment the count for the byte in the key at the current offset.
    uint8_t bin = (*keyPtr)[task.offset];
    uint32_t size = ++bins->size[bin];

    // Track non-empty bins when the count transitions from zero to one.
    if (size == 1) {
      bins->used += 1;
      if (bin < bins->first) {
        bins->first = bin;
      }
      if (bin > bins->last) {
        bins->last = bin;
      }
    }
  }
}

/**
 * Convert the bin sizes to pointers to where each pile goes.
 *
 *   pile[0] = firstKey + bin->size[0],
 *   pile[1] = pile[0]  + bin->size[1], etc.
 *
 * After the keys are moved to the appropriate pile, we'll need to sort
 * each of the piles by the next radix position.  A new task is put on the
 * stack for each pile containing lots of keys, or a new task is is put on
 * the list for each pile containing few keys.
 *
 * @param stack      pointer the top of the stack
 * @param endOfStack the end of the stack
 * @param list       pointer the head of the list
 * @param pile       array that will be filled pointers to the end of each pile
 * @param bins       the histogram of the sizes of each pile
 * @param firstKey   the first key of the stack
 * @param offset     the next radix position to sort by
 * @param length     the number of bytes remaining in the sort keys
 *
 * @return UDS_SUCCESS or an error code
 **/
static INLINE int pushBins(Task       **stack,
                           Task       *endOfStack,
                           Task       **list,
                           Key        *pile[],
                           Histogram  *bins,
                           Key        *firstKey,
                           uint16_t    offset,
                           uint16_t    length)
{
  Key *pileStart = firstKey;
  int bin;
  for (bin = bins->first; ; bin++) {
    uint32_t size = bins->size[bin];
    // Skip empty piles.
    if (size == 0) {
      continue;
    }
    // There's no need to sort empty keys.
    if (length > 0) {
      if (size > INSERTION_SORT_THRESHOLD) {
        if (*stack >= endOfStack) {
          return UDS_BAD_STATE;
        }
        pushTask(stack, pileStart, size, offset, length);
      } else if (size > 1) {
        pushTask(list, pileStart, size, offset, length);
      }
    }
    pileStart += size;
    pile[bin] = pileStart;
    if (--bins->used == 0) {
      break;
    }
  }
  return UDS_SUCCESS;
}

/**********************************************************************/
int makeRadixSorter(unsigned int count, RadixSorter **sorter)
{
  unsigned int stackSize = count / INSERTION_SORT_THRESHOLD;
  RadixSorter *radixSorter;
  int result = ALLOCATE_EXTENDED(RadixSorter, stackSize, Task, __func__,
                                 &radixSorter);
  if (result != UDS_SUCCESS) {
    return result;
  }
  radixSorter->count = count;
  radixSorter->endOfStack = radixSorter->stack + stackSize;
  *sorter = radixSorter;
  return UDS_SUCCESS;
}

/**********************************************************************/
void freeRadixSorter(RadixSorter *sorter)
{
  FREE(sorter);
}

/**********************************************************************/
int radixSort(RadixSorter         *sorter,
              const unsigned char *keys[],
              unsigned int         count,
              unsigned short       length)
{
  // All zero-length keys are identical and therefore already sorted.
  if ((count == 0) || (length == 0)) {
    return UDS_SUCCESS;
  }

  // The initial task is to sort the entire length of all the keys.
  Task start = {
    .firstKey = keys,
    .lastKey  = &keys[count - 1],
    .offset   = 0,
    .length   = length,
  };

  if (count <= INSERTION_SORT_THRESHOLD) {
    insertionSort(start);
    return UDS_SUCCESS;
  }

  if (count > sorter->count) {
    return UDS_INVALID_ARGUMENT;
  }

  Histogram  *bins  = &sorter->bins;
  Key       **pile  = sorter->pile;
  Task       *sp    = sorter->stack;

  /*
   * Repeatedly consume a sorting task from the stack and process it, pushing
   * new sub-tasks onto to the stack for each radix-sorted pile. When all
   * tasks and sub-tasks have been processed, the stack will be empty and all
   * the keys in the starting task will be fully sorted.
   */
  for (*sp = start; sp >= sorter->stack; sp--) {
    const Task task = *sp;
    measureBins(task, bins);

    // Now that we know how large each bin is, generate pointers for each of
    // the piles and push a new task to sort each pile by the next radix byte.
    Task *lp = sorter->isList;
    int result = pushBins(&sp, sorter->endOfStack, &lp, pile, bins,
                          task.firstKey, task.offset + 1, task.length - 1);
    if (result != UDS_SUCCESS) {
      memset(bins, 0, sizeof(*bins));
      return result;
    }
    // Now bins->used is zero again.

    // Don't bother processing the last pile--when piles 0..N-1 are all in
    // place, then pile N must also be in place.
    Key *end = task.lastKey - bins->size[bins->last];
    bins->size[bins->last] = 0;

    Key *fence;
    for (fence = task.firstKey; fence <= end; ) {
      uint8_t bin;
      Key key = *fence;
      // The radix byte of the key tells us which pile it belongs in. Swap it
      // for an unprocessed item just below that pile, and repeat.
      while (--pile[bin = key[task.offset]] > fence) {
        swapKeys(pile[bin], &key);
      }
      // The pile reached the fence. Put the key at the bottom of that pile.
      // completing it, and advance the fence to the next pile.
      *fence = key;
      fence += bins->size[bin];
      bins->size[bin] = 0;
    }
    // Now bins->size[] is all zero again.

    // When the number of keys in a task gets small enough, its faster to use
    // an insertion sort than to keep subdividing into tiny piles.
    while (--lp >= sorter->isList) {
      insertionSort(*lp);
    }
  }
  return UDS_SUCCESS;
}