/// This file handles the (context-free) encoding of x86 instructions.
/// It may substitute an instruction for a shorter one, but is unable to fuse instructions.

// REFERENCES:
//   http://ref.x86asm.net/index.html (geek64-abc)
//   https://wiki.osdev.org/X86-64_Instruction_Encoding
//   https://defuse.ca/online-x86-assembler.htm

#include "format.h"
#include "x86encode.h"

#define REX 0x40
// REX prefix with 64-bit operands set
#define REX_W 0x48
#define REX_R 0x44
#define REX_X 0x42
#define REX_B 0x41

// lower 3 bits of register (not including part encoded in REX)
#define REG(r) (r & 7)

#define MODRM_RR (3 << 6)
#define MODRM_RM (0 << 6)
#define MODRM_RM8 (1 << 6)
#define MODRM_RM32 (2 << 6)

static void x86_opt_rexr(reg reg) {
    if (reg >= R8) {
        append_u8(REX | REX_R);
    }
}

static void x86_rexwr(reg reg) {
    uint8_t rex = REX | REX_W;
    if (reg >= R8) rex |= REX_R;
    append_u8(rex);
}

static void x86_rexwb(reg b) {
    uint8_t rex = REX | REX_W;
    if (b >= R8) rex |= REX_B;
    append_u8(rex);
}

static void x86_rexwrb(reg r, reg b) {
    uint8_t rex = REX | REX_W;
    if (r >= R8) rex |= REX_R;
    if (b >= R8) rex |= REX_B;
    append_u8(rex);
}

static void x86_modrr(reg r, reg b) {
    append_u8(MODRM_RR | (REG(r) << 3) | REG(b));
}

static void x86_modrm(reg r, reg b) {
    append_u8(MODRM_RM | (REG(r) << 3) | REG(b));
}

static void x86_modrm8(reg r, reg b) {
    append_u8(MODRM_RM8 | (REG(r) << 3) | REG(b));
}

static void x86_modrm32(reg r, reg b) {
    append_u8(MODRM_RM32 | (REG(r) << 3) | REG(b));
}

static void x86_modxm(uint8_t ext, reg b) {
    append_u8(MODRM_RR | (ext << 3) | REG(b));
}

static void x86_enc_opr(uint8_t op, reg reg) {
    x86_opt_rexr(reg);
    append_u8(op + REG(reg));
}

static void x86_enc_rexw_opr(uint8_t op, reg reg) {
    x86_rexwr(reg);
    append_u8(op + REG(reg));
}

static void x86_enc_opr_imm32(uint8_t op, reg reg, uint32_t imm) {
    x86_opt_rexr(reg);
    append_u8(op + REG(reg));
    append_u32(imm);
}

static void x86_enc_rexw_opr_imm32(uint8_t op, reg reg, uint32_t imm) {
    x86_rexwr(reg);
    append_u8(op + REG(reg));
    append_u32(imm);
}

static void x86_enc_rexw_opr_imm64(uint8_t op, reg reg, uint64_t imm) {
    x86_rexwr(reg);
    append_u8(op + REG(reg));
    append_u64(imm);
}

static void x86_enc_rexw_modrr(uint8_t op, reg r, reg m) {
    x86_rexwrb(r, m);
    append_u8(op);
    x86_modrr(r, m);
}

static void x86_enc_rexw_modrm(uint8_t op, reg r, reg b) {
    x86_rexwrb(r, b);
    append_u8(op);
    x86_modrm(r, b);
}

static void x86_enc_rexw_modrm8(uint8_t op, reg r, reg b, int8_t disp) {
    x86_rexwrb(r, b);
    append_u8(op);
    x86_modrm8(r, b);
    append_u8(disp);
}

static void x86_enc_rexw_modrm32(uint8_t op, reg r, reg b, int32_t disp) {
    x86_rexwrb(r, b);
    append_u8(op);
    x86_modrm32(r, b);
    append_u32(disp);
}

static void x86_enc_rexw_modrmd(uint8_t op, reg r, reg b, int32_t disp) {
    if (disp == 0 && b != BP && b != R13) {
        x86_enc_rexw_modrm(op, r, b);
    } else if (disp >= INT8_MIN && disp <= INT8_MAX) {
        x86_enc_rexw_modrm8(op, r, b, (int8_t) disp);
    } else {
        x86_enc_rexw_modrm32(op, r, b, disp);
    }
}

static void x86_enc_rexw_modxm_imm8(uint8_t op, uint8_t ext, reg m, uint8_t imm) {
    x86_rexwb(m);
    append_u8(op);
    x86_modxm(ext, m);
    append_u8(imm);
}

static void x86_enc_rexw_modxm_imm32(uint8_t op, uint8_t ext, reg m, uint32_t imm) {
    x86_rexwb(m);
    append_u8(op);
    x86_modxm(ext, m);
    append_u32(imm);
}

static void x86_enc_rexw_modxm_imm(uint8_t op, uint8_t ext, reg m, uint32_t imm) {
    if (imm <= UINT8_MAX) {
        x86_enc_rexw_modxm_imm8(op, ext, m, (uint8_t) imm);
    } else {
        x86_enc_rexw_modxm_imm32(op, ext, m, imm);
    }
}

static void x86_enc_disp8(uint8_t op, int8_t disp) {
    uint8_t buf[2] = { op, (uint8_t) disp };
    append_data(2, buf);
}

static void x86_enc_disp32(uint8_t op, int32_t disp) {
    append_u8(op);
    append_u32((uint32_t) disp);
}

void x86_inst_mov_r64_imm64(reg dest, uint64_t imm) {
    x86_enc_rexw_opr_imm64(0xb8, dest, imm);
}

