「寄存器分配与图着色」

# 中间代码里的伪寄存器: 数量不受限
IR = [
    ("MOV", "t1", 5),           # t1 = 5
    ("MOV", "t2", 3),           # t2 = 3
    ("ADD", "t3", "t1", "t2"),  # t3 = t1 + t2
    ("MUL", "t4", "t3", 2),     # t4 = t3 * 2
    ("RET", "t4"),
]

# 目标机器只有 4 个物理寄存器可用
PHYSICAL_REGISTERS = [f"r{i}" for i in range(4)]

# 朴素分配: 遇到伪寄存器就占一个物理寄存器
def naive_allocate(ir, n_regs=4):
    table, used = {}, []
    for inst in ir:
        for opnd in inst[1:]:
            if isinstance(opnd, str) and opnd.startswith("t"):
                if opnd not in table:
                    if len(used) >= n_regs:      # 寄存器耗尽
                        table[opnd] = f"spill_{opnd}"
                    else:
                        reg = f"r{len(used)}"
                        used.append(reg)
                        table[opnd] = reg
    return table

print(naive_allocate(IR))

def liveness(cfg, use, defs):
    """use[n]: 块 n 中被读取的变量; defs[n]: 块 n 中被定义的变量."""
    live_in = {n: set() for n in cfg}
    live_out = {n: set() for n in cfg}
    changed = True
    while changed:
        changed = False
        for n in cfg:
            new_out = set()
            for succ in cfg[n]:
                new_out |= live_in[succ]
            new_in = use[n] | (new_out - defs[n])
            if new_in != live_in[n] or new_out != live_out[n]:
                live_in[n] = new_in
                live_out[n] = new_out
                changed = True
    return live_in, live_out

# 简单控制流图: 0 -> 1 -> 2 -> 3 -> 4
cfg = {0: [1], 1: [2], 2: [3], 3: [4], 4: []}
use = {0: set(), 1: {"a"}, 2: {"b"}, 3: {"a", "b"}, 4: set()}
defs = {0: {"a"}, 1: {"b"}, 2: {"c"}, 3: {"d"}, 4: set()}

li, lo = liveness(cfg, use, defs)
for n in cfg:
    print(f"块{n}: live_in={sorted(li[n])} live_out={sorted(lo[n])}")

def build_interference(cfg, live_in, live_out, defs):
    all_vars = set()
    for s in defs.values():
        all_vars |= s
    ig = {v: set() for v in all_vars}
    for n in cfg:
        defined = defs[n]
        live = live_out[n] | (live_in[n] - defined)
        for d in defined:
            for v in live:
                if v != d:
                    ig[d].add(v)
                    ig[v].add(d)
    return ig

ig = build_interference(cfg, li, lo, defs)
for v in sorted(ig):
    print(f"{v} 冲突: {sorted(ig[v])}")

# 冲突图可视化 (点->边)
def render_graph(ig):
    lines = ["graph G {"]
    for v, nbrs in ig.items():
        for n in nbrs:
            if v < n:
                lines.append(f'  "{v}" -- "{n}"')
    lines.append("}")
    return "\n".join(lines)

print(render_graph(ig))

def kempe_coloring(interference, k):
    stack = []
    remaining = {v: set(nbrs) for v, nbrs in interference.items()}
    while remaining:                      # 简化阶段
        node = min(remaining, key=lambda v: len(remaining[v]))
        stack.append(node)
        for nbr in list(remaining[node]):
            remaining[nbr].discard(node)
        del remaining[node]

    color = {}
    while stack:                          # 选择阶段
        n = stack.pop()
        used = {color[x] for x in interference[n] if x in color}
        for c in range(k):
            if c not in used:
                color[n] = c
                break
        else:
            color[n] = None               # 无法着色: 溢出
    return color

print(kempe_coloring(ig, k=3))

from dataclasses import dataclass

@dataclass
class Interval:
    v: str
    start: int
    end: int

def linear_scan(intervals, k):
    intervals = sorted(intervals, key=lambda iv: iv.start)
    active = []                     # 活跃中区间, 按 start 排序
    assign = {}
    for iv in intervals:
        active = [a for a in active if a.end >= iv.start]
        if len(active) < k:
            used = {assign[a.v] for a in active}
            for r in range(k):
                if r not in used:
                    assign[iv.v] = r
                    active.append(iv)
                    break
        else:
            spill = max(active, key=lambda a: a.end)
            if spill.end > iv.end:          # 新区间更长, 挤掉旧的
                assign[iv.v] = "spill"
            else:
                assign[spill.v] = "spill"
                active.remove(spill)
                assign[iv.v] = 0            # 简化: 占据被挤掉的寄存器
                active.append(iv)
    return assign

ivs = [
    Interval("a", 0, 3), Interval("b", 1, 4),
    Interval("c", 2, 5), Interval("d", 3, 6),
]
print(linear_scan(ivs, k=2))

def eliminate_phi(incoming, result):
    """把 phi 节点 (src1, src2 -> result) 翻译成沿各前驱的并行拷贝."""
    copies = [(src, result) for src in incoming]
    # 若存在循环依赖 (a=b; b=a), 用临时变量打破
    for i, (src, dst) in enumerate(copies):
        if any(dst == s for s, _ in copies[i + 1:]):
            tmp = f"__tmp{i}"
            copies[i] = (src, tmp)
            copies.append((tmp, dst))
    return copies

print("phi(x,y)->z:", eliminate_phi(("x", "y"), "z"))
print("phi(y,x)->x:", eliminate_phi(("y", "x"), "x"))  # 循环依赖

def coalesce(pairs, interference):
    """把无冲突的拷贝对合并到同一寄存器."""
    merged = {}
    for src, dst in pairs:
        if not (dst in interference.get(src, set())):
            merged[dst] = merged.get(src, src)
    return merged

def classify_args(param_count):
    """x86-64 SysV: 前 6 个整型参数用寄存器, 其余压栈."""
    regs = ["rdi", "rsi", "rdx", "rcx", "r8", "r9"]
    return {
        i: (regs[i] if i < 6 else f"stack[{i - 6}]")
        for i in range(param_count)
    }

print(classify_args(8))