/* For qsort_r */
#define _GNU_SOURCE

#include <assert.h>
#include <limits.h>
#include <stdlib.h>
#include <string.h>

#include "policy-internal.h"

#include "quorum-internal.h"
#include "xalloc.h"

static size_t
bytecode_size_with_prefix (size_t w)
{
  size_t i = 1;
  for (w >>= 6; w > 0; i++)
    w >>= 6;
  return i;
}

static void
bytecode_write_prefix (uint8_t *bytecode, size_t size, size_t prefix)
{
  if (size == 0)
    return;
  /* Rule out undefined shift count */
  assert (6 * size < CHAR_BIT * sizeof (prefix));
  assert ((prefix >> (6*size)) == 0);
  /* Prefix bits of first byte must be nonzero. */
  assert ((prefix >> (6*(size - 1))) > 0);
  while (size-- > 0)
    /* Since the top-most bits are set to 0b11, no need to mask the
       prefix bits. */
    *bytecode++ = (BYTECODE_OP_PREFIX << 6) | (prefix >> (size * 6));
}

static void
bytecode_write_witness (uint8_t *bytecode, size_t size, size_t witness)
{
  assert (size > 0);
  bytecode_write_prefix (bytecode, size - 1, witness >> 6);
  bytecode[size-1] = (BYTECODE_OP_X_P << 6) | (witness & 0x3f);
}

static void
bytecode_write_threshold (uint8_t *bytecode, size_t size, size_t threshold)
{
  assert (size > 0);
  bytecode_write_prefix (bytecode, size - 1, threshold >> 6);
  bytecode[size-1] = (BYTECODE_OP_GE_K << 6) | (threshold & 0x3f);
}

static int
member_cmp (const void *ap, const void *bp, void *ctx)
{
  const struct sigsum_policy_def *defs = ctx;
  size_t a_size = defs[*(const size_t *) ap].bytecode_size;
  size_t b_size = defs[*(const size_t *) bp].bytecode_size;

  if (a_size < b_size)
    return -1;
  return (a_size > b_size);
}

/* Set bytecode size, and sorts members by size, if the definition is
   a group. */
static void
bytecode_size (const struct sigsum_policy_def *defs, struct sigsum_policy_def *def)
{
  assert (!def->bytecode_size);
  switch (def->type)
    {
    case SIGSUM_POLICY_DEF_NONE:
      def->bytecode_size = 1;
      return;
    case SIGSUM_POLICY_DEF_WITNESS:
      def->bytecode_size = bytecode_size_with_prefix (def->witness_index);
      return;
    case SIGSUM_POLICY_DEF_GROUP: ;
      struct sigsum_policy_group *group = &def->group;
      assert (group->num_members > 0);
      if (group->num_members == 1)
	{
	  assert (group->threshold == 1);
	  def->bytecode_size = defs[group->members[0]].bytecode_size;
	  return;
	}

      /* Sort members by bytecode size. */
      qsort_r (group->members, group->num_members, sizeof (group->members[0]),
	       member_cmp, (void *) defs);

      def->bytecode_size = group->num_members - 1
	+ bytecode_size_with_prefix (group->threshold);

      for (unsigned i = 0; i < group->num_members; i++)
	{
	  unsigned member_bytecode_size = defs[group->members[i]].bytecode_size;
	  assert (member_bytecode_size > 0);
	  def->bytecode_size += member_bytecode_size;
	}
      break;
    }
}

static int
bytecode_cmp (const void *ap, const void *bp, void *ctx)
{
  size_t size = *(const unsigned *) ctx;
  return memcmp (ap, bp, size);
}

/* Sort the entries of same size, and append add instructions to each
   group but the very first. Returns total size including add
   instructions. */
static size_t
bytecode_sort (uint8_t *bytecode, bool is_first, size_t n, size_t size)
{
  assert (n > 0);

  if (n > 1)
    qsort_r (bytecode, n, size, bytecode_cmp, &size);

  if (is_first)
    {
      /* No add instruction after the very first member. */
      if (!--n)
	return size;
      bytecode += size;
    }

  /* Move to make space for add instruction after each member.
     Element i (i > 0) is moved from i*size to i*(size+1). */
  for (unsigned i = n-1; i > 0; i--)
    {
      memmove (bytecode + i * (size + 1),
	       bytecode + i * size, size);
      /* Append add instruction. */
      bytecode[i*(size + 1) + size] = BYTECODE_INSTR_ADD;
    }
  bytecode[size] = BYTECODE_INSTR_ADD;
  return n*(size+1) + (is_first ? size : 0);
}

