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

// --- Dominators (Lengauer-Tarjan) ---

struct LTContext {
mut:
    parent   []int   // DFS tree parent
    semi     []int   // Semidominator (BlockID)
    vertex   []int   // Map DFS number -> BlockID
    bucket   [][]int // bucket[w] = set of vertices v s.t. semi[v] = w
    dfnum    []int   // DFS number (0 means unvisited)
    ancestor []int   // DSU parent
    label    []int   // DSU label (min semi in path)
    idom     []int   // Immediate dominator (flat array, indexed by block ID)
    dom_tree [][]int // Dom tree children (flat array, indexed by block ID)
    n        int     // Counter
}

// DFS frame for iterative lt_dfs
struct DfsFrame {
mut:
    node     int
    succ_idx int // next successor to process
}

// DomInfo holds dominator results in flat arrays to avoid struct field access issues
struct DomInfo {
mut:
    idom     []int   // Immediate dominator, indexed by block ID
    dom_tree [][]int // Dom tree children, indexed by block ID
}

fn compute_dominators(mut m ssa.Module, cfg &CfgData) DomInfo {
    max_id := m.blocks.len

    mut ctx := LTContext{
        parent:   []int{len: max_id, init: -1}
        semi:     []int{len: max_id, init: -1}
        vertex:   []int{len: max_id + 1, init: -1}
        bucket:   [][]int{len: max_id}
        dfnum:    []int{len: max_id, init: 0}
        ancestor: []int{len: max_id, init: -1}
        label:    []int{len: max_id, init: -1}
        idom:     []int{len: max_id, init: -1}
        dom_tree: [][]int{len: max_id}
        n:        0
    }

    for fi in 0 .. m.funcs.len {
        func := m.funcs[fi]
        if func.blocks.len == 0 {
            continue
        }

        // Validate that all block IDs and their successor/predecessor
        // references are within bounds.
        mut valid := true
        n_func_blocks := func.blocks.len
        for fbi in 0 .. n_func_blocks {
            blk_id := func.blocks[fbi]
            if blk_id < 0 || blk_id >= max_id {
                valid = false
                break
            }
            n_succs := cfg.succs[blk_id].len
            for si in 0 .. n_succs {
                s := cfg.succs[blk_id][si]
                if s < 0 || s >= max_id {
                    valid = false
                    break
                }
            }
            if !valid {
                break
            }
            n_preds := cfg.preds[blk_id].len
            for pi in 0 .. n_preds {
                p := cfg.preds[blk_id][pi]
                if p < 0 || p >= max_id {
                    valid = false
                    break
                }
            }
            if !valid {
                break
            }
        }
        if !valid {
            continue
        }

        // Reset context for this function (only reset what's needed)
        ctx.n = 0
        for fbi2 in 0 .. n_func_blocks {
            blk_id := func.blocks[fbi2]
            ctx.parent[blk_id] = -1
            ctx.semi[blk_id] = blk_id
            ctx.vertex[blk_id] = -1
            ctx.bucket[blk_id].clear()
            ctx.dfnum[blk_id] = 0
            ctx.ancestor[blk_id] = -1
            ctx.label[blk_id] = blk_id
            ctx.idom[blk_id] = -1
            ctx.dom_tree[blk_id] = []
        }

        entry := func.blocks[0]
        lt_dfs(entry, cfg, mut ctx)

        // Process in reverse DFS order (skip root)
        for i := ctx.n; i >= 2; i-- {
            if i >= ctx.vertex.len {
                continue
            }
            w := ctx.vertex[i]
            if w < 0 || w >= max_id {
                continue
            }

            // 1. Calculate Semidominator - use cfg.preds instead of m.blocks[w].preds
            n_w_preds := cfg.preds[w].len
            for pi2 in 0 .. n_w_preds {
                p := cfg.preds[w][pi2]
                if p < 0 || p >= max_id || ctx.dfnum[p] == 0 {
                    continue
                }

                u := ctx.eval(p)
                if u < 0 || u >= max_id {
                    continue
                }
                semi_u := ctx.semi[u]
                semi_w := ctx.semi[w]
                if semi_u < 0 || semi_u >= max_id || semi_w < 0 || semi_w >= max_id {
                    continue
                }
                if ctx.dfnum[semi_u] < ctx.dfnum[semi_w] {
                    ctx.semi[w] = ctx.semi[u]
                }
            }

            // Add w to bucket of its semidominator
            semi_w2 := ctx.semi[w]
            if semi_w2 >= 0 && semi_w2 < max_id {
                ctx.bucket[semi_w2] << w
            }

            // Link to parent in forest
            ctx.link(ctx.parent[w], w)

            // 2. Implicitly compute IDom
            parent_w := ctx.parent[w]
            if parent_w < 0 || parent_w >= max_id {
                continue
            }
            // Drain bucket of parent
            n_bucket := ctx.bucket[parent_w].len
            for bvi in 0 .. n_bucket {
                v := ctx.bucket[parent_w][bvi]
                if v < 0 || v >= max_id {
                    continue
                }
                u := ctx.eval(v)
                if u < 0 || u >= max_id || v >= max_id {
                    continue
                }
                if ctx.semi[u] == ctx.semi[v] {
                    ctx.idom[v] = parent_w
                } else {
                    ctx.idom[v] = u // Deferred: idom[v] = idom[u]
                }
            }
            ctx.bucket[parent_w].clear()
        }

        // 3. Explicitly compute IDom
        for i := 2; i <= ctx.n; i++ {
            if i >= ctx.vertex.len {
                continue
            }
            w := ctx.vertex[i]
            if w < 0 || w >= max_id {
                continue
            }
            semi_w := ctx.semi[w]
            if semi_w < 0 || semi_w >= max_id {
                continue
            }
            dfnum_semi_w := ctx.dfnum[semi_w]
            if dfnum_semi_w < 0 || dfnum_semi_w >= ctx.vertex.len {
                continue
            }
            target := ctx.vertex[dfnum_semi_w]
            if target < 0 {
                continue
            }
            if ctx.idom[w] != target {
                idom_w := ctx.idom[w]
                if idom_w >= 0 && idom_w < max_id {
                    next_idom := ctx.idom[idom_w]
                    if next_idom >= 0 {
                        ctx.idom[w] = next_idom
                    }
                }
            }
        }

        ctx.idom[entry] = entry

        // Build Dom Tree Children
        for fbi4 in 0 .. n_func_blocks {
            blk_id := func.blocks[fbi4]
            idom := ctx.idom[blk_id]
            if idom != -1 && idom != blk_id {
                if idom >= 0 && idom < max_id {
                    ctx.dom_tree[idom] << blk_id
                }
            }
        }
    }

    return DomInfo{
        idom:     ctx.idom
        dom_tree: ctx.dom_tree
    }
}

