/**

 * vim: set ts=2 sw=2 tw=99 et:

 * =============================================================================

 * SourceMod

 * Copyright (C) 2004-2008 AlliedModders LLC.  All rights reserved.

 * =============================================================================

 *

 * This program is free software; you can redistribute it and/or modify it under

 * the terms of the GNU General Public License, version 3.0, as published by the

 * Free Software Foundation.

 * 

 * 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, see <http://www.gnu.org/licenses/>.

 *

 * As a special exception, AlliedModders LLC gives you permission to link the

 * code of this program (as well as its derivative works) to "Half-Life 2," the

 * "Source Engine," the "SourcePawn JIT," and any Game MODs that run on software

 * by the Valve Corporation.  You must obey the GNU General Public License in

 * all respects for all other code used.  Additionally, AlliedModders LLC grants

 * this exception to all derivative works.  AlliedModders LLC defines further

 * exceptions, found in LICENSE.txt (as of this writing, version JULY-31-2007),

 * or <http://www.sourcemod.net/license.php>.

 *

 * Version: $Id$

 */

#include <string.h>

#include <stdlib.h>

#include <assert.h>

#include "jit_x86.h"

#include "../sp_vm_engine.h"

#include "../engine2.h"

#include "../BaseRuntime.h"

#include "../sp_vm_basecontext.h"

#include "watchdog_timer.h"

using namespace Knight;

#if defined USE_UNGEN_OPCODES

#include "ungen_opcodes.h"

#endif

#define __ masm.

JITX86 g_Jit;

KeCodeCache *g_pCodeCache = NULL;

ISourcePawnEngine *engine = &g_engine1;

static inline uint8_t *

LinkCode(AssemblerX86 &masm)

{

  if (masm.outOfMemory())

    return NULL;

  void *code = Knight::KE_AllocCode(g_pCodeCache, masm.length());

  if (!code)

    return NULL;

  masm.emitToExecutableMemory(code);

  return reinterpret_cast<uint8_t *>(code);

}

static inline ConditionCode

OpToCondition(OPCODE op)

{

  switch (op) {

   case OP_EQ:

   case OP_JEQ:

    return equal;

   case OP_NEQ:

   case OP_JNEQ:

    return not_equal;

   case OP_SLESS:

   case OP_JSLESS:

    return less;

   case OP_SLEQ:

   case OP_JSLEQ:

    return less_equal;

   case OP_SGRTR:

   case OP_JSGRTR:

    return greater;

   case OP_SGEQ:

   case OP_JSGEQ:

    return greater_equal;

   default:

    assert(false);

    return negative;

  }

}

struct array_creation_t

{

  const cell_t *dim_list;     /* Dimension sizes */

  cell_t dim_count;           /* Number of dimensions */

  cell_t *data_offs;          /* Current offset AFTER the indirection vectors (data) */

  cell_t *base;               /* array base */

};

static cell_t

GenerateInnerArrayIndirectionVectors(array_creation_t *ar, int dim, cell_t cur_offs)

{

  cell_t write_offs = cur_offs;

  cell_t *data_offs = ar->data_offs;

  cur_offs += ar->dim_list[dim];

  // Dimension n-x where x > 2 will have sub-vectors.  

  // Otherwise, we just need to reference the data section.

  if (ar->dim_count > 2 && dim < ar->dim_count - 2) {

    // For each index at this dimension, write offstes to our sub-vectors.

    // After we write one sub-vector, we generate its sub-vectors recursively.

    // At the end, we're given the next offset we can use.

    for (int i = 0; i < ar->dim_list[dim]; i++) {

      ar->base[write_offs] = (cur_offs - write_offs) * sizeof(cell_t);

      write_offs++;

      cur_offs = GenerateInnerArrayIndirectionVectors(ar, dim + 1, cur_offs);

    }

  } else {

    // In this section, there are no sub-vectors, we need to write offsets 

    // to the data.  This is separate so the data stays in one big chunk.

    // The data offset will increment by the size of the last dimension, 

    // because that is where the data is finally computed as. 

    for (int i = 0; i < ar->dim_list[dim]; i++) {

      ar->base[write_offs] = (*data_offs - write_offs) * sizeof(cell_t);

      write_offs++;

      *data_offs = *data_offs + ar->dim_list[dim + 1];

    }

  }

  return cur_offs;

}

static cell_t

calc_indirection(const array_creation_t *ar, cell_t dim)

{

  cell_t size = ar->dim_list[dim];

  if (dim < ar->dim_count - 2)

    size += ar->dim_list[dim] * calc_indirection(ar, dim + 1);

  return size;

}

