/* RISC-V 64-bit (RV64I) instruction builder. */
#include <stdlib.h>

#include "riscv.h"
#include "types.h"

const char *reg_names[] = {
    "zero", "ra", "sp", "gp", "tp",  "t0",  "t1", "t2", "fp", "s1", "a0",
    "a1",   "a2", "a3", "a4", "a5",  "a6",  "a7", "s2", "s3", "s4", "s5",
    "s6",   "s7", "s8", "s9", "s10", "s11", "t3", "t4", "t5", "t6"
};

const bool caller_saved_registers[] = {
    false, false, false, false, false, true,  true, true,  false, false, true,
    true,  true,  true,  true,  true,  true,  true, false, false, false, false,
    false, false, false, false, false, false, true, true,  true,  true
};

const reg_t temp_registers[] = { T1, T2, T3, T4, T5, T6 };

i32 sign_extend(u32 value, int bit_width) {
    if ((value >> (bit_width - 1)) & 1) {
        /* Sign bit is 1, so extend with 1s. */
        return (i32)(value | (~0u << bit_width));
    }
    return (i32)value;
}

i32 align(i32 size, i32 alignment) {
    /* Verify alignment is a power of 2. */
    /* This rounds up to the next multiple of alignment. */
    return (size + alignment - 1) & ~(alignment - 1);
}

/* Creates an I-type instruction struct.
 * Used for immediate operations like ADDI, SLTI and loads like LW. */
static instr_t instr_i(
    opcode_t opcode, funct3_t fn3, reg_t rd, reg_t rs1, i32 imm
) {
    return (instr_t){ .i = { .opcode   = opcode,
                             .rd       = rd,
                             .rs1      = rs1,
                             .funct3   = fn3,
                             .imm_11_0 = (u32)imm } };
}

/* Creates a U-type instruction struct.
 * Used for LUI (Load Upper Immediate) and AUIPC. */
static instr_t instr_u(opcode_t opcode, reg_t rd, i32 imm) {
    return (instr_t){ .u = {
                          .opcode    = opcode,
                          .rd        = rd,
                          .imm_31_12 = (u32)imm & 0xFFFFF,
                      } };
}

/* Creates an R-type instruction struct.
 * Used for register-register operations like ADD, SUB, AND, OR. */
static instr_t instr_r(
    opcode_t opcode, funct3_t fn3, funct7_t fn7, reg_t rd, reg_t rs1, reg_t rs2
) {
    return (instr_t){ .r = { .opcode = opcode,
                             .rd     = rd,
                             .rs1    = rs1,
                             .rs2    = rs2,
                             .funct3 = fn3,
                             .funct7 = fn7 } };
}

/* Creates an S-type instruction struct.
 * Used for store instructions like SW, SH, SB. */
static instr_t instr_s(
    opcode_t opcode, funct3_t fn3, reg_t rs1, reg_t rs2, i32 imm
) {
    return (instr_t){ .s = { .opcode   = opcode,
                             .rs1      = rs1,
                             .rs2      = rs2,
                             .funct3   = fn3,
                             .imm_4_0  = (u32)imm & 0x1F,
                             .imm_11_5 = ((u32)imm >> 5) & 0x7F } };
}

/* Creates an SB-type (branch) instruction struct.
 * Modified S-type used for conditional branches like BEQ, BNE. */
static instr_t instr_sb(
    opcode_t opcode, funct3_t fn3, reg_t rs1, reg_t rs2, i32 imm
) {
    return (instr_t){ .b = { .opcode   = opcode,
                             .rs1      = rs1,
                             .rs2      = rs2,
                             .funct3   = fn3,
                             .imm_11   = (imm >> 11) & 0x1,
                             .imm_4_1  = (imm >> 1) & 0xF,
                             .imm_10_5 = (imm >> 5) & 0x3F,
                             .imm_12   = (imm >> 12) & 0x1 } };
}

/* Creates a UJ-type (jump) instruction struct.
 * Modified U-type used for JAL instruction. */
static instr_t instr_uj(opcode_t opcode, reg_t rd, i32 imm) {
    return (instr_t){ .j = { .opcode    = opcode,
                             .rd        = rd,
                             .imm_20    = (imm >> 20) & 0x1,
                             .imm_10_1  = (imm >> 1) & 0x3FF,
                             .imm_11    = (imm >> 11) & 0x1,
                             .imm_19_12 = (imm >> 12) & 0xFF } };
}

