/*

 * IR - Lightweight JIT Compilation Framework

 * (Native code patcher)

 * Copyright (C) 2022 Zend by Perforce.

 * Authors: Dmitry Stogov <dmitry@php.net>

 *

 * Based on Mike Pall's implementation for LuaJIT.

 */

#include "ir.h"

#include "ir_private.h"

#if defined(IR_TARGET_X86) || defined(IR_TARGET_X64)

static uint32_t _asm_x86_inslen(const uint8_t* p)

{

  static const uint8_t map_op1[256] = {

    0x92,0x92,0x92,0x92,0x52,0x45,0x51,0x51,0x92,0x92,0x92,0x92,0x52,0x45,0x51,0x20,

    0x92,0x92,0x92,0x92,0x52,0x45,0x51,0x51,0x92,0x92,0x92,0x92,0x52,0x45,0x51,0x51,

    0x92,0x92,0x92,0x92,0x52,0x45,0x10,0x51,0x92,0x92,0x92,0x92,0x52,0x45,0x10,0x51,

    0x92,0x92,0x92,0x92,0x52,0x45,0x10,0x51,0x92,0x92,0x92,0x92,0x52,0x45,0x10,0x51,

#ifdef IR_TARGET_X64

    0x10,0x10,0x10,0x10,0x10,0x10,0x10,0x10,0x14,0x14,0x14,0x14,0x14,0x14,0x14,0x14,

#else

    0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,

#endif

    0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,

    0x51,0x51,0x92,0x92,0x10,0x10,0x12,0x11,0x45,0x86,0x52,0x93,0x51,0x51,0x51,0x51,

    0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,

    0x93,0x86,0x93,0x93,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,

    0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x47,0x51,0x51,0x51,0x51,0x51,

#ifdef IR_TARGET_X64

    0x59,0x59,0x59,0x59,0x51,0x51,0x51,0x51,0x52,0x45,0x51,0x51,0x51,0x51,0x51,0x51,

#else

    0x55,0x55,0x55,0x55,0x51,0x51,0x51,0x51,0x52,0x45,0x51,0x51,0x51,0x51,0x51,0x51,

#endif

    0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x05,0x05,0x05,0x05,0x05,0x05,0x05,0x05,

    0x93,0x93,0x53,0x51,0x70,0x71,0x93,0x86,0x54,0x51,0x53,0x51,0x51,0x52,0x51,0x51,

    0x92,0x92,0x92,0x92,0x52,0x52,0x51,0x51,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,

    0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x45,0x45,0x47,0x52,0x51,0x51,0x51,0x51,

    0x10,0x51,0x10,0x10,0x51,0x51,0x63,0x66,0x51,0x51,0x51,0x51,0x51,0x51,0x92,0x92

  };

  static const uint8_t map_op2[256] = {

    0x93,0x93,0x93,0x93,0x52,0x52,0x52,0x52,0x52,0x52,0x51,0x52,0x51,0x93,0x52,0x94,

    0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,

    0x53,0x53,0x53,0x53,0x53,0x53,0x53,0x53,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,

    0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x34,0x51,0x35,0x51,0x51,0x51,0x51,0x51,

    0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,

    0x53,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,

    0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,

    0x94,0x54,0x54,0x54,0x93,0x93,0x93,0x52,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,

    0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,

    0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,

    0x52,0x52,0x52,0x93,0x94,0x93,0x51,0x51,0x52,0x52,0x52,0x93,0x94,0x93,0x93,0x93,

    0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x94,0x93,0x93,0x93,0x93,0x93,

    0x93,0x93,0x94,0x93,0x94,0x94,0x94,0x93,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,

    0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,

    0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,

    0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x52

  };

  uint32_t result = 0;

  uint32_t prefixes = 0;

  uint32_t x = map_op1[*p];

  for (;;) {

    switch (x >> 4) {

      case 0:

        return result + x + (prefixes & 4);

      case 1:

        prefixes |= x;

        x = map_op1[*++p];

        result++;

        break;

      case 2:

        x = map_op2[*++p];

        break;

      case 3:

        p++;

        goto mrm;

      case 4:

        result -= (prefixes & 2);

        /* fallthrough */

      case 5:

        return result + (x & 15);

      case 6: /* Group 3. */

        if (p[1] & 0x38) {

          x = 2;

        } else if ((prefixes & 2) && (x == 0x66)) {

          x = 4;

        }

        goto mrm;

      case 7: /* VEX c4/c5. */

#ifdef IR_TARGET_X86

        if (p[1] < 0xc0) {

          x = 2;

          goto mrm;

        }

#endif

        if (x == 0x70) {

          x = *++p & 0x1f;

          result++;

          if (x >= 2) {

            p += 2;

            result += 2;

            goto mrm;

          }

        }

        p++;

        result++;

        x = map_op2[*++p];

        break;

      case 8:

        result -= (prefixes & 2);

        /* fallthrough */

      case 9:

mrm:

        /* ModR/M and possibly SIB. */

        result += (x & 15);

        x = *++p;

        switch (x >> 6) {

          case 0:

            if ((x & 7) == 5) {

              return result + 4;

            }

            break;

          case 1:

            result++;

            break;

          case 2:

            result += 4;

            break;

          case 3:

            return result;

        }

        if ((x & 7) == 4) {

          result++;

          if (x < 0x40 && (p[1] & 7) == 5) {

            result += 4;

          }

        }

        return result;

    }

  }

}

