// 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

// --- Mem2Reg (Promote Allocas) ---
// All data structures use flat arrays indexed by value_id or block_id
// to avoid map operations which are extremely slow in the ARM64 backend.
struct Mem2RegCtx {
mut:
    // Indexed by alloc_id (value_id): list of block_ids where this alloc is stored to
    defs [][]int
    // Indexed by alloc_id (value_id): list of block_ids where this alloc is loaded from
    uses [][]int
    // Indexed by block_id: list of alloc_ids that need phi nodes in this block
    phi_placements [][]int
    // Indexed by block_id: parallel to phi_placements, stores the phi value_ids
    // phi_vals[blk_id][i] is the phi val_id for phi_placements[blk_id][i]
    phi_vals [][]int
    // Indexed by alloc_id (value_id): stack of current SSA values during renaming
    stacks [][]int
    // Track which alloc_ids have been initialized (for cleanup)
    is_promotable []bool
    // Track which block_ids have phi placements (for cleanup and iteration)
    phi_blocks []int
    // Deferred phi operands: flat list of [instr_idx, val, bvid, instr_idx, val, bvid, ...].
    // Applied after rename to avoid m.instrs[idx].operands << x which is broken
    // in ARM64 self-hosted binaries (chained array-element-field-append bug).
    deferred_phi_ops []int
}

// Helper: append to an inner array of a [][]int with proper write-back.
// Direct `arr[idx] << val` is broken in ARM64 self-hosted binaries
// (chained array-element append copies the struct instead of modifying in-place).
fn array2d_append(mut arr [][]int, idx int, val int) {
    mut inner := arr[idx]
    inner << val
    arr[idx] = inner
}