/* Instruction definitions table. */
/* Maps each instruction to its parameters for easy lookup. */
typedef struct {
    ifmt_t   type;    /* Instruction type */
    opcode_t opcode;  /* Opcode value */
    funct3_t funct3;  /* Function3 value (if applicable) */
    funct7_t funct7;  /* Function7 value (if applicable) */
    bool     special; /* Special handling flag */
} idef_t;

/* Instruction definition table */
static const idef_t idefs[] = {
    /* Upper immediate instructions. */
    [I_LUI]   = { IFMT_U, OP_LUI, 0, 0, true },
    [I_AUIPC] = { IFMT_U, OP_AUIPC, 0, 0, true },
    /* Jump instructions. */
    [I_JAL]  = { IFMT_J, OP_JAL, 0, 0, true },
    [I_JALR] = { IFMT_I, OP_JALR, 0, 0, false },
    /* Branch instructions. */
    [I_BEQ]  = { IFMT_B, OP_BRANCH, FUNCT3_BYTE, 0, false },
    [I_BNE]  = { IFMT_B, OP_BRANCH, FUNCT3_HALF, 0, false },
    [I_BLT]  = { IFMT_B, OP_BRANCH, FUNCT3_BYTE_U, 0, false },
    [I_BGE]  = { IFMT_B, OP_BRANCH, FUNCT3_HALF_U, 0, false },
    [I_BLTU] = { IFMT_B, OP_BRANCH, FUNCT3_OR, 0, false },
    [I_BGEU] = { IFMT_B, OP_BRANCH, FUNCT3_AND, 0, false },
    /* Load instructions. */
    [I_LB]  = { IFMT_I, OP_LOAD, FUNCT3_BYTE, 0, false },
    [I_LH]  = { IFMT_I, OP_LOAD, FUNCT3_HALF, 0, false },
    [I_LW]  = { IFMT_I, OP_LOAD, FUNCT3_WORD, 0, false },
    [I_LBU] = { IFMT_I, OP_LOAD, FUNCT3_BYTE_U, 0, false },
    [I_LHU] = { IFMT_I, OP_LOAD, FUNCT3_HALF_U, 0, false },
    /* Store instructions. */
    [I_SB] = { IFMT_S, OP_STORE, FUNCT3_BYTE, 0, false },
    [I_SH] = { IFMT_S, OP_STORE, FUNCT3_HALF, 0, false },
    [I_SW] = { IFMT_S, OP_STORE, FUNCT3_WORD, 0, false },
    /* ALU immediate operations. */
    [I_ADDI]  = { IFMT_I, OP_IMM, FUNCT3_ADD, 0, false },
    [I_SLTI]  = { IFMT_I, OP_IMM, FUNCT3_SLT, 0, false },
    [I_SLTIU] = { IFMT_I, OP_IMM, FUNCT3_SLTU, 0, false },
    [I_XORI]  = { IFMT_I, OP_IMM, FUNCT3_XOR, 0, false },
    [I_ORI]   = { IFMT_I, OP_IMM, FUNCT3_OR, 0, false },
    [I_ANDI]  = { IFMT_I, OP_IMM, FUNCT3_AND, 0, false },
    [I_SLLI]  = { IFMT_I, OP_IMM, FUNCT3_SLL, 0, true },
    [I_SRLI]  = { IFMT_I, OP_IMM, FUNCT3_SRL, 0, true },
    [I_SRAI]  = { IFMT_I, OP_IMM, FUNCT3_SRL, 0, true },
    /* ALU register operations. */
    [I_ADD]  = { IFMT_R, OP_OP, FUNCT3_ADD, FUNCT7_NORMAL, false },
    [I_SUB]  = { IFMT_R, OP_OP, FUNCT3_ADD, FUNCT7_SUB, false },
    [I_SLL]  = { IFMT_R, OP_OP, FUNCT3_SLL, FUNCT7_NORMAL, false },
    [I_SLT]  = { IFMT_R, OP_OP, FUNCT3_SLT, FUNCT7_NORMAL, false },
    [I_SLTU] = { IFMT_R, OP_OP, FUNCT3_SLTU, FUNCT7_NORMAL, false },
    [I_XOR]  = { IFMT_R, OP_OP, FUNCT3_XOR, FUNCT7_NORMAL, false },
    [I_SRL]  = { IFMT_R, OP_OP, FUNCT3_SRL, FUNCT7_NORMAL, false },
    [I_AND]  = { IFMT_R, OP_OP, FUNCT3_AND, FUNCT7_NORMAL, false },
    [I_OR]   = { IFMT_R, OP_OP, FUNCT3_OR, FUNCT7_NORMAL, false },
    /* M extension - multiply and divide. */
    [I_MUL]    = { IFMT_R, OP_OP, FUNCT3_ADD, FUNCT7_MUL, false },
    [I_MULH]   = { IFMT_R, OP_OP, FUNCT3_SLL, FUNCT7_MUL, false },
    [I_MULHSU] = { IFMT_R, OP_OP, FUNCT3_SLT, FUNCT7_MUL, false },
    [I_MULHU]  = { IFMT_R, OP_OP, FUNCT3_SLTU, FUNCT7_MUL, false },
    [I_DIV]    = { IFMT_R, OP_OP, FUNCT3_XOR, FUNCT7_MUL, false },
    [I_DIVU]   = { IFMT_R, OP_OP, FUNCT3_SRL, FUNCT7_MUL, false },
    [I_REM]    = { IFMT_R, OP_OP, FUNCT3_OR, FUNCT7_MUL, false },
    [I_REMU]   = { IFMT_R, OP_OP, FUNCT3_AND, FUNCT7_MUL, false },
    /* Pseudo-instructions. */
    [I_MV]  = { IFMT_I, OP_IMM, FUNCT3_ADD, 0, true },
    [I_JMP] = { IFMT_J, OP_JAL, 0, 0, true },
    [I_NOP] = { IFMT_I, OP_IMM, FUNCT3_ADD, 0, true },
    [I_NOT] = { IFMT_I, OP_IMM, FUNCT3_XOR, 0, true },
    [I_NEG] = { IFMT_R, OP_OP, FUNCT3_ADD, FUNCT7_SUB, true },
    /* System instructions */
    [I_EBREAK] = { IFMT_I, OP_SYSTEM, 0, 0, true },
    [I_ECALL]  = { IFMT_I, OP_SYSTEM, 0, 0, true },
    /* F Extension floating-point instructions */
    [I_FADD_S] = { IFMT_R, OP_OP_FP, FUNCT3_ADD, FUNCT7_FADD_S, false },
    [I_FSUB_S] = { IFMT_R, OP_OP_FP, FUNCT3_ADD, FUNCT7_FSUB_S, false },
    [I_FMUL_S] = { IFMT_R, OP_OP_FP, FUNCT3_ADD, FUNCT7_FMUL_S, false },
    [I_FDIV_S] = { IFMT_R, OP_OP_FP, FUNCT3_ADD, FUNCT7_FDIV_S, false },
    [I_FEQ_S]  = { IFMT_R, OP_OP_FP, FUNCT3_FEQ, FUNCT7_FEQ_S, false },
    [I_FLT_S]  = { IFMT_R, OP_OP_FP, FUNCT3_FLT, FUNCT7_FLT_S, false },
    [I_FLE_S]  = { IFMT_R, OP_OP_FP, FUNCT3_FLE, FUNCT7_FLE_S, false },
    [I_FLW]    = { IFMT_I, OP_LOAD_FP, FUNCT3_WORD_FP, 0, false },
    [I_FSW]    = { IFMT_S, OP_STORE_FP, FUNCT3_WORD_FP, 0, false },
    /* RV64I load/store */
    [I_LWU] = { IFMT_I, OP_LOAD, FUNCT3_WORD_U, 0, false },
    [I_LD]  = { IFMT_I, OP_LOAD, FUNCT3_DOUBLE, 0, false },
    [I_SD]  = { IFMT_S, OP_STORE, FUNCT3_DOUBLE, 0, false },
    /* RV64I immediate W-ops */
    [I_ADDIW] = { IFMT_I, OP_IMM_32, FUNCT3_ADD, 0, false },
    [I_SLLIW] = { IFMT_I, OP_IMM_32, FUNCT3_SLL, 0, true },
    [I_SRLIW] = { IFMT_I, OP_IMM_32, FUNCT3_SRL, 0, true },
    [I_SRAIW] = { IFMT_I, OP_IMM_32, FUNCT3_SRL, 0, true },
    /* RV64I register W-ops */
    [I_ADDW] = { IFMT_R, OP_OP_32, FUNCT3_ADD, FUNCT7_NORMAL, false },
    [I_SUBW] = { IFMT_R, OP_OP_32, FUNCT3_ADD, FUNCT7_SUB, false },
    [I_SLLW] = { IFMT_R, OP_OP_32, FUNCT3_SLL, FUNCT7_NORMAL, false },
    [I_SRLW] = { IFMT_R, OP_OP_32, FUNCT3_SRL, FUNCT7_NORMAL, false },
    [I_SRAW] = { IFMT_R, OP_OP_32, FUNCT3_SRL, FUNCT7_SRA, false },
    /* RV64M W-ops */
    [I_MULW]  = { IFMT_R, OP_OP_32, FUNCT3_ADD, FUNCT7_MUL, false },
    [I_DIVW]  = { IFMT_R, OP_OP_32, FUNCT3_XOR, FUNCT7_MUL, false },
    [I_DIVUW] = { IFMT_R, OP_OP_32, FUNCT3_SRL, FUNCT7_MUL, false },
    [I_REMW]  = { IFMT_R, OP_OP_32, FUNCT3_OR, FUNCT7_MUL, false },
    [I_REMUW] = { IFMT_R, OP_OP_32, FUNCT3_AND, FUNCT7_MUL, false },
};