typedef IR_SET_ALIGNED(1, uint16_t unaligned_uint16_t);

typedef IR_SET_ALIGNED(1, int32_t unaligned_int32_t);

static int ir_patch_code(const void *code, size_t size, const void *from_addr, const void *to_addr)

{

  int ret = 0;

  uint8_t *p, *end;

  p = (uint8_t*)code;

  end = p + size - 4;

  while (p < end) {

    if ((*(unaligned_uint16_t*)p & 0xf0ff) == 0x800f && p + *(unaligned_int32_t*)(p+2) == (uint8_t*)from_addr - 6) {

      *(unaligned_int32_t*)(p+2) = ((uint8_t*)to_addr - (p + 6));

      ret++;

    } else if (*p == 0xe9 && p + *(unaligned_int32_t*)(p+1) == (uint8_t*)from_addr - 5) {

      *(unaligned_int32_t*)(p+1) = ((uint8_t*)to_addr - (p + 5));

      ret++;

    }

    p += _asm_x86_inslen(p);

  }

  if (ret) {

    ir_mem_flush((void*)code, size);

  }

  return ret;

}

#elif defined(IR_TARGET_AARCH64)

static int ir_patch_code(const void *code, size_t size, const void *from_addr, const void *to_addr)

{

  int ret = 0;

  uint8_t *p, *end;

  const void *veneer = NULL;

  ptrdiff_t delta;

  end = (uint8_t*)code;

  p = end + size;

  while (p > end) {

    uint32_t *ins_ptr;

    uint32_t ins;

    p -= 4;

    ins_ptr = (uint32_t*)p;

    ins = *ins_ptr;

    if ((ins & 0xfc000000u) == 0x14000000u) {

      // B (imm26:0..25)

      delta = (uint32_t*)from_addr - ins_ptr;

      if (((ins ^ (uint32_t)delta) & 0x01ffffffu) == 0) {

        delta = (uint32_t*)to_addr - ins_ptr;

        if (((delta + 0x02000000) >> 26) != 0) {

          abort(); // branch target out of range

        }

        *ins_ptr = (ins & 0xfc000000u) | ((uint32_t)delta & 0x03ffffffu);

        ret++;

        if (!veneer) {

          veneer = p;

        }

      }

    } else if ((ins & 0xff000000u) == 0x54000000u ||

               (ins & 0x7e000000u) == 0x34000000u) {

      // B.cond, CBZ, CBNZ (imm19:5..23)

      delta = (uint32_t*)from_addr - ins_ptr;

      if (((ins ^ ((uint32_t)delta << 5)) & 0x00ffffe0u) == 0) {

        delta = (uint32_t*)to_addr - ins_ptr;

        if (((delta + 0x40000) >> 19) != 0) {

          if (veneer) {

            delta = (uint32_t*)veneer - ins_ptr;

            if (((delta + 0x40000) >> 19) != 0) {

              abort(); // branch target out of range

            }

          } else {

            abort(); // branch target out of range

          }

        }

        *ins_ptr = (ins & 0xff00001fu) | (((uint32_t)delta & 0x7ffffu) << 5);

        ret++;

      }

    } else if ((ins & 0x7e000000u) == 0x36000000u) {

      // TBZ, TBNZ (imm14:5..18)

      delta = (uint32_t*)from_addr - ins_ptr;

      if (((ins ^ ((uint32_t)delta << 5)) & 0x0007ffe0u) == 0) {

        delta = (uint32_t*)to_addr - ins_ptr;

        if (((delta + 0x2000) >> 14) != 0) {

          if (veneer) {

            delta = (uint32_t*)veneer - ins_ptr;

            if (((delta + 0x2000) >> 14) != 0) {

              abort(); // branch target out of range

            }

          } else {

            abort(); // branch target out of range

          }

        }

        *ins_ptr = (ins & 0xfff8001fu) | (((uint32_t)delta & 0x3fffu) << 5);

        ret++;

      }

    }

  }

  if (ret) {

    ir_mem_flush((void*)code, size);

  }

  return ret;

}

#endif

int ir_patch(const void *code, size_t size, uint32_t jmp_table_size, const void *from_addr, const void *to_addr)

{

  int ret = 0;

  if (jmp_table_size) {

    const void **jmp_slot = (const void **)((char*)code + IR_ALIGNED_SIZE(size, sizeof(void*)));

    do {

      if (*jmp_slot == from_addr) {

        *jmp_slot = to_addr;

        ret++;

      }

      jmp_slot++;

    } while (--jmp_table_size);

  }

  ret += ir_patch_code(code, size, from_addr, to_addr);

  return ret;

}