module optimize

import v3.ssa

// --- Simplify trivial phi nodes ---
// A phi is trivial if all its non-self operands resolve to a single value, or if
// it has no uses at all. Trivial phis are replaced and marked dead (nop bitcast).
fn simplify_phi_nodes(mut m ssa.Module) bool {
    mut any_changed := false
    mut changed := true
    for changed {
        changed = false
        for fi in 0 .. m.funcs.len {
            func := m.funcs[fi]
            for blk_id in func.blocks {
                blk := m.blocks[blk_id]
                for val_id in blk.instrs {
                    val := m.values[val_id]
                    if val.kind != .instruction {
                        continue
                    }
                    instr := m.instrs[val.index]
                    if instr.op != .phi {
                        continue
                    }

                    // Dead phi: no uses -> remove.
                    if val.uses.len == 0 {
                        for i := 0; i < instr.operands.len; i += 2 {
                            op_val := instr.operands[i]
                            if op_val != val_id {
                                remove_use(mut m, op_val, val_id)
                            }
                        }
                        mut dead_phi := m.instrs[val.index]
                        dead_phi.op = .bitcast
                        dead_phi.operands = []
                        m.instrs[val.index] = dead_phi
                        changed = true
                        any_changed = true
                        continue
                    }

                    // Trivial phi: all non-self value operands identical.
                    mut unique_val := -1
                    mut is_trivial := true
                    for i := 0; i < instr.operands.len; i += 2 {
                        op_val := instr.operands[i]
                        if op_val == val_id {
                            continue
                        }
                        if unique_val == -1 {
                            unique_val = op_val
                        } else if unique_val != op_val {
                            is_trivial = false
                            break
                        }
                    }
                    if is_trivial && unique_val != -1 {
                        m.replace_uses(val_id, unique_val)
                        mut triv_phi := m.instrs[val.index]
                        triv_phi.op = .bitcast
                        triv_phi.operands = []
                        m.instrs[val.index] = triv_phi
                        changed = true
                        any_changed = true
                    }
                }
            }
        }
    }
    return any_changed
}

// prune_phi_operands drops phi operand pairs whose predecessor block is no longer
// an actual predecessor (e.g. after branch folding / unreachable-block removal),
// keeping phis consistent with the current CFG. A phi reduced to a single distinct
// incoming value is replaced by that value. Requires up-to-date preds + uses.
fn prune_phi_operands(mut m ssa.Module) {
    for fi in 0 .. m.funcs.len {
        for blk_id in m.funcs[fi].blocks {
            if blk_id < 0 || blk_id >= m.blocks.len {
                continue
            }
            preds := m.blocks[blk_id].preds
            for val_id in m.blocks[blk_id].instrs {
                if val_id <= 0 || val_id >= m.values.len || m.values[val_id].kind != .instruction {
                    continue
                }
                idx := m.values[val_id].index
                if m.instrs[idx].op != .phi {
                    continue
                }
                ops := m.instrs[idx].operands
                mut kept := []ssa.ValueID{}
                mut distinct := []ssa.ValueID{}
                for i := 0; i + 1 < ops.len; i += 2 {
                    pred := int(ops[i + 1])
                    if pred in preds {
                        kept << ops[i]
                        kept << ops[i + 1]
                        if ops[i] !in distinct {
                            distinct << ops[i]
                        }
                    }
                }
                if kept.len == ops.len && distinct.len > 1 {
                    continue // unchanged, still a real phi
                }
                if distinct.len == 1 {
                    // Trivial phi -> its single incoming value.
                    m.replace_uses(val_id, distinct[0])
                    mut nop := m.instrs[idx]
                    nop.op = .bitcast
                    nop.operands = []
                    m.instrs[idx] = nop
                } else if distinct.len == 0 {
                    mut nop := m.instrs[idx]
                    nop.op = .bitcast
                    nop.operands = []
                    m.instrs[idx] = nop
                } else {
                    mut pruned := m.instrs[idx]
                    pruned.operands = kept
                    m.instrs[idx] = pruned
                }
            }
        }
    }
}

