SSA 形式与编译器优化

class LinkEval:
    """带路径压缩的 Link-Eval 数据结构"""

    def __init__(self, n):
        self.ancestor = [0] * n
        self.label = list(range(n))
        self.semi = list(range(n))

    def find(self, v):
        if self.ancestor[v] == 0:
            return v
        u = self.find(self.ancestor[v])
        if self.semi[self.label[self.ancestor[v]]] < self.semi[self.label[v]]:
            self.label[v] = self.label[self.ancestor[v]]
        self.ancestor[v] = self.ancestor[self.ancestor[v]]
        return self.label[v]

    def link(self, v, w):
        self.ancestor[w] = v

def lengauer_tarjan(cfg, entry):
    """Lengauer-Tarjan 支配树算法，O(E·α(V))"""
    dfn = {}
    vertex = []
    parent = {}
    pred = {v: [] for v in cfg}
    counter = [0]

    def dfs(v):
        dfn[v] = counter[0]
        vertex.append(v)
        counter[0] += 1
        for w in cfg.get(v, []):
            pred[w].append(v)
            if w not in dfn:
                parent[w] = v
                dfs(w)

    dfs(entry)
    n = counter[0]

    semi = list(range(n))
    idom = [0] * n
    ancestor = [0] * n
    label = list(range(n))
    bucket = [[] for _ in range(n)]

    def find(v):
        if ancestor[v] == 0:
            return v
        u = find(ancestor[v])
        if semi[label[ancestor[v]]] < semi[label[v]]:
            label[v] = label[ancestor[v]]
        ancestor[v] = ancestor[ancestor[v]]
        return label[v]

    # 按 DFS 逆序计算 sdom
    for i in range(n - 1, 0, -1):
        w = vertex[i]
        for v in pred.get(w, []):
            if v not in dfn:
                continue
            u = find(dfn[v])
            if semi[u] < semi[i]:
                semi[i] = semi[u]
        bucket[semi[i]].append(i)
        ancestor[i] = dfn[parent[w]]
        p = dfn[parent[w]]
        for v in bucket[p]:
            u = find(v)
            idom[v] = p if semi[u] == semi[v] else u
        bucket[p].clear()

    # 修正 idom
    for i in range(1, n):
        if idom[i] != semi[i]:
            idom[i] = idom[idom[i]]

    return {vertex[i]: vertex[idom[i]] for i in range(1, n)}

def compute_df(cfg, idom):
    """O(E) 时间计算支配边界"""
    df = {b: set() for b in cfg}
    preds = {b: [] for b in cfg}
    for v in cfg:
        for w in cfg[v]:
            preds[w].append(v)

    for y in cfg:
        if len(preds[y]) < 2:
            continue
        for x in preds[y]:
            runner = x
            while runner is not None and runner != idom.get(y):
                df[runner].add(y)
                runner = idom.get(runner)
    return df

def place_phi(cfg, df, var_defs):
    """Cytron 算法放置 φ 函数"""
    phi = {v: set() for v in var_defs}
    for v, def_blocks in var_defs.items():
        worklist = list(def_blocks)
        visited = set(def_blocks)
        while worklist:
            block = worklist.pop()
            for y in df.get(block, set()):
                if y not in phi[v]:
                    phi[v].add(y)
                    if y not in visited:
                        visited.add(y)
                        worklist.append(y)
    return phi

def place_pruned_phi(cfg, df, var_defs, live_in):
    """Pruned SSA：只在变量活跃的位置放置 φ 函数"""
    phi = {v: set() for v in var_defs}
    for v, def_blocks in var_defs.items():
        worklist = list(def_blocks)
        visited = set(def_blocks)
        while worklist:
            block = worklist.pop()
            for y in df.get(block, set()):
                if y not in phi[v] and v in live_in.get(y, set()):
                    phi[v].add(y)
                    if y not in visited:
                        visited.add(y)
                        worklist.append(y)
    return phi

def rename(cfg, dom_tree, phi_blocks, entry):
    """SSA 变量重命名（支配树前序遍历）"""
    counter = {}
    stack = {}

    def new_name(var):
        i = counter.get(var, 0)
        counter[var] = i + 1
        stack.setdefault(var, []).append(i)
        return f"{var}.{i}"

    def cur(var):
        s = stack.get(var, [])
        return f"{var}.{s[-1]}" if s else f"{var}.undef"

    def visit(block):
        saved = {v: len(stack.get(v, [])) for v in counter}

        # φ 函数定义
        for var in phi_blocks.get(block, []):
            new_name(var)

        # 普通指令：先重命名 use，再重命名 def
        for instr in get_instrs(block):
            instr.uses = [cur(u) for u in instr.uses]
            if instr.dst:
                instr.dst = new_name(instr.dst)

        # 填充后继的 φ 参数
        for succ in cfg[block]:
            for var in phi_blocks.get(succ, []):
                add_phi_arg(succ, var, block, cur(var))

        # 递归处理支配树子节点
        for child in dom_tree.get(block, []):
            visit(child)

        # 回溯栈
        for v in counter:
            target = saved.get(v, 0)
            while len(stack.get(v, [])) > target:
                stack[v].pop()

    visit(entry)