fn lt_dfs(root int, cfg &CfgData, mut ctx LTContext) {
    if root < 0 || root >= ctx.dfnum.len {
        return
    }
    n_total_blocks := cfg.succs.len

    mut stack := []DfsFrame{}
    // Visit root
    ctx.n++
    if ctx.n >= ctx.vertex.len {
        return
    }
    ctx.dfnum[root] = ctx.n
    ctx.vertex[ctx.n] = root

    stack << DfsFrame{
        node: root
    }

    for stack.len > 0 {
        top := stack.len - 1
        node := stack[top].node
        if node < 0 || node >= n_total_blocks {
            stack.pop()
            continue
        }
        // Use cfg.succs flat array instead of m.blocks[node].succs
        n_succs := cfg.succs[node].len
        si := stack[top].succ_idx
        if si < n_succs {
            // Avoid stack[top].succ_idx++ -- chained field increment broken in ARM64 self-hosted
            mut frame := stack[top]
            frame.succ_idx++
            stack[top] = frame
            w := cfg.succs[node][si]
            if w >= 0 && w < ctx.dfnum.len && ctx.dfnum[w] == 0 {
                ctx.parent[w] = node
                ctx.n++
                if ctx.n >= ctx.vertex.len {
                    return
                }
                ctx.dfnum[w] = ctx.n
                ctx.vertex[ctx.n] = w
                stack << DfsFrame{
                    node: w
                }
            }
        } else {
            stack.pop()
        }
    }
}

fn (mut ctx LTContext) compress(v int) {
    if v < 0 || v >= ctx.ancestor.len {
        return
    }
    mut chain := []int{}
    mut cur := v
    for cur >= 0 && cur < ctx.ancestor.len {
        av := ctx.ancestor[cur]
        if av < 0 || av >= ctx.ancestor.len {
            break
        }
        if ctx.ancestor[av] == -1 {
            break
        }
        chain << cur
        cur = av
    }
    for ci := chain.len - 1; ci >= 0; ci-- {
        node := chain[ci]
        if node < 0 || node >= ctx.ancestor.len {
            continue
        }
        anc := ctx.ancestor[node]
        if anc < 0 || anc >= ctx.label.len || node >= ctx.label.len {
            continue
        }
        label_anc := ctx.label[anc]
        label_node := ctx.label[node]
        if label_anc < 0 || label_anc >= ctx.semi.len || label_node < 0
            || label_node >= ctx.semi.len {
            continue
        }
        semi_label_anc := ctx.semi[label_anc]
        semi_label_node := ctx.semi[label_node]
        if semi_label_anc < 0 || semi_label_anc >= ctx.dfnum.len || semi_label_node < 0
            || semi_label_node >= ctx.dfnum.len {
            continue
        }
        if ctx.dfnum[semi_label_anc] < ctx.dfnum[semi_label_node] {
            ctx.label[node] = ctx.label[anc]
        }
        aav := ctx.ancestor[anc]
        if aav >= 0 {
            ctx.ancestor[node] = aav
        }
    }
}

fn (mut ctx LTContext) eval(v int) int {
    if v < 0 || v >= ctx.ancestor.len {
        if v >= 0 && v < ctx.label.len {
            return ctx.label[v]
        }
        return 0
    }
    if ctx.ancestor[v] == -1 {
        return ctx.label[v]
    }
    ctx.compress(v)
    av := ctx.ancestor[v]
    if av < 0 || av >= ctx.label.len {
        if v < ctx.label.len {
            return ctx.label[v]
        }
        return 0
    }
    if v >= ctx.label.len {
        return 0
    }
    label_av := ctx.label[av]
    label_v := ctx.label[v]
    if label_av < 0 || label_av >= ctx.semi.len || label_v < 0 || label_v >= ctx.semi.len {
        return label_v
    }
    semi_lav := ctx.semi[label_av]
    semi_lv := ctx.semi[label_v]
    if semi_lav < 0 || semi_lav >= ctx.dfnum.len || semi_lv < 0 || semi_lv >= ctx.dfnum.len {
        return label_v
    }
    if ctx.dfnum[semi_lav] >= ctx.dfnum[semi_lv] {
        return label_v
    }
    return label_av
}

fn (mut ctx LTContext) link(v int, w int) {
    ctx.ancestor[w] = v
}