fn promote_memory_to_register(mut m ssa.Module, dom DomInfo, cfg &CfgData) {
    n_values := m.values.len
    n_blocks := m.blocks.len

    // Pre-allocate flat arrays ONCE outside the function loop.
    mut ctx := Mem2RegCtx{
        defs:           [][]int{len: n_values}
        uses:           [][]int{len: n_values}
        phi_placements: [][]int{len: n_blocks}
        phi_vals:       [][]int{len: n_blocks}
        stacks:         [][]int{len: n_values}
        is_promotable:  []bool{len: n_values}
        phi_blocks:     []int{}
    }

    // Pre-allocate stack_counts array (for rename_recursive).
    mut stack_counts := []int{len: n_values}
    for sci in 0 .. n_values {
        stack_counts[sci] = -1
    }

    // Dominance frontier: indexed by block_id
    mut df := [][]int{len: n_blocks}
    // Track which blocks have DF entries for cleanup
    mut df_blocks := []int{}

    // Per-alloc visited/has_phi bitsets (reused across allocs within a function)
    mut visited := []bool{len: n_blocks}
    mut has_phi := []bool{len: n_blocks}

    mut total_allocas := 0
    mut total_promoted := 0
    mut total_not_promoted := 0
    mut total_blocks_visited := 0
    mut total_blocks_skipped := 0
    mut total_blocks_in_funcs := 0

    for fi in 0 .. m.funcs.len {
        // 1. Analyze Allocas
        mut promotable := []int{}
        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]
                instr := m.instrs[val.index]
                if instr.op == .alloca {
                    total_allocas += 1
                    if is_promotable(m, val_id) {
                        promotable << val_id
                        ctx.is_promotable[val_id] = true
                        total_promoted += 1
                    } else {
                        total_not_promoted += 1
                    }
                }

                if instr.op == .store {
                    ptr := instr.operands[1]
                    // Avoid duplicate def entries for same block
                    if ptr < n_values && blk_id !in ctx.defs[ptr] {
                        array2d_append(mut ctx.defs, ptr, blk_id)
                    }
                } else if instr.op == .load {
                    ptr := instr.operands[0]
                    // Avoid duplicate use entries for same block
                    if ptr < n_values && blk_id !in ctx.uses[ptr] {
                        array2d_append(mut ctx.uses, ptr, blk_id)
                    }
                }
            }
        }

        // 2. Compute Dominance Frontier (flat array version)
        compute_dominance_frontier_flat(m, fi, &dom, cfg, mut df, mut df_blocks)

        // 3. Insert Phis (Dominance Frontier)
        n_promotable := promotable.len
        for pri in 0 .. n_promotable {
            alloc_id := promotable[pri]
            mut worklist := ctx.defs[alloc_id].clone()
            // Track blocks we mark visited/has_phi for cleanup
            mut visited_list := []int{}
            mut has_phi_list := []int{}

            for worklist.len > 0 {
                b := worklist.pop()
                if b < 0 || b >= n_blocks {
                    continue
                }
                n_df_b := df[b].len
                for di in 0 .. n_df_b {
                    d := df[b][di]
                    if !has_phi[d] {
                        array2d_append(mut ctx.phi_placements, d, alloc_id)
                        if d !in ctx.phi_blocks {
                            ctx.phi_blocks << d
                        }
                        has_phi[d] = true
                        has_phi_list << d
                        if !visited[d] {
                            visited[d] = true
                            visited_list << d
                            worklist << d
                        }
                    }
                }
            }

            // Reset visited/has_phi for this alloc
            for vli in 0 .. visited_list.len {
                visited[visited_list[vli]] = false
            }
            for hli in 0 .. has_phi_list.len {
                has_phi[has_phi_list[hli]] = false
            }
        }

        // Insert Phis
        n_phi_blocks := ctx.phi_blocks.len
        for pbi in 0 .. n_phi_blocks {
            blk_id := ctx.phi_blocks[pbi]
            n_phi_allocs := ctx.phi_placements[blk_id].len
            for pai in 0 .. n_phi_allocs {
                alloc_id := ctx.phi_placements[blk_id][pai]
                alloc_val := m.values[alloc_id]
                alloc_typ := m.type_store.types[alloc_val.typ]
                typ := alloc_typ.elem_type
                phi_val := m.add_instr_front(.phi, blk_id, typ, [])
                // Avoid m.values[phi_val].name = ... -- chained field assign broken in ARM64 self-hosted
                mut pv := m.values[phi_val]
                pv.name = '${alloc_val.name}.phi_${blk_id}'
                m.values[phi_val] = pv
                // Store phi_val_id in parallel array for direct lookup (no string matching)
                array2d_append(mut ctx.phi_vals, blk_id, phi_val)
            }
        }

        // 4. Rename Variables
        func2 := m.funcs[fi]
        if func2.blocks.len > 0 {
            entry := func2.blocks[0]
            total_blocks_in_funcs += func2.blocks.len
            bv, bs :=
                rename_recursive(mut m, entry, mut ctx, promotable, mut stack_counts, &dom, cfg)
            total_blocks_visited += bv
            total_blocks_skipped += bs
        }

        // 4b. Apply deferred phi operands.
        // Builds complete operand lists for each phi and assigns them all at once.
        // This avoids m.instrs[idx].operands << x which is broken in v3-compiled binaries
        // (chained array-element-field-append copies the struct instead of modifying in-place).
        if ctx.deferred_phi_ops.len > 0 {
            n_deferred := ctx.deferred_phi_ops.len
            n_phi_blocks3 := ctx.phi_blocks.len
            mut n_applied := 0
            mut n_verified := 0
            mut n_mismatch := 0
            for pbi3 in 0 .. n_phi_blocks3 {
                pb_id := ctx.phi_blocks[pbi3]
                n_phi_allocs3 := ctx.phi_vals[pb_id].len
                for pai3 in 0 .. n_phi_allocs3 {
                    phi_vid := ctx.phi_vals[pb_id][pai3]
                    phi_v := m.values[phi_vid]
                    phi_instr_idx := phi_v.index
                    // Collect all operands for this phi from deferred list
                    mut ops := []int{}
                    mut dj := 0
                    for dj + 2 < n_deferred {
                        if ctx.deferred_phi_ops[dj] == phi_instr_idx {
                            ops << ctx.deferred_phi_ops[dj + 1]
                            ops << ctx.deferred_phi_ops[dj + 2]
                        }
                        dj += 3
                    }
                    if ops.len > 0 {
                        // Avoid m.instrs[X].operands = ops -- chained field assign broken in ARM64 self-hosted
                        mut phi_instr := m.instrs[phi_instr_idx]
                        phi_instr.operands = ops
                        m.instrs[phi_instr_idx] = phi_instr
                        n_applied += 1
                        // Verify the assignment was actually written
                        actual_len := m.instrs[phi_instr_idx].operands.len
                        if actual_len != ops.len {
                            n_mismatch += 1
                            if n_mismatch <= 5 {
                                eprintln('  DEFERRED MISMATCH: phi_vid=${phi_vid} instr=${phi_instr_idx} wanted=${ops.len} got=${actual_len}')
                            }
                        } else {
                            n_verified += 1
                            // Also verify first operand value
                            if ops.len >= 2 {
                                actual0 := m.instrs[phi_instr_idx].operands[0]
                                actual1 := m.instrs[phi_instr_idx].operands[1]
                                if actual0 != ops[0] || actual1 != ops[1] {
                                    n_mismatch += 1
                                    if n_mismatch <= 5 {
                                        eprintln('  DEFERRED VALUE MISMATCH: phi_vid=${phi_vid} instr=${phi_instr_idx} wanted[0]=${ops[0]},${ops[1]} got=${actual0},${actual1}')
                                    }
                                }
                            }
                        }
                    }
                }
            }
            if n_mismatch > 0 {
                eprintln('  DEFERRED phi ops: applied=${n_applied} verified=${n_verified} MISMATCH=${n_mismatch}')
            }
            ctx.deferred_phi_ops = []
        }
        // Update uses for phi operands (deferred from step 3)
        n_phi_blocks4 := ctx.phi_blocks.len
        for pbi4 in 0 .. n_phi_blocks4 {
            pb_id := ctx.phi_blocks[pbi4]
            n_phi_allocs4 := ctx.phi_vals[pb_id].len
            for pai4 in 0 .. n_phi_allocs4 {
                phi_vid := ctx.phi_vals[pb_id][pai4]
                phi_v := m.values[phi_vid]
                phi_instr_idx := phi_v.index
                instr := m.instrs[phi_instr_idx]
                for oi := 0; oi < instr.operands.len; oi += 2 {
                    val_op := instr.operands[oi]
                    if val_op < m.values.len && phi_vid !in m.values[val_op].uses {
                        // Read whole struct, modify, write back (chained broken in ARM64)
                        mut vop := m.values[val_op]
                        vop.uses << phi_vid
                        m.values[val_op] = vop
                    }
                }
            }
        }

        // 5. Cleanup for next function: reset only entries we touched
        for cli in 0 .. promotable.len {
            alloc_id := promotable[cli]
            ctx.defs[alloc_id] = []
            ctx.uses[alloc_id] = []
            ctx.stacks[alloc_id] = []
            ctx.is_promotable[alloc_id] = false
        }
        n_phi_blocks2 := ctx.phi_blocks.len
        for cli2 in 0 .. n_phi_blocks2 {
            ctx.phi_placements[ctx.phi_blocks[cli2]] = []
            ctx.phi_vals[ctx.phi_blocks[cli2]] = []
        }
        ctx.phi_blocks = []
        n_df_blocks := df_blocks.len
        for cli3 in 0 .. n_df_blocks {
            df[df_blocks[cli3]] = []
        }
        df_blocks = []
    }
    eprintln('  mem2reg stats: allocas=${total_allocas} promoted=${total_promoted} not_promoted=${total_not_promoted} blocks_total=${total_blocks_in_funcs} blocks_visited=${total_blocks_visited} blocks_skipped=${total_blocks_skipped}')
}

