package org.jcnc.snow.compiler.backend;

import org.jcnc.snow.compiler.ir.core.IRFunction;
import org.jcnc.snow.compiler.ir.core.IRInstruction;
import org.jcnc.snow.compiler.ir.core.IRValue;
import org.jcnc.snow.compiler.ir.value.IRVirtualRegister;

import java.util.HashMap;
import java.util.Map;

/**
 * 线性扫描寄存器分配器。
 * <p>
 * 分配策略：
 * <ol>
 *   <li>将函数的形式参数按照声明顺序分配到槽编号 0 到 n-1。</li>
 *   <li>随后扫描函数体中的指令列表，每遇到尚未分配槽号的虚拟寄存器，
 *       即为其分配下一个可用的槽编号。</li>
 * </ol>
 */
public final class RegisterAllocator {

    /**
     * 存储虚拟寄存器到物理槽的映射表。
     * 键为虚拟寄存器，值为分配的槽编号。
     */
    private final Map<IRVirtualRegister, Integer> map = new HashMap<>();

    /**
     * 为指定的 IR 函数计算寄存器到槽的映射关系。
     * <p>
     * 先为函数的形式参数分配槽号，再为函数体中出现的其他虚拟寄存器依次分配。
     * </p>
     *
     * @param fn 需要进行寄存器分配的中间表示函数
     * @return   一个不可变的映射表，表示每个虚拟寄存器对应的槽编号
     */
    public Map<IRVirtualRegister, Integer> allocate(IRFunction fn) {
        int next = 0;

        // 1. 为形式参数分配槽编号
        for (IRVirtualRegister param : fn.parameters()) {
            map.put(param, next++);
        }

        // 2. 扫描函数体，分配其他虚拟寄存器的槽编号
        for (IRInstruction inst : fn.body()) {
            if (inst.dest() != null && !map.containsKey(inst.dest())) {
                map.put(inst.dest(), next++);
            }
            for (IRValue operand : inst.operands()) {
                if (operand instanceof IRVirtualRegister vr && !map.containsKey(vr)) {
                    map.put(vr, next++);
                }
            }
        }

        return Map.copyOf(map);
    }
}
