module optimize

import v3.ssa

// --- Mem2Reg (promote allocas to SSA values + phi nodes) ---
// Flat arrays indexed by value_id / block_id. Phi predecessor operands store
// raw block ids, matching the rest of the v3 SSA (and the arm64 backend's
// emit_phi_edge_copies, which reads operand[odd] as a block id).
struct Mem2RegCtx {
mut:
    defs           [][]int // alloc_id -> blocks storing to it
    uses           [][]int // alloc_id -> blocks loading from it
    phi_placements [][]int // block_id -> alloc_ids needing a phi here
    phi_vals       [][]int // block_id -> parallel phi value ids
    stacks         [][]int // alloc_id -> reaching-definition stack
    is_promotable  []bool
    phi_blocks     []int
}

// array2d_append supports array2d append handling for optimize.
fn array2d_append(mut arr [][]int, idx int, val int) {
    mut inner := arr[idx]
    inner << val
    arr[idx] = inner
}

// promote_memory_to_register promotes scalar allocas to SSA registers, inserting
// phi nodes at dominance frontiers and renaming load/store chains.
fn promote_memory_to_register(mut m ssa.Module, dom DomInfo, cfg &CfgData) {
    n_values := m.values.len
    n_blocks := m.blocks.len

    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{}
    }

    mut df := [][]int{len: n_blocks}
    mut df_blocks := []int{}
    mut visited := []bool{len: n_blocks}
    mut has_phi := []bool{len: n_blocks}

    for fi in 0 .. m.funcs.len {
        mut promotable := []int{}
        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 == .alloca {
                    if is_promotable(m, val_id) {
                        promotable << val_id
                        ctx.is_promotable[val_id] = true
                    }
                }
                if instr.op == .store {
                    ptr := instr.operands[1]
                    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]
                    if ptr < n_values && blk_id !in ctx.uses[ptr] {
                        array2d_append(mut ctx.uses, ptr, blk_id)
                    }
                }
            }
        }

        compute_dominance_frontier_flat(m, fi, &dom, cfg, mut df, mut df_blocks)

        // Insert phis at iterated dominance frontiers of each promotable alloca.
        for alloc_id in promotable {
            mut worklist := ctx.defs[alloc_id].clone()
            mut visited_list := []int{}
            mut has_phi_list := []int{}
            for worklist.len > 0 {
                b := worklist.pop()
                if b < 0 || b >= n_blocks {
                    continue
                }
                for d in df[b] {
                    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
                        }
                    }
                }
            }
            for d in visited_list {
                visited[d] = false
            }
            for d in has_phi_list {
                has_phi[d] = false
            }
        }

        // Materialize the phi instructions (empty operands, filled during rename).
        for blk_id in ctx.phi_blocks {
            for alloc_id in ctx.phi_placements[blk_id] {
                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, [])
                mut pv := m.values[phi_val]
                pv.name = '${alloc_val.name}.phi_${blk_id}'
                m.values[phi_val] = pv
                array2d_append(mut ctx.phi_vals, blk_id, phi_val)
            }
        }

        // Rename: walk the dominator tree, threading reaching definitions.
        if func.blocks.len > 0 {
            rename_blocks(mut m, func.blocks[0], mut ctx, dom, cfg)
        }

        // Reset per-function state.
        for alloc_id in promotable {
            ctx.defs[alloc_id] = []
            ctx.uses[alloc_id] = []
            ctx.stacks[alloc_id] = []
            ctx.is_promotable[alloc_id] = false
        }
        for pb in ctx.phi_blocks {
            ctx.phi_placements[pb] = []
            ctx.phi_vals[pb] = []
        }
        ctx.phi_blocks = []
        for d in df_blocks {
            df[d] = []
        }
        df_blocks = []
    }
}

// is_promotable returns true if every use of the alloca is a direct load/store,
// i.e. the pointer never escapes. Array- and pointer-backed slots stay in memory.
fn is_promotable(m &ssa.Module, alloc_id int) bool {
    if alloc_id > 0 && alloc_id < m.values.len {
        alloc_typ_id := m.values[alloc_id].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]
                    // Keep pointer-, array- and aggregate-backed slots in memory.
                    // Promoting them to SSA values would require the backend to
                    // thread by-value aggregates through phi copies, which the v3
                    // arm64 lowering does not yet handle reliably.
                    if elem_typ.kind in [.ptr_t, .array_t, .struct_t] {
                        return false
                    }
                }
            }
        }
    }

    for u in m.values[alloc_id].uses {
        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 {
                if instr.operands.len < 2 || instr.operands[1] != alloc_id {
                    return false
                }
            }
            else {
                return false
            }
        }
    }
    return true
}