// --- Critical edge splitting ---
// Inserts an empty block on every edge whose source has multiple successors and
// whose target has multiple predecessors, so phi-elimination copies have a unique
// home. Operands here are raw block ids (v3 convention).
fn split_critical_edges(mut m ssa.Module) {
    build_cfg(mut m)
    for fi in 0 .. m.funcs.len {
        mut edges_to_split := [][]int{}
        func := m.funcs[fi]
        for blk_id in func.blocks {
            if blk_id < 0 || blk_id >= m.blocks.len {
                continue
            }
            if m.blocks[blk_id].succs.len > 1 {
                for succ_id in m.blocks[blk_id].succs {
                    if succ_id < 0 || succ_id >= m.blocks.len {
                        continue
                    }
                    if m.blocks[succ_id].preds.len > 1 {
                        edges_to_split << [int(blk_id), int(succ_id)]
                    }
                }
            }
        }

        for edge in edges_to_split {
            pred_id := edge[0]
            succ_id := edge[1]
            if pred_id < 0 || pred_id >= m.blocks.len || succ_id < 0 || succ_id >= m.blocks.len {
                continue
            }
            split_blk := m.add_block(fi, 'split_${pred_id}_${succ_id}')
            m.add_instr(.jmp, split_blk, 0, [ssa.ValueID(succ_id)])

            // Re-point predecessor terminator from succ to the split block.
            pred_blk := m.blocks[pred_id]
            if pred_blk.instrs.len > 0 {
                term_val_id := pred_blk.instrs[pred_blk.instrs.len - 1]
                if term_val_id >= 0 && term_val_id < m.values.len {
                    idx := m.values[term_val_id].index
                    mut term := m.instrs[idx]
                    for i in 0 .. term.operands.len {
                        if int(term.operands[i]) == succ_id {
                            term.operands[i] = ssa.ValueID(split_blk)
                        }
                    }
                    m.instrs[idx] = term
                }
            }

            // Re-point phi predecessors in succ from pred to the split block.
            for vid in m.blocks[succ_id].instrs {
                if vid < 0 || vid >= m.values.len || m.values[vid].kind != .instruction {
                    continue
                }
                pidx := m.values[vid].index
                if m.instrs[pidx].op != .phi {
                    continue
                }
                mut phi := m.instrs[pidx]
                for i := 1; i < phi.operands.len; i += 2 {
                    if int(phi.operands[i]) == pred_id {
                        phi.operands[i] = ssa.ValueID(split_blk)
                    }
                }
                m.instrs[pidx] = phi
            }
        }
    }
    build_cfg(mut m)
}

// eliminate_phi_nodes lowers phis to `assign` copies placed in predecessor blocks,
// sequencing the per-predecessor parallel copies (Briggs) to respect dependencies.
fn eliminate_phi_nodes(mut m ssa.Module) {
    split_critical_edges(mut m)

    n_blocks := m.blocks.len
    mut pred_copy_dests := [][]int{len: n_blocks}
    mut pred_copy_srcs := [][]int{len: n_blocks}

    for fi in 0 .. m.funcs.len {
        func := m.funcs[fi]
        mut pred_copy_blocks := []int{}

        for blk_id in func.blocks {
            if blk_id < 0 || blk_id >= m.blocks.len {
                continue
            }
            for val_id in m.blocks[blk_id].instrs {
                if val_id < 0 || val_id >= m.values.len {
                    continue
                }
                val := m.values[val_id]
                if val.kind != .instruction {
                    continue
                }
                instr := m.instrs[val.index]
                if instr.op != .phi {
                    continue
                }
                for i := 0; i + 1 < instr.operands.len; i += 2 {
                    val_in := instr.operands[i]
                    pred_blk_idx := int(instr.operands[i + 1])
                    if pred_blk_idx < 0 || pred_blk_idx >= n_blocks {
                        continue
                    }
                    if pred_copy_dests[pred_blk_idx].len == 0 {
                        pred_copy_blocks << pred_blk_idx
                    }
                    array2d_append(mut pred_copy_dests, pred_blk_idx, val_id)
                    array2d_append(mut pred_copy_srcs, pred_blk_idx, val_in)
                }
            }
        }

        for pred_blk in pred_copy_blocks {
            resolve_parallel_copies(mut m, pred_blk, pred_copy_dests[pred_blk],
                pred_copy_srcs[pred_blk])
        }
        for pred_blk in pred_copy_blocks {
            pred_copy_dests[pred_blk] = []
            pred_copy_srcs[pred_blk] = []
        }

        // Remove phi instructions (nop them).
        for blk_id in func.blocks {
            if blk_id < 0 || blk_id >= m.blocks.len {
                continue
            }
            for val_id in m.blocks[blk_id].instrs {
                if val_id < 0 || val_id >= m.values.len {
                    continue
                }
                val := m.values[val_id]
                if val.kind != .instruction {
                    continue
                }
                idx := val.index
                if m.instrs[idx].op == .phi {
                    mut cleanup := m.instrs[idx]
                    cleanup.op = .bitcast
                    cleanup.operands = []
                    m.instrs[idx] = cleanup
                }
            }
        }
    }
}

