/* vm_opt.c - 字节码（VM 层）优化 pass
 *
 * 从 stage0/spl_mcode.c 迁出：对已生成的 spl_prog_t 做 prog->prog 优化。
 * 默认不启用（由 splc0 --opt 驱动）。 */

#include "vm_opt.h"

#include "../stage0/include/core_map.h"

#include <stdlib.h>

/* 消除“store 后立刻 load 同一槽位”的往返。
 *
 * 模式（5 条）：laddr X; swap; store k; laddr X; load k
 * 变换（3 条）：laddr X; pick 1; store k
 *   [.., v] -> laddr X   => [.., v, X]
 *           -> pick 1    => [.., v, X, v]
 *           -> store k   => [.., v]        （存 v 到 X，同时栈上保留 v）
 * 与原来 store 后再 load 的效果等价，但去掉 2 条指令与一次内存读。
 *
 * 会真的删除指令，因此必须重定位：相对跳转、函数表 address、
 * 以及作为函数指针的 `push ptr <addr>`。仅当被删/被改写的指令不是
 * 跳转目标或函数入口时才变换（已收集 mark）。 */
static usize peep_hash_int(usize k) { return MAP_HASH_INT(k); }
static int peep_cmp_int(usize a, usize b) { return MAP_CMP_INT(a, b); }

void spl_prog_peephole(spl_prog_t *prog) {
    if (!prog)
        return;
    usize n = vec_size(prog->insns);
    if (n < 5)
        return;
    spl_vm_ins_t *ins = prog->insns.data;

    unsigned char *mark = (unsigned char *)calloc(n, 1); /* 跳转目标/函数入口，受保护 */
    unsigned char *rem = (unsigned char *)calloc(n, 1);
    if (!mark || !rem) {
        free(mark);
        free(rem);
        return;
    }

    for (usize i = 0; i < n; i++) {
        spl_vm_ins_t *x = &ins[i];
        if (x->opcode == SPL_JMP || x->opcode == SPL_BZ || x->opcode == SPL_BNZ) {
            isize t = (isize)i + 1 + (isize)x->imm;
            if (t >= 0 && (usize)t < n)
                mark[(usize)t] = 1;
        }
    }
    vec_for(prog->funcs, fi) {
        spl_vm_func_t *f = &vec_at(prog->funcs, fi);
        if (f->address < n)
            mark[f->address] = 1;
    }

    for (usize i = 0; i + 4 < n; i++) {
        spl_vm_ins_t *a = &ins[i];
        spl_vm_ins_t *b = &ins[i + 1];
        spl_vm_ins_t *c = &ins[i + 2];
        spl_vm_ins_t *d = &ins[i + 3];
        spl_vm_ins_t *e = &ins[i + 4];
        if (a->opcode == SPL_LADDR && b->opcode == SPL_SWAP && c->opcode == SPL_STORE &&
            d->opcode == SPL_LADDR && e->opcode == SPL_LOAD && a->imm == d->imm &&
            c->type == e->type) {
            if (mark[i + 1] || mark[i + 3] || mark[i + 4])
                continue;
            b->opcode = SPL_PICK;
            b->type = SPL_VOID;
            b->imm = 1;
            rem[i + 3] = 1;
            rem[i + 4] = 1;
            i += 4;
        }
    }

    usize kept = 0;
    for (usize i = 0; i < n; i++)
        if (!rem[i])
            kept++;
    if (kept == n) {
        free(mark);
        free(rem);
        return;
    }

    spl_vm_ins_t *fresh = (spl_vm_ins_t *)malloc(sizeof(spl_vm_ins_t) * (kept ? kept : 1));
    usize *map = (usize *)malloc(sizeof(usize) * (n + 1)); /* old index -> new index */
    usize *old_of_new = (usize *)malloc(sizeof(usize) * (kept ? kept : 1));
    if (!fresh || !map || !old_of_new) {
        free(mark);
        free(rem);
        free(fresh);
        free(map);
        free(old_of_new);
        return;
    }

    usize out = 0;
    for (usize i = 0; i < n; i++) {
        if (!rem[i]) {
            map[i] = out;
            old_of_new[out] = i;
            fresh[out] = ins[i];
            out++;
        }
    }
    map[n] = kept;
    for (usize i = n; i-- > 0;)
        if (rem[i])
            map[i] = map[i + 1];

    /* 函数地址 old->new */
    MAP(usize, usize) addr_map;
    map_init(addr_map, peep_hash_int, peep_cmp_int);
    vec_for(prog->funcs, fi) {
        spl_vm_func_t *f = &vec_at(prog->funcs, fi);
        usize old_a = f->address;
        if (old_a < n) {
            usize new_a = map[old_a];
            map_put(addr_map, old_a, new_a);
            f->address = new_a;
        }
    }

    /* 修正相对跳转与函数指针常量 */
    for (usize j = 0; j < kept; j++) {
        spl_vm_ins_t *x = &fresh[j];
        if (x->opcode == SPL_JMP || x->opcode == SPL_BZ || x->opcode == SPL_BNZ) {
            isize t = (isize)old_of_new[j] + 1 + (isize)x->imm;
            if (t >= 0 && (usize)t <= n) {
                usize nt = map[(usize)t];
                x->imm = (spl_vm_val_t)(isize)((isize)nt - (isize)(j + 1));
            }
        } else if (x->opcode == SPL_PUSH && x->type == SPL_PTR && x->imm != 0) {
            usize nv = 0;
            if (map_get(addr_map, (usize)x->imm, &nv))
                x->imm = nv;
        }
    }
    map_free(addr_map);

    /* 指令删除后函数长度变化，重算 ninsns（函数按地址升序排列）。 */
    {
        usize nf = vec_size(prog->funcs);
        for (usize i = 0; i < nf; i++) {
            spl_vm_func_t *f = &vec_at(prog->funcs, i);
            usize next = (i + 1 < nf) ? vec_at(prog->funcs, i + 1).address : kept;
            f->ninsns = next > f->address ? (spl_vm_val_t)(next - f->address) : 0;
        }
    }

    free(prog->insns.data);
    prog->insns.data = fresh;
    prog->insns.size = kept;
    prog->insns.cap = kept;

    free(mark);
    free(rem);
    free(map);
    free(old_of_new);
}

/* 优化 pass 流水线（纯 prog->prog）：当前 P1，后续可追加。 */
void spl_prog_opt(spl_prog_t *prog) { spl_prog_peephole(prog); }