static void
bytecode_compile (uint8_t *bytecode, const struct sigsum_policy_def *defs,
		  size_t index)
{
  const struct sigsum_policy_def *def = &defs[index];
  switch (def->type)
    {
    case SIGSUM_POLICY_DEF_NONE:
      assert (def->bytecode_size == 1);
      *bytecode = BYTECODE_INSTR_TRUE;
      return;
    case SIGSUM_POLICY_DEF_WITNESS:
      bytecode_write_witness (bytecode, def->bytecode_size, def->witness_index);
      return;
    case SIGSUM_POLICY_DEF_GROUP: ;
      const struct sigsum_policy_group *group = &def->group;
      assert (group->num_members > 0);

      if (group->num_members == 1)
	{
	  assert (group->threshold == 1);
	  assert (def->bytecode_size == defs[group->members[0]].bytecode_size);
	  bytecode_compile (bytecode, defs, group->members[0]);
	  return;
	}
      size_t left = def->bytecode_size;

      const size_t *members = group->members;
      for (size_t done = 0; done < group->num_members; )
	{
	  /* Find interval of groups of the same compiled size */
	  size_t size = defs[members[done]].bytecode_size;
	  bytecode_compile (bytecode, defs, members[done]);

	  size_t n;
	  for (n = 1; done + n < group->num_members ; n++)
	    {
	      size_t next_size = defs[members[done + n]].bytecode_size;
	      if (next_size == size)
		bytecode_compile (bytecode + n * size, defs, members[done + n]);
	      else
		{
		  assert (next_size > size);
		  break;
		}
	    }

	  /* Sort bytecodes of these n members lexicographically, and
	     interleave add instructions. */
	  size_t written = bytecode_sort (bytecode, done == 0, n, size);
	  assert (written < left);

	  bytecode += written;
	  left -= written;

	  done += n;
	}
      bytecode_write_threshold (bytecode, left, group->threshold);
    }
}

/* Computes needed stack size. Returns (size_t) -1 if byte code is invalid. */
static size_t
bytecode_stack_size (size_t n, const uint8_t *bytecode)
{
  size_t stack_size = 0;
  size_t max_size = 0;
  for (size_t i = 0; i < n; i++)
    {
      uint8_t instr = bytecode[i];
      switch (quorum_bytecode_op (instr))
	{
	case BYTECODE_OP_MISC:
	  switch (instr)
	    {
	    case BYTECODE_INSTR_ADD:
	      if (stack_size < 2)
		return (size_t) -1;
	      stack_size--;
	      break;
	    case BYTECODE_INSTR_TRUE:
	      stack_size++;
	      break;
	    default:
	      return (size_t) -1;
	    }
	  break;
	case BYTECODE_OP_X_P:
	  stack_size++;
	  break;
	case BYTECODE_OP_GE_K:
	  if (stack_size < 1)
	    return (size_t) -1;
	  break;
	case BYTECODE_OP_PREFIX:
	  break;
	}
      if (stack_size > max_size)
	max_size = stack_size;
    }
  return max_size;
}

void
_sigsum_bytecode_compile_quorum (struct sigsum_quorum *quorum,
				 struct sigsum_policy_def *defs,
				 size_t index)
{
  /* Compute the compiled size of each group or witness, */
  for (size_t i = 0; i <= index; i++)
    bytecode_size (defs, &defs[i]);

  quorum->bytecode_size = defs[index].bytecode_size;
  uint8_t *bytecode = _sigsum_xalloc (quorum->bytecode_size);
  bytecode_compile (bytecode, defs, index);

  quorum->stack_alloc = bytecode_stack_size (quorum->bytecode_size, bytecode);
  assert (quorum->stack_alloc != (size_t) -1);

  quorum->bytecode = bytecode;
}