static cell_t

GenerateArrayIndirectionVectors(cell_t *arraybase, cell_t dims[], cell_t _dimcount, bool autozero)

{

  array_creation_t ar;

  cell_t data_offs;

  /* Reverse the dimensions */

  cell_t dim_list[sDIMEN_MAX];

  int cur_dim = 0;

  for (int i = _dimcount - 1; i >= 0; i--)

    dim_list[cur_dim++] = dims[i];

  ar.base = arraybase;

  ar.dim_list = dim_list;

  ar.dim_count = _dimcount;

  ar.data_offs = &data_offs;

  data_offs = calc_indirection(&ar, 0);

  GenerateInnerArrayIndirectionVectors(&ar, 0, 0);

  return data_offs;

}

static int

PopTrackerAndSetHeap(BaseRuntime *rt)

{

  sp_context_t *ctx = rt->GetBaseContext()->GetCtx();

  tracker_t *trk = (tracker_t *)(ctx->vm[JITVARS_TRACKER]);

  assert(trk->pCur > trk->pBase);

  trk->pCur--;

  if (trk->pCur < trk->pBase)

    return SP_ERROR_TRACKER_BOUNDS;

  ucell_t amt = *trk->pCur;

  if (amt > (ctx->hp - rt->plugin()->data_size))

    return SP_ERROR_HEAPMIN;

  ctx->hp -= amt;

  return SP_ERROR_NONE;

}

static int

PushTracker(sp_context_t *ctx, size_t amount)

{

  tracker_t *trk = (tracker_t *)(ctx->vm[JITVARS_TRACKER]);

  if ((size_t)(trk->pCur - trk->pBase) >= trk->size)

    return SP_ERROR_TRACKER_BOUNDS;

  if (trk->pCur + 1 - (trk->pBase + trk->size) == 0) {

    size_t disp = trk->size - 1;

    trk->size *= 2;

    trk->pBase = (ucell_t *)realloc(trk->pBase, trk->size * sizeof(cell_t));

    if (!trk->pBase)

      return SP_ERROR_TRACKER_BOUNDS;

    trk->pCur = trk->pBase + disp;

  }

  *trk->pCur++ = amount;

  return SP_ERROR_NONE;

}

static int

GenerateArray(BaseRuntime *rt, uint32_t argc, cell_t *argv, int autozero)

{

  sp_context_t *ctx = rt->GetBaseContext()->GetCtx();

  // Calculate how many cells are needed.

  if (argv[0] <= 0)

    return SP_ERROR_ARRAY_TOO_BIG;

  uint32_t dim = 1;         // second to last dimension

  uint32_t cells = argv[0];

  for (uint32_t dim = 1; dim < argc; dim++) {

    cell_t dimsize = argv[dim];

    if (dimsize <= 0)

      return SP_ERROR_ARRAY_TOO_BIG;

    if (!ke::IsUint32MultiplySafe(cells, dimsize))

      return SP_ERROR_ARRAY_TOO_BIG;

    cells *= uint32_t(dimsize);

    if (!ke::IsUint32AddSafe(cells, dimsize))

      return SP_ERROR_ARRAY_TOO_BIG;

    cells += uint32_t(dimsize);

  }

  if (!ke::IsUint32MultiplySafe(cells, 4))

    return SP_ERROR_ARRAY_TOO_BIG;

  uint32_t bytes = cells * 4;

  if (!ke::IsUint32AddSafe(ctx->hp, bytes))

    return SP_ERROR_ARRAY_TOO_BIG;

  uint32_t new_hp = ctx->hp + bytes;

  cell_t *dat_hp = reinterpret_cast<cell_t *>(rt->plugin()->memory + new_hp);

  // argv, coincidentally, is STK.

  if (dat_hp >= argv - STACK_MARGIN)

    return SP_ERROR_HEAPLOW;

  if (int err = PushTracker(rt->GetBaseContext()->GetCtx(), bytes))

    return err;

  cell_t *base = reinterpret_cast<cell_t *>(rt->plugin()->memory + ctx->hp);

  cell_t offs = GenerateArrayIndirectionVectors(base, argv, argc, !!autozero);

  assert(size_t(offs) == cells);

  argv[argc - 1] = ctx->hp;

  ctx->hp = new_hp;

  return SP_ERROR_NONE;

}

#if !defined NDEBUG

static const char *

GetFunctionName(const sp_plugin_t *plugin, uint32_t offs)