/* Generates a RISC-V instruction based on the instruction definition */
instr_t instr(iname_t iop, reg_t rd, reg_t rs1, reg_t rs2, i32 imm) {
    const idef_t *def = &idefs[iop];

    /* Handle special cases that need specific processing. */
    if (def->special) {
        switch (iop) {
        case I_LUI:
            return instr_u(OP_LUI, rd, imm);
        case I_AUIPC:
            return instr_u(OP_AUIPC, rd, imm);
        case I_JAL:
            return instr_uj(OP_JAL, rd, imm);
        case I_SLLI:
            return instr_i(OP_IMM, FUNCT3_SLL, rd, rs1, imm & 0x3F);
        case I_SRLI:
            return instr_i(OP_IMM, FUNCT3_SRL, rd, rs1, imm & 0x3F);
        case I_SRAI:
            return instr_i(OP_IMM, FUNCT3_SRL, rd, rs1, (imm & 0x3F) + 0x400);
        case I_SLLIW:
            return instr_i(OP_IMM_32, FUNCT3_SLL, rd, rs1, imm & 0x1F);
        case I_SRLIW:
            return instr_i(OP_IMM_32, FUNCT3_SRL, rd, rs1, imm & 0x1F);
        case I_SRAIW:
            return instr_i(
                OP_IMM_32, FUNCT3_SRL, rd, rs1, (imm & 0x1F) + 0x400
            );
        case I_MV:
            return instr_i(OP_IMM, FUNCT3_ADD, rd, rs1, 0);
        case I_JMP:
            return instr_uj(OP_JAL, ZERO, imm);
        case I_NOP:
            return instr_i(OP_IMM, FUNCT3_ADD, ZERO, ZERO, 0);
        case I_NOT:
            return instr_i(OP_IMM, FUNCT3_XOR, rd, rs1, 1);
        case I_NEG:
            return instr_r(OP_OP, FUNCT3_ADD, FUNCT7_SUB, rd, ZERO, rs1);
        case I_EBREAK:
            /* EBREAK is encoded as all zeros except for the opcode */
            return instr_i(OP_SYSTEM, 0, 0, 0, 1);
        case I_ECALL:
            /* ECALL is encoded as all zeros including immediate */
            return instr_i(OP_SYSTEM, 0, 0, 0, 0);
        default:
            break;
        }
    }

    /* Regular instructions by type. */
    switch (def->type) {
    case IFMT_I:
        return instr_i(def->opcode, def->funct3, rd, rs1, imm);
    case IFMT_R:
        return instr_r(def->opcode, def->funct3, def->funct7, rd, rs1, rs2);
    case IFMT_S:
        return instr_s(def->opcode, def->funct3, rs1, rs2, imm);
    case IFMT_B:
        return instr_sb(def->opcode, def->funct3, rs1, rs2, imm);
    case IFMT_U:
        return instr_u(def->opcode, rd, imm);
    case IFMT_J:
        return instr_uj(def->opcode, rd, imm);
    default:
        abort();
    }
}