fn is_promotable(m &ssa.Module, alloc_id int) bool {
    // Keep array-backed slots in memory. Promoting pointer-to-array allocas can
    // lose correct addressing semantics for fixed-array literals/indexing.
    if alloc_id > 0 && alloc_id < m.values.len {
        alloc_val := m.values[alloc_id]
        alloc_typ_id := alloc_val.typ
        if alloc_typ_id > 0 && alloc_typ_id < m.type_store.types.len {
            alloc_typ := m.type_store.types[alloc_typ_id]
            if alloc_typ.kind == .ptr_t {
                elem_typ_id := alloc_typ.elem_type
                if elem_typ_id > 0 && elem_typ_id < m.type_store.types.len {
                    elem_typ := m.type_store.types[elem_typ_id]
                    if elem_typ.kind == .ptr_t {
                        return false
                    }
                    if elem_typ.kind == .array_t {
                        return false
                    }
                }
            }
        }
    }

    alloc_val2 := m.values[alloc_id]
    uses := alloc_val2.uses
    n_uses := uses.len
    for ui in 0 .. n_uses {
        u := uses[ui]
        if u >= m.values.len {
            continue
        }
        user := m.values[u]
        if user.kind != .instruction {
            return false
        }
        instr := m.instrs[user.index]
        match instr.op {
            .load {
                if instr.operands.len == 0 || instr.operands[0] != alloc_id {
                    return false
                }
            }
            .store {
                // Only safe if used as pointer (index 1)
                if instr.operands.len < 2 || instr.operands[1] != alloc_id {
                    return false
                }
            }
            else {
                // Escape (GEP, Call, Phi, etc.)
                return false
            }
        }
    }
    return true
}