{

  if (!plugin->debug.unpacked) {

    uint32_t max, iter;

    sp_fdbg_symbol_t *sym;

    sp_fdbg_arraydim_t *arr;

    uint8_t *cursor = (uint8_t *)(plugin->debug.symbols);

    max = plugin->debug.syms_num;

    for (iter = 0; iter < max; iter++) {

      sym = (sp_fdbg_symbol_t *)cursor;

      if (sym->ident == SP_SYM_FUNCTION && sym->codestart <= offs && sym->codeend > offs)

        return plugin->debug.stringbase + sym->name;

      if (sym->dimcount > 0) {

        cursor += sizeof(sp_fdbg_symbol_t);

        arr = (sp_fdbg_arraydim_t *)cursor;

        cursor += sizeof(sp_fdbg_arraydim_t) * sym->dimcount;

        continue;

      }

      cursor += sizeof(sp_fdbg_symbol_t);

    }

  } else {

    uint32_t max, iter;

    sp_u_fdbg_symbol_t *sym;

    sp_u_fdbg_arraydim_t *arr;

    uint8_t *cursor = (uint8_t *)(plugin->debug.symbols);

    max = plugin->debug.syms_num;

    for (iter = 0; iter < max; iter++) {

      sym = (sp_u_fdbg_symbol_t *)cursor;

      if (sym->ident == SP_SYM_FUNCTION && sym->codestart <= offs && sym->codeend > offs)

        return plugin->debug.stringbase + sym->name;

      if (sym->dimcount > 0) {

        cursor += sizeof(sp_u_fdbg_symbol_t);

        arr = (sp_u_fdbg_arraydim_t *)cursor;

        cursor += sizeof(sp_u_fdbg_arraydim_t) * sym->dimcount;

        continue;

      }

      cursor += sizeof(sp_u_fdbg_symbol_t);

    }

  }

  return NULL;

}

#endif

static int

CompileFromThunk(BaseRuntime *runtime, cell_t pcode_offs, void **addrp, char *pc)

{

  // If the watchdog timer has declared a timeout, we must process it now,

  // and possibly refuse to compile, since otherwise we will compile a

  // function that is not patched for timeouts.

  if (!g_WatchdogTimer.HandleInterrupt())

    return SP_ERROR_TIMEOUT;

  JitFunction *fn = runtime->GetJittedFunctionByOffset(pcode_offs);

  if (!fn) {

    int err;

    fn = g_Jit.CompileFunction(runtime, pcode_offs, &err);

    if (!fn)

      return err;

  }

#if !defined NDEBUG

  g_engine1.GetDebugHook()->OnDebugSpew(

      "Patching thunk to %s::%s\n",

      runtime->plugin()->name,

      GetFunctionName(runtime->plugin(), pcode_offs));

#endif

  *addrp = fn->GetEntryAddress();

  /* Right now, we always keep the code RWE */

  *(intptr_t *)(pc - 4) = intptr_t(fn->GetEntryAddress()) - intptr_t(pc);

  return SP_ERROR_NONE;

}

static cell_t

NativeCallback(sp_context_t *ctx, ucell_t native_idx, cell_t *params)

{

  sp_native_t *native;

  cell_t save_sp = ctx->sp;

  cell_t save_hp = ctx->hp;

  sp_plugin_t *pl = (sp_plugin_t *)(ctx->vm[JITVARS_PLUGIN]);

  ctx->n_idx = native_idx;

  if (ctx->hp < (cell_t)pl->data_size) {

    ctx->n_err = SP_ERROR_HEAPMIN;

    return 0;

  }

  if (ctx->hp + STACK_MARGIN > ctx->sp) {

    ctx->n_err = SP_ERROR_STACKLOW;

    return 0;

  }

  if ((uint32_t)ctx->sp >= pl->mem_size) {

    ctx->n_err = SP_ERROR_STACKMIN;

    return 0;

  }

  native = &pl->natives[native_idx];

  if (native->status == SP_NATIVE_UNBOUND) {

    ctx->n_err = SP_ERROR_INVALID_NATIVE;

    return 0;

  }

  cell_t result = native->pfn(GET_CONTEXT(ctx), params);

  if (ctx->n_err != SP_ERROR_NONE)

    return result;

  if (save_sp != ctx->sp) {

    ctx->n_err = SP_ERROR_STACKLEAK;

    return result;

  }

  if (save_hp != ctx->hp) {

    ctx->n_err = SP_ERROR_HEAPLEAK;

    return result;

  }

  return result;

}