// resolve_parallel_copies sequences a set of parallel copies dest[i] = src[i]
// using Briggs' worklist algorithm, breaking cycles with a temp.
fn resolve_parallel_copies(mut m ssa.Module, blk_id int, dests []int, srcs []int) {
    if dests.len == 0 {
        return
    }
    mut p_dest := []int{}
    mut p_src := []int{}
    for ci in 0 .. dests.len {
        if dests[ci] != srcs[ci] {
            p_dest << dests[ci]
            p_src << srcs[ci]
        }
    }
    np := p_dest.len
    if np == 0 {
        return
    }

    mut max_id := m.values.len
    for i in 0 .. np {
        if p_dest[i] >= max_id {
            max_id = p_dest[i] + 1
        }
        if p_src[i] >= max_id {
            max_id = p_src[i] + 1
        }
    }
    mut src_ref_count := []int{len: max_id + np + 4}

    for i in 0 .. np {
        src_ref_count[p_src[i]]++
    }

    mut alive := []int{len: np}
    for i in 0 .. np {
        alive[i] = 1
    }

    mut worklist := []int{}
    for i in 0 .. np {
        if src_ref_count[p_dest[i]] == 0 {
            worklist << i
        }
    }

    mut s_dest := []int{}
    mut s_src := []int{}
    mut remaining := np

    for remaining > 0 {
        if worklist.len > 0 {
            idx := worklist.pop()
            if idx < 0 || idx >= np || alive[idx] == 0 {
                continue
            }
            alive[idx] = 0
            remaining--
            d := p_dest[idx]
            s := p_src[idx]
            s_dest << d
            s_src << s
            if src_ref_count[s] > 0 {
                src_ref_count[s]--
            }
            if src_ref_count[s] == 0 {
                for j in 0 .. np {
                    if alive[j] == 1 && p_dest[j] == s {
                        worklist << j
                    }
                }
            }
            continue
        }

        // Cycle: break it by copying one source into a temp.
        mut ci := -1
        for k in 0 .. np {
            if alive[k] == 1 {
                ci = k
                break
            }
        }
        if ci < 0 {
            break
        }
        cycle_src := p_src[ci]
        mut typ := 0
        if cycle_src >= 0 && cycle_src < m.values.len {
            typ = m.values[cycle_src].typ
        }
        if typ == 0 && p_dest[ci] >= 0 && p_dest[ci] < m.values.len {
            typ = m.values[p_dest[ci]].typ
        }
        temp := insert_temp_in_block(mut m, blk_id, cycle_src, typ)
        for temp >= src_ref_count.len {
            src_ref_count << 0
        }
        src_ref_count[temp] = 0
        for i in 0 .. np {
            if alive[i] == 1 && p_src[i] == cycle_src {
                p_src[i] = temp
                src_ref_count[temp]++
            }
        }
        src_ref_count[cycle_src] = 0
        for j in 0 .. np {
            if alive[j] == 1 && p_dest[j] == cycle_src {
                worklist << j
            }
        }
    }

    for si in 0 .. s_dest.len {
        insert_copy_in_block(mut m, blk_id, s_dest[si], s_src[si])
    }
}

// insert_temp_in_block updates insert temp in block state for optimize.
fn insert_temp_in_block(mut m ssa.Module, blk_id int, src int, typ int) int {
    m.instrs << ssa.Instruction{
        op:       .bitcast
        block:    blk_id
        typ:      typ
        operands: [ssa.ValueID(src)]
    }
    temp_id := m.add_value(.instruction, typ, 'phi_tmp_${m.values.len}', m.instrs.len - 1)
    insert_before_terminator(mut m, blk_id, temp_id)
    return temp_id
}

// insert_copy_in_block updates insert copy in block state for optimize.
fn insert_copy_in_block(mut m ssa.Module, blk_id int, dest int, src int) {
    typ := m.values[dest].typ
    m.instrs << ssa.Instruction{
        op:       .assign
        block:    blk_id
        typ:      typ
        operands: [ssa.ValueID(dest), ssa.ValueID(src)]
    }
    val_id := m.add_value(.instruction, typ, 'copy', m.instrs.len - 1)
    insert_before_terminator(mut m, blk_id, val_id)
}

// insert_before_terminator updates insert before terminator state for optimize.
fn insert_before_terminator(mut m ssa.Module, blk_id int, new_val int) {
    mut blk := m.blocks[blk_id]
    mut insert_idx := blk.instrs.len
    if insert_idx > 0 {
        last_val_id := blk.instrs[blk.instrs.len - 1]
        last_val := m.values[last_val_id]
        last_instr := m.instrs[last_val.index]
        if last_instr.op in [.ret, .br, .jmp, .switch_, .unreachable] {
            insert_idx = blk.instrs.len - 1
        }
    }
    mut new_instrs := []ssa.ValueID{cap: blk.instrs.len + 1}
    for ii in 0 .. insert_idx {
        new_instrs << blk.instrs[ii]
    }
    new_instrs << ssa.ValueID(new_val)
    for ii in insert_idx .. blk.instrs.len {
        new_instrs << blk.instrs[ii]
    }
    blk.instrs = new_instrs
    m.blocks[blk_id] = blk
}