def sccp(blocks, cfg, entry):
    """Sparse Conditional Constant Propagation"""
    TOP, BOT = "__top__", "__bot__"
    lat = {}
    exe_edges = set()
    visited = set()
    cfg_wl, ssa_wl = [], []

    def get(v):
        return lat.get(v, TOP)

    def meet(a, b):
        if a == TOP: return b
        if b == TOP: return a
        return a if a == b else BOT

    def set_lat(v, val):
        old = get(v)
        if old == BOT:
            return False
        nv = meet(old, val)
        if nv != old:
            lat[v] = nv
            return True
        return False

    def eval_op(op, args):
        vals = [get(a) for a in args]
        if any(v == BOT for v in vals): return BOT
        if any(v == TOP for v in vals): return TOP
        a, b = vals[0], vals[1] if len(vals) > 1 else None
        if op == "add": return a + b
        if op == "sub": return a - b
        if op == "mul": return a * b
        if op == "lt":  return 1 if a < b else 0
        return BOT

    cfg_wl.append((None, entry))

    while cfg_wl or ssa_wl:
        while cfg_wl:
            src, dst = cfg_wl.pop()
            if (src, dst) in exe_edges:
                continue
            exe_edges.add((src, dst))
            if dst not in visited:
                visited.add(dst)
                for instr in blocks[dst]:
                    if instr.op == "const":
                        if set_lat(instr.dst, instr.value):
                            ssa_wl.append(instr.dst)
                    elif instr.op == "phi":
                        val = TOP
                        for pred, arg in instr.args:
                            if (pred, dst) in exe_edges:
                                val = meet(val, get(arg))
                        if set_lat(instr.dst, val):
                            ssa_wl.append(instr.dst)
                    elif instr.dst:
                        val = eval_op(instr.op, instr.args)
                        if set_lat(instr.dst, val):
                            ssa_wl.append(instr.dst)
                    if instr.op == "br":
                        succs = cfg[dst]
                        if len(succs) == 1:
                            cfg_wl.append((dst, succs[0]))
                        else:
                            c = get(instr.args[0])
                            if c == BOT:
                                for s in succs:
                                    cfg_wl.append((dst, s))
                            elif c != TOP:
                                cfg_wl.append((dst, succs[0 if c else 1]))

        while ssa_wl:
            var = ssa_wl.pop()
            for bname in visited:
                for instr in blocks[bname]:
                    if var in instr.args and instr.dst:
                        val = eval_op(instr.op, instr.args)
                        if set_lat(instr.dst, val):
                            ssa_wl.append(instr.dst)

    return lat, visited

def gvn(blocks, dom_tree, entry):
    """基于支配树的 GVN（hash-consing）"""
    vn_table = {}   # (op, vn1, vn2) -> value number
    var_vn = {}     # ssa_var -> value number
    next_vn = [0]

    def get_vn(var):
        return var_vn.get(var, var)

    def lookup_or_add(expr):
        if expr in vn_table:
            return vn_table[expr]
        vn = next_vn[0]
        next_vn[0] += 1
        vn_table[expr] = vn
        return vn

    def process(block):
        for instr in blocks[block]:
            if instr.dst:
                expr = (instr.op,) + tuple(get_vn(a) for a in instr.args)
                var_vn[instr.dst] = lookup_or_add(expr)
        for child in dom_tree.get(block, []):
            process(child)

    process(entry)
    return var_vn

def aggressive_dce(blocks):
    """积极死代码消除：从有副作用的指令反向标记"""
    live = set()
    worklist = []

    # 有副作用的指令始终活
    for bname, instrs in blocks.items():
        for instr in instrs:
            if instr.op in ("store", "call", "ret", "br"):
                live.add(instr)
                worklist.append(instr)

    # 反向传播
    while worklist:
        instr = worklist.pop()
        for operand in instr.args:
            defn = find_def(operand)
            if defn and defn not in live:
                live.add(defn)
                worklist.append(defn)

    # 删除非活指令
    for bname in blocks:
        blocks[bname] = [i for i in blocks[bname] if i in live or not i.dst]

def cse(blocks, dom_tree, entry):
    """基于支配树的 CSE"""
    def process(block, available):
        local = dict(available)
        for instr in blocks[block]:
            if instr.dst and instr.op not in ("call", "store", "load"):
                key = (instr.op,) + tuple(instr.args)
                if key in local:
                    replace_all_uses(instr.dst, local[key])
                else:
                    local[key] = instr.dst
        for child in dom_tree.get(block, []):
            process(child, local)

    process(entry, {})