void x86_inst_mov_r64_imm32(reg dest, uint32_t imm) {
    x86_enc_opr_imm32(0xb8, dest, imm);
}

void x86_inst_mov_r64_imm(reg dest, uint64_t imm) {
    // TODO: xor if 0, use inc and dec, 16-bit and 8-bit immediates, sign extension
    if (imm <= UINT32_MAX) {
        x86_inst_mov_r64_imm32(dest, (uint32_t) imm);
    } else {
        x86_inst_mov_r64_imm64(dest, imm);
    }
}

void x86_inst_mov_r64_r64(reg dest, reg src) {
    x86_enc_rexw_modrr(0x8b, dest, src);
}

void x86_inst_mov_r64_m64(reg dest, reg src) {
    x86_enc_rexw_modrm(0x8b, dest, src);
}

void x86_inst_mov_r64_m64_disp8(reg dest, reg src, int8_t disp) {
    x86_enc_rexw_modrm8(0x8b, dest, src, disp);
}

void x86_inst_mov_r64_m64_disp32(reg dest, reg src, int32_t disp) {
    x86_enc_rexw_modrm32(0x8b, dest, src, disp);
}

void x86_inst_mov_r64_m64_disp(reg dest, reg src, int32_t disp) {
    x86_enc_rexw_modrmd(0x8b, dest, src, disp);
}

void x86_inst_mov_m64_r64(reg dest, reg src) {
    x86_enc_rexw_modrm(0x8a, src, dest);
}

void x86_inst_mov_m64_r64_disp8(reg dest, reg src, int8_t disp) {
    x86_enc_rexw_modrm8(0x8a, src, dest, disp);
}

void x86_inst_mov_m64_r64_disp32(reg dest, reg src, int32_t disp) {
    x86_enc_rexw_modrm32(0x8a, src, dest, disp);
}

void x86_inst_mov_m64_r64_disp(reg dest, reg src, int32_t disp) {
    x86_enc_rexw_modrmd(0x8a, src, dest, disp);
}

void x86_inst_push_r64(reg reg) {
    x86_enc_opr(0x50, reg);
}

void x86_inst_pop_r64(reg reg) {
    x86_enc_opr(0x58, reg);
}

void x86_inst_test_r64_r64(reg r1, reg r2) {
    x86_enc_rexw_modrr(0x85, r1, r2);
}

void x86_inst_jmp_disp8(int8_t disp) {
    x86_enc_disp8(0xeb, disp);
}

void x86_inst_jmp_disp32(int32_t disp) {
    x86_enc_disp32(0xe9, disp);
}

void x86_inst_jmp_disp(int32_t disp) {
    int32_t disp8 = disp + X86_JMP_DISP8_SIZE;
    if (disp8 >= INT8_MIN && disp8 <= INT8_MAX) {
        x86_inst_jmp_disp8((int8_t) disp8);
    } else {
        x86_inst_jmp_disp32(disp + X86_JMP_DISP32_SIZE);
    }
}

void x86_inst_jmp_disp32_op(void) {
    append_u8(0xe9);
}

void x86_inst_jnz_disp8(int8_t disp) {
    x86_enc_disp8(0x75, disp);
}

void x86_inst_jnz_disp32(int32_t disp) {
    append_u8(0x0f);
    x86_enc_disp32(0x85, disp);
}

void x86_inst_jnz_disp(int32_t disp) {
    int32_t disp8 = disp + X86_JNZ_DISP8_SIZE;
    if (disp8 >= INT8_MIN && disp8 <= INT8_MAX) {
        x86_inst_jnz_disp8((int8_t) disp);
    } else {
        x86_inst_jnz_disp32(disp + X86_JNZ_DISP32_SIZE);
    }
}

void x86_inst_jnz_disp32_op(void) {
    append_u8(0x0f);
    append_u8(0x84);
}

void x86_inst_jz_disp8(int8_t disp) {
    x86_enc_disp8(0x74, disp);
}

void x86_inst_jz_disp32(int32_t disp) {
    append_u8(0x0f);
    x86_enc_disp32(0x84, disp);
}

void x86_inst_jz_disp(int32_t disp) {
    int32_t disp8 = disp + X86_JZ_DISP8_SIZE;
    if (disp8 >= INT8_MIN && disp8 <= INT8_MAX) {
        x86_inst_jz_disp8((int8_t) disp);
    } else {
        x86_inst_jz_disp32(disp + X86_JZ_DISP32_SIZE);
    }
}

void x86_inst_jz_disp32_op(void) {
    append_u8(0x0f);
    append_u8(0x84);
}

// TODO: special instructions for AX
void x86_inst_sub_r64_imm8(reg dest, int8_t imm) {
    x86_enc_rexw_modxm_imm8(0x83, 5, dest, (uint8_t) imm);
}

void x86_inst_sub_r64_imm32(reg dest, int32_t imm) {
    x86_enc_rexw_modxm_imm32(0x81, 5, dest, (uint32_t) imm);
}

void x86_inst_sub_r64_imm(reg dest, int32_t imm) {
    x86_enc_rexw_modxm_imm(0x81, 5, dest, (uint32_t) imm);
}

void x86_inst_sub_r64_r64(reg dest, reg src) {
    x86_enc_rexw_modrr(0x2B, dest, src);
}

// TODO: in case of -128, choice of add vs. sub matters!
void x86_inst_add_r64_imm8(reg dest, int8_t imm) {
    x86_enc_rexw_modxm_imm8(0x83, 0, dest, (uint8_t) imm);
}

void x86_inst_add_r64_r64(reg dest, reg src) {
    x86_enc_rexw_modrr(0x01, src, dest);
}

void x86_inst_syscall(void) {
    const uint8_t buf[2] = { 0x0f, 0x05 };
    append_data(2, buf);
}