// compute_dominance_frontier_flat supports compute dominance frontier flat handling for optimize.
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]
    for blk_id in func.blocks {
        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]
                for runner != -1 && runner != idom {
                    if runner < 0 || runner >= m.blocks.len {
                        break
                    }
                    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 represents rename frame data used by optimize.
struct RenameFrame {
mut:
    blk_id        int
    child_idx     int
    pushed_allocs []int
    processed     bool
}

// rename_blocks supports rename blocks handling for optimize.
fn rename_blocks(mut m ssa.Module, root_blk int, mut ctx Mem2RegCtx, dom DomInfo, cfg &CfgData) {
    mut work := []RenameFrame{}
    work << RenameFrame{
        blk_id: root_blk
    }
    mut visited := []bool{len: m.blocks.len}

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

        if !work[fi].processed {
            if blk_id >= 0 && blk_id < visited.len && visited[blk_id] {
                work.pop()
                continue
            }
            if blk_id >= 0 && blk_id < visited.len {
                visited[blk_id] = true
            }
            mut frame2 := work[fi]
            frame2.processed = true
            work[fi] = frame2

            // 1. Push phis defined in this block.
            if blk_id < ctx.phi_placements.len && ctx.phi_placements[blk_id].len > 0 {
                n_phi := ctx.phi_placements[blk_id].len
                for pai in 0 .. n_phi {
                    alloc_id := ctx.phi_placements[blk_id][pai]
                    if pai < ctx.phi_vals[blk_id].len {
                        phi_val_id := ctx.phi_vals[blk_id][pai]
                        mut new_stack := ctx.stacks[alloc_id].clone()
                        new_stack << phi_val_id
                        ctx.stacks[alloc_id] = new_stack
                        mut wf := work[fi]
                        wf.pushed_allocs << alloc_id
                        work[fi] = wf
                    }
                }
            }

            // 2. Process instructions: rewrite loads, record store values, nop them.
            mut instrs_to_nop := []int{}
            blk2 := m.blocks[blk_id]
            for val_id in blk2.instrs {
                val := m.values[val_id]
                if val.kind != .instruction {
                    continue
                }
                instr := m.instrs[val.index]
                if instr.op == .store {
                    ptr := instr.operands[1]
                    if ptr < ctx.is_promotable.len && ctx.is_promotable[ptr] {
                        mut new_stack := ctx.stacks[ptr].clone()
                        new_stack << instr.operands[0]
                        ctx.stacks[ptr] = new_stack
                        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 {
                            repl = m.get_or_add_const(val.typ, '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 vid in instrs_to_nop {
                vid_val := m.values[vid]
                mut nop_instr := m.instrs[vid_val.index]
                nop_instr.op = .bitcast
                nop_instr.operands = []
                m.instrs[vid_val.index] = nop_instr
            }

            // 3. Fill phi operands in CFG successors (raw block id as predecessor).
            for succ_id in cfg.succs[blk_id] {
                if succ_id < ctx.phi_placements.len && ctx.phi_placements[succ_id].len > 0 {
                    n_succ_phi := ctx.phi_placements[succ_id].len
                    for spai in 0 .. n_succ_phi {
                        alloc_id := ctx.phi_placements[succ_id][spai]
                        if spai < ctx.phi_vals[succ_id].len {
                            vid := ctx.phi_vals[succ_id][spai]
                            phi_v := m.values[vid]
                            phi_val := if ctx.stacks[alloc_id].len > 0 {
                                ctx.stacks[alloc_id].last()
                            } else {
                                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')
                            }
                            m.append_phi_operands(phi_v.index, phi_val, blk_id)
                        }
                    }
                }
            }
        }

        // 4. Recurse into dominator-tree children.
        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]
            mut frame := work[fi]
            frame.child_idx++
            work[fi] = frame
            work << RenameFrame{
                blk_id: child
            }
        } else {
            // 5. Pop reaching definitions pushed in this block.
            pushed := work[fi].pushed_allocs.clone()
            work.pop()
            for i := pushed.len - 1; i >= 0; i-- {
                mut s := ctx.stacks[pushed[i]].clone()
                s.pop()
                ctx.stacks[pushed[i]] = s
            }
        }
    }
}