def linearize_parallel_copy(copies):
    """将并行拷贝 [(dst, src), ...] 线性化为顺序拷贝"""
    result = []
    pending = [(d, s) for d, s in copies if d != s]

    while pending:
        progress = False
        next_round = []
        sources = {s for _, s in pending}

        for dst, src in pending:
            if dst not in sources:
                result.append((dst, src))
                progress = True
            else:
                next_round.append((dst, src))
        pending = next_round

        if not progress and pending:
            # 循环依赖——引入临时变量打破
            dst, src = pending[0]
            temp = f"__tmp_{dst}"
            result.append((temp, src))
            pending = [(d, temp if s == src else s)
                       for d, s in pending]

    return result

def deconstruct_ssa(blocks, phis, cfg):
    """SSA 解构：消除 φ 函数"""
    for block, phi_list in phis.items():
        if not phi_list:
            continue
        pred_copies = {}
        for phi_var, _, args in phi_list:
            for pred, src in args.items():
                pred_copies.setdefault(pred, []).append((phi_var, src))
        for pred, copies in pred_copies.items():
            seq = linearize_parallel_copy(copies)
            insert_copies_at_end(pred, seq)
        phis[block] = []

define i32 @sum(i32 %n) {
entry:
  br label %loop
loop:
  %i = phi i32 [0, %entry], [%i.next, %loop]
  %s = phi i32 [0, %entry], [%s.next, %loop]
  %s.next = add i32 %s, %i
  %i.next = add i32 %i, 1
  %cond = icmp slt i32 %i.next, %n
  br i1 %cond, label %loop, label %exit
exit:
  ret i32 %s.next
}

; 提升前                          ; 提升后
%ptr = alloca i32                 ; (deleted)
store i32 %x, i32* %ptr          ; (deleted)
store i32 42, i32* %ptr          ; (deleted)
%val = load i32, i32* %ptr       ret i32 42
ret i32 %val

clang -emit-llvm -S -O0 example.c -o example.ll
opt -mem2reg -S example.ll -o example_ssa.ll
opt -sccp -S example_ssa.ll -o example_sccp.ll
opt -gvn -S example_ssa.ll -o example_gvn.ll

"""SSA Construction + SCCP：完整可运行实现"""
from collections import defaultdict, deque


class Instr:
    def __init__(self, op, dst=None, args=None, value=None):
        self.op = op
        self.dst = dst
        self.args = args or []
        self.value = value
        self._orig = None  # φ 的原始变量名

    def __repr__(self):
        if self.op == "const":
            return f"{self.dst} = {self.value}"
        if self.op == "phi":
            pairs = ", ".join(f"[{v}, {b}]" for b, v in self.args)
            return f"{self.dst} = phi({pairs})"
        if self.op == "br":
            return f"br {', '.join(str(a) for a in self.args)}"
        if self.op == "ret":
            return f"ret {self.args[0]}"
        if self.op in ("add", "sub", "mul", "lt"):
            return f"{self.dst} = {self.op} {self.args[0]}, {self.args[1]}"
        return f"{self.op} {self.dst} {self.args}"


class BB:
    def __init__(self, name):
        self.name = name
        self.instrs = []
        self.succs = []
        self.preds = []


class CFG:
    def __init__(self):
        self.blocks = {}
        self.entry = None

    def add(self, bb):
        self.blocks[bb.name] = bb
        if not self.entry:
            self.entry = bb.name

    def edge(self, a, b):
        self.blocks[a].succs.append(b)
        self.blocks[b].preds.append(a)


# ---------- 支配树（Cooper-Harvey-Kennedy 迭代算法） ----------

def compute_doms(cfg):
    names = list(cfg.blocks.keys())
    entry = cfg.entry
    doms = {b: set(names) for b in names}
    doms[entry] = {entry}
    changed = True
    while changed:
        changed = False
        for b in names:
            if b == entry or not cfg.blocks[b].preds:
                continue
            new = set.intersection(*(doms[p] for p in cfg.blocks[b].preds))
            new.add(b)
            if new != doms[b]:
                doms[b] = new
                changed = True
    return doms


def compute_idom(cfg, doms):
    idom = {}
    for b in cfg.blocks:
        if b == cfg.entry:
            continue
        strict = doms[b] - {b}
        for d in strict:
            if all(d2 == d or d not in doms[d2] or d2 not in strict
                   for d2 in strict):
                idom[b] = d
                break
    return idom


def dom_tree(idom):
    ch = defaultdict(list)
    for b, d in idom.items():
        ch[d].append(b)
    return dict(ch)


# ---------- 支配边界 ----------