fn compute_dominance_frontier_flat(m &ssa.Module, func_idx int, dom &DomInfo, cfg &CfgData, mut df [][]int, mut df_blocks []int) {
    func := m.funcs[func_idx]
    n_func_blocks := func.blocks.len
    for bi in 0 .. n_func_blocks {
        blk_id := func.blocks[bi]
        if blk_id < 0 || blk_id >= m.blocks.len {
            continue
        }
        num_preds := cfg.preds[blk_id].len
        if num_preds >= 2 {
            for pi in 0 .. num_preds {
                mut runner := cfg.preds[blk_id][pi]
                idom := dom.idom[blk_id]
                // Safety check: idom != -1
                for runner != -1 && runner != idom {
                    if runner < 0 || runner >= m.blocks.len {
                        break
                    }
                    // Avoid duplicate entries in dominance frontier
                    if blk_id !in df[runner] {
                        array2d_append(mut df, runner, blk_id)
                        if runner !in df_blocks {
                            df_blocks << runner
                        }
                    }
                    if runner == dom.idom[runner] {
                        break
                    }
                    runner = dom.idom[runner]
                }
            }
        }
    }
}

// RenameFrame holds the state for one level of the iterative dom tree walk.
struct RenameFrame {
mut:
    blk_id        int
    child_idx     int   // next child to process in dom_tree
    pushed_allocs []int // alloc_ids pushed in this block (for cleanup)
    processed     bool  // whether steps 1-3 have been run for this block
}