Compiler::Compiler(BaseRuntime *rt, cell_t pcode_offs)

  : rt_(rt),

    plugin_(rt->plugin()),

    error_(SP_ERROR_NONE),

    pcode_start_(pcode_offs),

    code_start_(reinterpret_cast<cell_t *>(plugin_->pcode + pcode_start_)),

    cip_(code_start_),

    code_end_(reinterpret_cast<cell_t *>(plugin_->pcode + plugin_->pcode_size))

{

  size_t nmaxops = plugin_->pcode_size / sizeof(cell_t) + 1;

  jump_map_ = new Label[nmaxops];

}

Compiler::~Compiler()

{

  delete [] jump_map_;

}

JitFunction *

Compiler::emit(int *errp)

{

  if (cip_ >= code_end_ || *cip_ != OP_PROC) {

    *errp = SP_ERROR_INVALID_INSTRUCTION;

    return NULL;

  }

#if !defined NDEBUG

  g_engine1.GetDebugHook()->OnDebugSpew(

      "Compiling function %s::%s\n",

      plugin_->name,

      GetFunctionName(plugin_, pcode_start_));

  SpewOpcode(plugin_, code_start_, cip_);

#endif

  cell_t *codeseg = reinterpret_cast<cell_t *>(plugin_->pcode);

  cip_++;

  if (!emitOp(OP_PROC)) {

      *errp = (error_ == SP_ERROR_NONE) ? SP_ERROR_OUT_OF_MEMORY : error_;

      return NULL;

  }

  while (cip_ < code_end_) {

    // If we reach the end of this function, or the beginning of a new

    // procedure, then stop.

    if (*cip_ == OP_PROC || *cip_ == OP_ENDPROC)

      break;

#if !defined NDEBUG

    SpewOpcode(plugin_, code_start_, cip_);

#endif

    // We assume every instruction is a jump target, so before emitting

    // an opcode, we bind its corresponding label.

    __ bind(&jump_map_[cip_ - codeseg]);

    OPCODE op = (OPCODE)readCell();

    if (!emitOp(op) || error_ != SP_ERROR_NONE) {

      *errp = (error_ == SP_ERROR_NONE) ? SP_ERROR_OUT_OF_MEMORY : error_;

      return NULL;

    }

  }

  emitCallThunks();

  emitErrorPaths();

  uint8_t *code = LinkCode(masm);

  if (!code) {

    *errp = SP_ERROR_OUT_OF_MEMORY;

    return NULL;

  }

  LoopEdge *edges = new LoopEdge[backward_jumps_.length()];

  for (size_t i = 0; i < backward_jumps_.length(); i++) {

    edges[i].offset = backward_jumps_[i];

    edges[i].disp32 = *reinterpret_cast<int32_t *>(code + edges[i].offset - 4);

  }

  return new JitFunction(code, pcode_start_, edges, backward_jumps_.length());

}

bool

Compiler::emitOp(OPCODE op)

{

  switch (op) {

    case OP_MOVE_PRI:

      __ movl(pri, alt);

      break;

    case OP_MOVE_ALT:

      __ movl(alt, pri);

      break;

    case OP_XCHG:

      __ xchgl(pri, alt);

      break;

    case OP_ZERO:

    {

      cell_t offset = readCell();

      __ movl(Operand(dat, offset), 0);

      break;

    }

    case OP_ZERO_S:

    {

      cell_t offset = readCell();

      __ movl(Operand(frm, offset), 0);

      break;

    }

    case OP_PUSH_PRI:

    case OP_PUSH_ALT:

    {

      Register reg = (op == OP_PUSH_PRI) ? pri : alt;

      __ movl(Operand(stk, -4), reg);

      __ subl(stk, 4);

      break;

    }

    case OP_PUSH_C:

    case OP_PUSH2_C:

    case OP_PUSH3_C:

    case OP_PUSH4_C:

    case OP_PUSH5_C:

    {

      int n = 1;

      if (op >= OP_PUSH2_C)

        n = ((op - OP_PUSH2_C) / 4) + 2;

      int i = 1;

      do {

        cell_t val = readCell();

        __ movl(Operand(stk, -(4 * i)), val);

      } while (i++ < n);

      __ subl(stk, 4 * n);

      break;

    }

    case OP_PUSH_ADR:

    case OP_PUSH2_ADR:

    case OP_PUSH3_ADR:

    case OP_PUSH4_ADR:

    case OP_PUSH5_ADR:

    {

      int n = 1;

      if (op >= OP_PUSH2_ADR)

        n = ((op - OP_PUSH2_ADR) / 4) + 2;

      int i = 1;

      // We temporarily relocate FRM to be a local address instead of an

      // absolute address.