def compute_df(cfg, idom):
    df = {b: set() for b in cfg.blocks}
    for y in cfg.blocks:
        preds = cfg.blocks[y].preds
        if len(preds) < 2:
            continue
        for x in preds:
            r = x
            while r and r != idom.get(y):
                df[r].add(y)
                r = idom.get(r)
    return df


# ---------- φ 放置 ----------

def collect_defs(cfg):
    defs = defaultdict(set)
    for bn, bb in cfg.blocks.items():
        for ins in bb.instrs:
            if ins.dst:
                defs[ins.dst].add(bn)
    return defs


def place_phi(df, defs):
    phi = defaultdict(set)  # block -> {var, ...}
    for v, dblks in defs.items():
        wl = deque(dblks)
        vis = set(dblks)
        while wl:
            b = wl.popleft()
            for y in df.get(b, set()):
                if v not in phi[y]:
                    phi[y].add(v)
                    if y not in vis:
                        vis.add(y)
                        wl.append(y)
    return phi


# ---------- 重命名 ----------

def rename(cfg, dtree, phi_map, entry):
    cnt = defaultdict(int)
    stk = defaultdict(list)

    def nn(v):
        i = cnt[v]; cnt[v] += 1
        name = f"{v}.{i}"; stk[v].append(name); return name

    def cur(v):
        return stk[v][-1] if stk[v] else f"{v}.undef"

    def visit(bn):
        bb = cfg.blocks[bn]
        saved = {v: len(stk[v]) for v in cnt}

        # 插入 φ 指令
        new_instrs = []
        for v in sorted(phi_map.get(bn, set())):
            ins = Instr("phi", dst=nn(v))
            ins._orig = v
            new_instrs.append(ins)

        # 重命名已有指令
        for ins in bb.instrs:
            ins.args = [
                cur(a) if isinstance(a, str) and a.startswith("%") else a
                for a in ins.args
            ]
            if ins.dst:
                ins.dst = nn(ins.dst)
            new_instrs.append(ins)
        bb.instrs = new_instrs

        # 填充后继 φ 参数
        for s in bb.succs:
            for ins in cfg.blocks[s].instrs:
                if ins.op == "phi":
                    ins.args.append((bn, cur(ins._orig)))

        for ch in dtree.get(bn, []):
            visit(ch)

        for v in cnt:
            while len(stk[v]) > saved.get(v, 0):
                stk[v].pop()

    visit(entry)


# ---------- 测试 ----------

def build_example():
    """
    Entry: x=1, y=2 → B1
    B1: if x<10 → B2 else B3
    B2: x=x+1, y=y+x → B1
    B3: return y
    """
    cfg = CFG()
    e = BB("entry")
    e.instrs = [Instr("const", "%x", value=1),
                Instr("const", "%y", value=2),
                Instr("br", args=["B1"])]
    cfg.add(e)

    b1 = BB("B1")
    b1.instrs = [Instr("lt", "%cond", ["%x", 10]),
                 Instr("br", args=["%cond", "B2", "B3"])]
    cfg.add(b1)

    b2 = BB("B2")
    b2.instrs = [Instr("add", "%x", ["%x", 1]),
                 Instr("add", "%y", ["%y", "%x"]),
                 Instr("br", args=["B1"])]
    cfg.add(b2)

    b3 = BB("B3")
    b3.instrs = [Instr("ret", args=["%y"])]
    cfg.add(b3)

    cfg.edge("entry","B1"); cfg.edge("B1","B2")
    cfg.edge("B1","B3");    cfg.edge("B2","B1")
    return cfg


def main():
    print("=" * 55)
    print("SSA Construction Demo")
    print("=" * 55)

    cfg = build_example()
    doms = compute_doms(cfg)
    idom = compute_idom(cfg, doms)
    dt = dom_tree(idom)
    df = compute_df(cfg, idom)
    defs = collect_defs(cfg)
    phi_map = place_phi(df, defs)

    print("\nidom:", {b: idom[b] for b in sorted(idom)})
    print("DF:  ", {b: sorted(df[b]) for b in sorted(df) if df[b]})
    print("phi: ", {b: sorted(phi_map[b]) for b in sorted(phi_map) if phi_map[b]})

    rename(cfg, dt, phi_map, cfg.entry)

    print("\n--- SSA IR ---")
    for bn in ["entry", "B1", "B2", "B3"]:
        bb = cfg.blocks[bn]
        print(f"{bn}:  (preds={bb.preds})")
        for ins in bb.instrs:
            print(f"  {ins}")
    print("\nSSA verification:", end=" ")
    seen = set()
    ok = True
    for bb in cfg.blocks.values():
        for ins in bb.instrs:
            if ins.dst:
                if ins.dst in seen:
                    print(f"FAIL: {ins.dst} defined twice"); ok = False
                seen.add(ins.dst)
    if ok:
        print("PASS (every variable defined exactly once)")


if __name__ == "__main__":
    main()