fn rename_iterative(mut m ssa.Module, root_blk int, mut ctx Mem2RegCtx, _promotable []int, mut _stack_counts []int, dom &DomInfo, cfg &CfgData) (int, int) {
    // Pre-build block_val_ids[] by scanning values for basic_block kind.
    // This avoids m.blocks[blk_id].val_id which produces wrong results
    // in ARM64-compiled binaries (large struct field access bug).
    n_blks := m.blocks.len
    mut block_val_ids := []int{len: n_blks}
    n_vals := m.values.len
    mut n_bb_found := 0
    mut n_bb_valid := 0
    for vi in 0 .. n_vals {
        if m.values[vi].kind == .basic_block {
            bid := m.values[vi].index
            n_bb_found += 1
            if bid >= 0 && bid < n_blks {
                block_val_ids[bid] = vi
                n_bb_valid += 1
            }
        }
    }
    // Verify: also check using m.blocks[bid].val_id approach and compare
    mut n_match := 0
    mut n_mismatch := 0
    for bid in 0 .. n_blks {
        bval := m.blocks[bid].val_id
        if block_val_ids[bid] != bval {
            n_mismatch += 1
            if n_mismatch <= 3 {
                eprintln('  mem2reg block_val_ids MISMATCH: bid=${bid} scan=${block_val_ids[bid]} blocks=${bval}')
            }
        } else {
            n_match += 1
        }
    }
    if n_mismatch > 0 {
        eprintln('  mem2reg block_val_ids: bb_found=${n_bb_found} valid=${n_bb_valid} match=${n_match} mismatch=${n_mismatch}')
    }

    mut work := []RenameFrame{}
    work << RenameFrame{
        blk_id: root_blk
    }
    mut visited := []bool{len: m.blocks.len}
    mut blocks_visited := 0
    mut blocks_skipped := 0

    for work.len > 0 {
        fi := work.len - 1
        blk_id := work[fi].blk_id

        if !work[fi].processed {
            // Cycle detection: skip blocks already visited in this traversal
            if blk_id >= 0 && blk_id < visited.len && visited[blk_id] {
                blocks_skipped += 1
                work.pop()
                continue
            }
            if blk_id >= 0 && blk_id < visited.len {
                visited[blk_id] = true
            }
            blocks_visited += 1
            // Steps 1-3: process this block (runs once per block)
            // Avoid work[fi].processed = true — chained field assignment broken in ARM64 self-hosted.
            mut frame2 := work[fi]
            frame2.processed = true
            work[fi] = frame2

            // 1. Push Phis to stack (use phi_vals[] for direct lookup, no string matching)
            if blk_id < ctx.phi_placements.len && ctx.phi_placements[blk_id].len > 0 {
                n_phi_allocs := ctx.phi_placements[blk_id].len
                for pai in 0 .. n_phi_allocs {
                    alloc_id := ctx.phi_placements[blk_id][pai]
                    // Direct lookup: phi_vals[blk_id][pai] is the phi for this alloc
                    if blk_id < ctx.phi_vals.len && pai < ctx.phi_vals[blk_id].len {
                        phi_val_id := ctx.phi_vals[blk_id][pai]
                        // Avoid ctx.stacks[X] << Y — chained append broken in ARM64 self-hosted.
                        mut new_stack := ctx.stacks[alloc_id].clone()
                        new_stack << phi_val_id
                        ctx.stacks[alloc_id] = new_stack
                        // Avoid work[fi].pushed_allocs = X — chained field assignment broken in ARM64.
                        mut wf := work[fi]
                        wf.pushed_allocs << alloc_id
                        work[fi] = wf
                    }
                }
            }

            // 2. Process Instructions
            mut instrs_to_nop := []int{}

            blk2 := m.blocks[blk_id]
            n_instrs2 := blk2.instrs.len
            for ii in 0 .. n_instrs2 {
                val_id := blk2.instrs[ii]
                val := m.values[val_id]
                instr := m.instrs[val.index]
                if instr.op == .store {
                    ptr := instr.operands[1]
                    if ptr < ctx.is_promotable.len && ctx.is_promotable[ptr] {
                        // Avoid ctx.stacks[X] << Y — chained append broken in ARM64 self-hosted.
                        mut new_stack := ctx.stacks[ptr].clone()
                        new_stack << instr.operands[0]
                        ctx.stacks[ptr] = new_stack
                        // Avoid work[fi].pushed_allocs = X — chained field assignment broken in ARM64.
                        mut wf := work[fi]
                        wf.pushed_allocs << ptr
                        work[fi] = wf
                        instrs_to_nop << val_id
                    }
                } else if instr.op == .load {
                    ptr := instr.operands[0]
                    if ptr < ctx.is_promotable.len && ctx.is_promotable[ptr] {
                        stack := ctx.stacks[ptr]
                        mut repl := 0
                        if stack.len > 0 {
                            repl = stack.last()
                        } else {
                            // Undef - reading uninitialized memory
                            res_type := val.typ
                            repl = m.get_or_add_const(res_type, 'undef')
                        }
                        m.replace_uses(val_id, repl)
                        instrs_to_nop << val_id
                    }
                } else if instr.op == .alloca {
                    if val_id < ctx.is_promotable.len && ctx.is_promotable[val_id] {
                        instrs_to_nop << val_id
                    }
                }
            }

            for nopi in 0 .. instrs_to_nop.len {
                vid := instrs_to_nop[nopi]
                vid_val := m.values[vid]
                // Avoid m.instrs[X].op/operands = ... -- chained field assign broken in ARM64 self-hosted
                mut nop_instr := m.instrs[vid_val.index]
                nop_instr.op = .bitcast
                nop_instr.operands = []
                m.instrs[vid_val.index] = nop_instr
            }

            // 3. Update Successor Phi Operands
            n_succs := cfg.succs[blk_id].len
            for si in 0 .. n_succs {
                succ_id := cfg.succs[blk_id][si]
                if succ_id < ctx.phi_placements.len && ctx.phi_placements[succ_id].len > 0 {
                    n_succ_phi_allocs := ctx.phi_placements[succ_id].len
                    for spai in 0 .. n_succ_phi_allocs {
                        alloc_id := ctx.phi_placements[succ_id][spai]
                        // Direct lookup: phi_vals[succ_id][spai] is the phi for this alloc
                        if succ_id < ctx.phi_vals.len && spai < ctx.phi_vals[succ_id].len {
                            vid := ctx.phi_vals[succ_id][spai]
                            phi_v := m.values[vid]
                            mut phi_val := if ctx.stacks[alloc_id].len > 0 {
                                ctx.stacks[alloc_id].last()
                            } else {
                                // Undef - reading uninitialized memory
                                alloc_v := m.values[alloc_id]
                                alloc_typ := m.type_store.types[alloc_v.typ]
                                m.get_or_add_const(alloc_typ.elem_type, 'undef')
                            }
                            bvid := block_val_ids[blk_id]
                            // Defer phi operand append to avoid m.instrs[idx].operands << x
                            // which is broken in ARM64 self-hosted binaries.
                            ctx.deferred_phi_ops << phi_v.index
                            ctx.deferred_phi_ops << phi_val
                            ctx.deferred_phi_ops << bvid
                            // FIX: Update uses so DCE doesn't remove the value
                            // Avoid m.values[X].uses << Y — chained append broken in ARM64 self-hosted.
                            if phi_val < m.values.len {
                                if vid !in m.values[phi_val].uses {
                                    // Read whole struct, modify, write back (chained broken in ARM64)
                                    mut vv := m.values[phi_val]
                                    vv.uses << vid
                                    m.values[phi_val] = vv
                                }
                            }
                        }
                    }
                }
            }
        }

        // 4. Push next unvisited dom child
        n_children := dom.dom_tree[blk_id].len
        child_idx := work[fi].child_idx
        if child_idx < n_children {
            child := dom.dom_tree[blk_id][child_idx]
            // Avoid work[fi].child_idx++ — chained field increment broken in ARM64 self-hosted.
            mut frame := work[fi]
            frame.child_idx++
            work[fi] = frame
            work << RenameFrame{
                blk_id: child
            }
        } else {
            // 5. All children processed - pop stacks (cleanup)
            pushed := work[fi].pushed_allocs.clone()
            work.pop()
            for i := pushed.len - 1; i >= 0; i-- {
                // Avoid ctx.stacks[X].pop() — chained method broken in ARM64 self-hosted.
                mut s := ctx.stacks[pushed[i]].clone()
                s.pop()
                ctx.stacks[pushed[i]] = s
            }
        }
    }
    return blocks_visited, blocks_skipped
}

fn rename_recursive(mut m ssa.Module, blk_id int, mut ctx Mem2RegCtx, promotable []int, mut stack_counts []int, dom &DomInfo, cfg &CfgData) (int, int) {
    return rename_iterative(mut m, blk_id, mut ctx, promotable, mut stack_counts, dom, cfg)
}