// Copyright (c) 2026 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

module optimize

import v2.ssa

// --- Simplify Phi Nodes ---
// Remove trivial phi nodes where all operands are the same or self-referential.
// A phi is trivial if all its non-self operands resolve to the same value.
// This reduces unnecessary instructions before phi elimination.
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]
            n_func_blocks := func.blocks.len
            for fbi in 0 .. n_func_blocks {
                blk_id := func.blocks[fbi]
                blk := m.blocks[blk_id]
                n_blk_instrs := blk.instrs.len
                for ii in 0 .. n_blk_instrs {
                    val_id := blk.instrs[ii]
                    val := m.values[val_id]
                    if val.kind != .instruction {
                        continue
                    }
                    instr := m.instrs[val.index]
                    if instr.op != .phi {
                        continue
                    }

                    // Dead phi removal: phi with no uses can be removed
                    if val.uses.len == 0 {
                        // Remove uses from this phi's operands
                        for i := 0; i < instr.operands.len; i += 2 {
                            op_val := instr.operands[i]
                            if op_val != val_id { // Don't try to remove self-reference
                                remove_phi_use(mut m, op_val, val_id)
                            }
                        }
                        // Mark phi as dead
                        // Avoid m.instrs[X].field = ... -- chained field assign broken in ARM64 self-hosted
                        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
                    }

                    // Check if phi is trivial (all non-self operands are the same)
                    mut unique_val := -1
                    mut is_trivial := true

                    for i := 0; i < instr.operands.len; i += 2 {
                        op_val := instr.operands[i]
                        // Skip self-references (phi refers to itself)
                        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 trivial and we found a unique value, replace all uses
                    if is_trivial && unique_val != -1 {
                        // Replace all uses of this phi with the unique value
                        m.replace_uses(val_id, unique_val)
                        // Mark phi as dead (will be cleaned up by DCE or ignored)
                        // Avoid m.instrs[X].field = ... -- chained field assign broken in ARM64 self-hosted
                        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
}

// Helper to remove a use from a value's uses list (for dead phi removal)
fn remove_phi_use(mut m ssa.Module, val_id int, user_id int) {
    if val_id >= m.values.len {
        return
    }
    mut val := &m.values[val_id]
    for i := val.uses.len - 1; i >= 0; i-- {
        if val.uses[i] == user_id {
            val.uses.delete(i)
        }
    }
}

// get_block_val_id returns the value ID for a block, using a pre-built lookup table.
// This avoids m.blocks[bid].val_id which can return wrong values for the large
// BasicBlock struct (168 bytes) in ARM64-compiled binaries.
fn get_block_val_id(_m &ssa.Module, bid int, block_val_ids []int) int {
    if bid >= 0 && bid < block_val_ids.len {
        return block_val_ids[bid]
    }
    return 0
}

// build_block_val_ids scans values for basic_block kind and builds block_id → val_id mapping.
fn build_block_val_ids(m &ssa.Module) []int {
    n_blks := m.blocks.len
    mut block_val_ids := []int{len: n_blks}
    n_vals := m.values.len
    for vi in 0 .. n_vals {
        if m.values[vi].kind == .basic_block {
            bid := m.values[vi].index
            if bid >= 0 && bid < n_blks {
                block_val_ids[bid] = vi
            }
        }
    }
    return block_val_ids
}

// --- Critical Edge Splitting ---
fn split_critical_edges(mut m ssa.Module) {
    mut cfg := build_cfg(mut m)

    for fi in 0 .. m.funcs.len {
        mut new_blocks := []ssa.BlockID{}

        // Collect edges to split (can't modify while iterating)
        mut edges_to_split := [][]ssa.BlockID{} // [pred_id, succ_id]

        // Find all critical edges
        func := m.funcs[fi]
        func_blocks_len := func.blocks.len
        for bi in 0 .. func_blocks_len {
            blk_id := func.blocks[bi]
            if blk_id < 0 || blk_id >= m.blocks.len {
                continue
            }
            n_blk_succs := cfg.succs[blk_id].len
            if n_blk_succs > 1 {
                for si in 0 .. n_blk_succs {
                    succ_id := cfg.succs[blk_id][si]
                    if succ_id < 0 || succ_id >= m.blocks.len {
                        continue
                    }
                    if cfg.preds[succ_id].len > 1 {
                        edges_to_split << [blk_id, succ_id]
                    }
                }
            }
        }

        // Split each critical edge
        n_edges := edges_to_split.len
        for ei in 0 .. n_edges {
            pred_id := edges_to_split[ei][0]
            succ_id := edges_to_split[ei][1]

            if pred_id < 0 || pred_id >= m.blocks.len || succ_id < 0 || succ_id >= m.blocks.len {
                continue
            }

            // Create new intermediate block
            func2 := m.funcs[fi]
            split_blk := m.add_block(func2.id, 'split_${pred_id}_${succ_id}')
            new_blocks << split_blk

            // Re-build block_val_ids after adding the split block
            // (new blocks get new val_ids that need to be in the table)
            block_val_ids := build_block_val_ids(m)

            succ_val := get_block_val_id(m, succ_id, block_val_ids)
            // Add unconditional jump from split block to original successor
            m.add_instr(.jmp, split_blk, 0, [succ_val])

            // Update predecessor's terminator to jump to split block instead of successor
            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 {
                    term_val := m.values[term_val_id]
                    mut term := &m.instrs[term_val.index]

                    old_succ_val := succ_val
                    new_succ_val := get_block_val_id(m, split_blk, block_val_ids)

                    // Replace ALL occurrences (handles switch with duplicate targets)
                    for i in 0 .. term.operands.len {
                        if term.operands[i] == old_succ_val {
                            term.operands[i] = new_succ_val
                        }
                    }
                }
            }

            // Update phi nodes in successor to reference split block instead of pred
            succ_blk2 := m.blocks[succ_id]
            n_succ_instrs := succ_blk2.instrs.len
            for vi in 0 .. n_succ_instrs {
                val_id := succ_blk2.instrs[vi]
                if val_id < 0 || val_id >= m.values.len {
                    continue
                }
                val := m.values[val_id]
                if val.kind != .instruction {
                    continue
                }
                idx := val.index
                if idx < 0 || idx >= m.instrs.len {
                    continue
                }
                mut instr := &m.instrs[idx]
                if instr.op == .phi {
                    old_pred_val := get_block_val_id(m, pred_id, block_val_ids)
                    new_pred_val := get_block_val_id(m, split_blk, block_val_ids)
                    // Replace all occurrences (defensive - handles edge cases)
                    for i := 1; i < instr.operands.len; i += 2 {
                        if instr.operands[i] == old_pred_val {
                            instr.operands[i] = new_pred_val
                        }
                    }
                }
            }
        }
    }

    // Rebuild CFG after splitting
    cfg = build_cfg(mut m)
    _ = cfg
}

fn eliminate_phi_nodes(mut m ssa.Module) {
    // First split critical edges to ensure correct copy placement
    split_critical_edges(mut m)

    n_blocks := m.blocks.len
    n_funcs := m.funcs.len

    // Build reverse map: val_id → block_id for block-kind values.
    // Scan values array for basic_block kind instead of using m.blocks[bid].val_id
    // which returns wrong results in ARM64-compiled binaries (large struct field access bug).
    mut val_to_block := []int{len: m.values.len, init: -1}
    n_vals := m.values.len
    for vi in 0 .. n_vals {
        if m.values[vi].kind == .basic_block {
            bid := m.values[vi].index
            if bid >= 0 && bid < n_blocks {
                val_to_block[vi] = bid
            }
        }
    }

    // Use flat parallel arrays indexed by block_id instead of [][]ParallelCopy.
    // pred_copy_dests[blk_id] and pred_copy_srcs[blk_id] store the dest/src pairs.
    mut pred_copy_dests := [][]int{len: n_blocks}
    mut pred_copy_srcs := [][]int{len: n_blocks}
    mut pred_copy_blocks := []int{}

    // Pre-allocate shared src_ref_count array for resolve_parallel_copies_flat.
    // Sized to m.values.len + headroom for temps. Reused across all calls.
    mut src_ref_count := []int{len: m.values.len + 1024}
    mut touched_ids := []int{cap: 256}

    for fi in 0 .. n_funcs {
        func := m.funcs[fi]
        func_blocks_len2 := func.blocks.len
        for bi2 in 0 .. func_blocks_len2 {
            blk_id := func.blocks[bi2]
            if blk_id < 0 || blk_id >= m.blocks.len {
                continue
            }
            blk := m.blocks[blk_id]
            n_instrs := blk.instrs.len
            for ii in 0 .. n_instrs {
                val_id := blk.instrs[ii]
                if val_id < 0 || val_id >= m.values.len {
                    continue
                }
                val := m.values[val_id]
                if val.kind != .instruction {
                    continue
                }
                idx := val.index
                if idx < 0 || idx >= m.instrs.len {
                    continue
                }
                instr := m.instrs[idx]
                if instr.op == .phi {
                    is_do_loop := func.name == 'do_loop'
                    if is_do_loop {
                        eprintln('  do_loop phi: val=${val_id} blk=${blk_id} nops=${instr.operands.len}')
                    }
                    // Phi operands: [val0, blk0, val1, blk1, ...]
                    for i := 0; i < instr.operands.len; i += 2 {
                        if i + 1 >= instr.operands.len {
                            break
                        }
                        val_in := instr.operands[i]
                        blk_val := instr.operands[i + 1]
                        if blk_val < 0 || blk_val >= m.values.len {
                            if is_do_loop {
                                eprintln('    op[${i}]: val_in=${val_in} blk_val=${blk_val} SKIP(oob)')
                            }
                            continue
                        }
                        // Use val_to_block reverse map instead of blk_v.index
                        pred_blk_idx := val_to_block[blk_val]
                        if is_do_loop {
                            blk_kind := m.values[blk_val].kind
                            blk_index := m.values[blk_val].index
                            eprintln('    op[${i}]: val_in=${val_in} blk_val=${blk_val} kind=${blk_kind} index=${blk_index} pred_blk_idx=${pred_blk_idx}')
                        }

                        if pred_blk_idx >= 0 && pred_blk_idx < n_blocks {
                            if pred_copy_dests[pred_blk_idx].len == 0 {
                                pred_copy_blocks << pred_blk_idx
                            }
                            mut pcd := pred_copy_dests[pred_blk_idx]
                            pcd << val_id
                            pred_copy_dests[pred_blk_idx] = pcd
                            mut pcs := pred_copy_srcs[pred_blk_idx]
                            pcs << val_in
                            pred_copy_srcs[pred_blk_idx] = pcs
                        }
                    }
                }
            }
        }

        // For each predecessor that has copies, resolve with Briggs algorithm
        for pki in 0 .. pred_copy_blocks.len {
            pred_blk := pred_copy_blocks[pki]
            resolve_parallel_copies_flat(mut m, pred_blk, pred_copy_dests[pred_blk],
                pred_copy_srcs[pred_blk], mut src_ref_count, mut touched_ids)
        }

        // Cleanup: clear pred_copies for blocks we touched
        for pki2 in 0 .. pred_copy_blocks.len {
            pred_copy_dests[pred_copy_blocks[pki2]] = []
            pred_copy_srcs[pred_copy_blocks[pki2]] = []
        }
        pred_copy_blocks.clear()

        // Remove phi instructions (mark as nop/bitcast with no operands)
        func2 := m.funcs[fi]
        func_blocks_len3 := func2.blocks.len
        for bi3 in 0 .. func_blocks_len3 {
            blk_id := func2.blocks[bi3]
            if blk_id < 0 || blk_id >= m.blocks.len {
                continue
            }
            blk := m.blocks[blk_id]
            n_instrs3 := blk.instrs.len
            for ii3 in 0 .. n_instrs3 {
                val_id := blk.instrs[ii3]
                if val_id < 0 || val_id >= m.values.len {
                    continue
                }
                val := m.values[val_id]
                if val.kind != .instruction {
                    continue
                }
                idx := val.index
                if idx < 0 || idx >= m.instrs.len {
                    continue
                }
                if m.instrs[idx].op == .phi {
                    // Avoid m.instrs[X].field = ... -- chained field assign broken in ARM64 self-hosted
                    mut cleanup_instr := m.instrs[idx]
                    cleanup_instr.op = .bitcast
                    cleanup_instr.operands = []
                    m.instrs[idx] = cleanup_instr
                }
            }
        }
    }
}

// resolve_parallel_copies_flat sequences a set of parallel copies (dest[i] = src[i])
// using the Briggs algorithm with worklist-based scheduling.
// A copy dest←src is "ready" when no other pending copy reads from dest.
// Uses caller-provided flat array src_ref_count[] indexed by value ID (no maps).
// touched_ids tracks which entries were modified so we can zero only those on cleanup.
fn resolve_parallel_copies_flat(mut m ssa.Module, blk_id int, dests []int, srcs []int, mut src_ref_count []int, mut touched_ids []int) {
    if dests.len == 0 {
        return
    }

    n := dests.len

    // Build pending copies, filtering self-copies
    mut p_dest := []int{cap: n}
    mut p_src := []int{cap: n}
    for ci in 0 .. n {
        if dests[ci] != srcs[ci] {
            p_dest << dests[ci]
            p_src << srcs[ci]
        }
    }

    np := p_dest.len
    if np == 0 {
        return
    }

    // Ensure src_ref_count is large enough for all value IDs we'll encounter.
    // Check current max needed against array length.
    mut need_len := m.values.len + np + 4
    for i in 0 .. np {
        if p_dest[i] >= need_len {
            need_len = p_dest[i] + np + 4
        }
        if p_src[i] >= need_len {
            need_len = p_src[i] + np + 4
        }
    }
    if need_len > src_ref_count.len {
        // Grow array
        for _ in 0 .. need_len - src_ref_count.len {
            src_ref_count << 0
        }
    }

    // src_ref_count[val] = how many pending copies use val as source
    // (array is pre-zeroed or cleaned from previous call via touched_ids)
    touched_ids.clear()
    for i in 0 .. np {
        s := p_src[i]
        if src_ref_count[s] == 0 {
            touched_ids << s
        }
        src_ref_count[s] = src_ref_count[s] + 1
    }
    // Also track dests we'll read from src_ref_count
    for i in 0 .. np {
        d := p_dest[i]
        if src_ref_count[d] == 0 {
            // dest not in touched_ids yet but we read it; track for safety
            // (it's already 0, but if cycle-breaking modifies it we need to clean it)
        }
    }

    // Use int instead of bool for alive - ARM64 codegen may mishandle []bool{init: true}
    // alive[i] = 1 if copy i is still pending, 0 if done
    mut alive := []int{len: np}
    for ai in 0 .. np {
        alive[ai] = 1
    }

    // Worklist: use manual stack pointer instead of .delete() which may be buggy in ARM64
    mut worklist := []int{len: np + np + 4}
    mut wl_top := 0 // stack pointer: worklist[0..wl_top] are valid entries
    for i in 0 .. np {
        d := p_dest[i]
        if src_ref_count[d] == 0 {
            worklist[wl_top] = i
            wl_top += 1
        }
    }

    // Sequenced output copies
    mut s_dest := []int{cap: np + 4}
    mut s_src := []int{cap: np + 4}

    mut remaining := np
    for remaining > 0 {
        if wl_top > 0 {
            // Pop a ready copy from the worklist
            wl_top -= 1
            idx := worklist[wl_top]
            if idx < 0 || idx >= np || alive[idx] == 0 {
                continue
            }
            alive[idx] = 0
            remaining -= 1

            d := p_dest[idx]
            s := p_src[idx]
            s_dest << d
            s_src << s

            // Decrement ref count for the source
            src_ref_count[s] = src_ref_count[s] - 1
            if src_ref_count[s] == 0 {
                // s is no longer used as source; any pending copy with dest==s is now ready
                for j in 0 .. np {
                    if alive[j] == 1 && p_dest[j] == s {
                        // Grow worklist if needed
                        if wl_top >= worklist.len {
                            worklist << 0
                        }
                        worklist[wl_top] = j
                        wl_top += 1
                    }
                }
            }
            continue
        }

        // Worklist empty but copies remain → cycle. Break it with 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)

        // Grow src_ref_count if temp exceeds array bounds
        if temp >= src_ref_count.len {
            for _ in 0 .. temp + 1 - src_ref_count.len {
                src_ref_count << 0
            }
        }

        // Replace all references to cycle_src with temp in pending copies
        src_ref_count[temp] = 0
        touched_ids << temp
        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[temp] + 1
            }
        }
        src_ref_count[cycle_src] = 0

        // Check if any copy whose dest was cycle_src is now ready
        for j in 0 .. np {
            if alive[j] == 1 && p_dest[j] == cycle_src {
                if wl_top >= worklist.len {
                    worklist << 0
                }
                worklist[wl_top] = j
                wl_top += 1
            }
        }
    }

    // Clean up src_ref_count: zero only the entries we touched
    for ti in 0 .. touched_ids.len {
        tid := touched_ids[ti]
        if tid < src_ref_count.len {
            src_ref_count[tid] = 0
        }
    }

    // Emit the sequenced copies
    for si in 0 .. s_dest.len {
        insert_copy_in_block(mut m, blk_id, s_dest[si], s_src[si])
    }
}

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_node(.instruction, typ, 'phi_tmp_${m.values.len}', m.instrs.len - 1)

    // Read whole struct, modify, write back (chained field assign broken in ARM64)
    if src < m.values.len && temp_id !in m.values[src].uses {
        mut sv := m.values[src]
        sv.uses << temp_id
        m.values[src] = sv
    }

    // Insert before terminator — build new instrs array
    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 := []int{cap: blk.instrs.len + 1}
    for ii in 0 .. insert_idx {
        new_instrs << blk.instrs[ii]
    }
    new_instrs << temp_id
    for ii in insert_idx .. blk.instrs.len {
        new_instrs << blk.instrs[ii]
    }
    blk.instrs = new_instrs
    m.blocks[blk_id] = blk
    return temp_id
}

fn insert_copy_in_block(mut m ssa.Module, blk_id int, dest int, src int) {
    dest_val := m.values[dest]
    typ := dest_val.typ
    m.instrs << ssa.Instruction{
        op:       .assign
        block:    blk_id
        typ:      typ
        operands: [ssa.ValueID(dest), src]
    }
    val_id := m.add_value_node(.instruction, typ, 'copy', m.instrs.len - 1)
    // Read whole struct, modify, write back (chained field assign broken in ARM64)
    if dest < m.values.len && val_id !in m.values[dest].uses {
        mut dv := m.values[dest]
        dv.uses << val_id
        m.values[dest] = dv
    }
    if src < m.values.len && val_id !in m.values[src].uses {
        mut sv := m.values[src]
        sv.uses << val_id
        m.values[src] = sv
    }

    // Insert before terminator — build new instrs array
    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 := []int{cap: blk.instrs.len + 1}
    for ii in 0 .. insert_idx {
        new_instrs << blk.instrs[ii]
    }
    new_instrs << val_id
    for ii in insert_idx .. blk.instrs.len {
        new_instrs << blk.instrs[ii]
    }
    blk.instrs = new_instrs
    m.blocks[blk_id] = blk
}