// 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 os
import v2.ssa
import time

pub struct OptimizeOptions {
pub:
    verify_each_pass bool
    strict_verify    bool
}

// Optimize Module
pub fn optimize(mut m ssa.Module) {
    strict_verify := os.getenv('V2_VERIFY_STRICT') != ''
    optimize_with_options(mut m, OptimizeOptions{
        verify_each_pass: os.getenv('V2_VERIFY') != ''
        strict_verify:    strict_verify
    })
}

pub fn optimize_with_options(mut m ssa.Module, opts OptimizeOptions) {
    mut t0 := time.ticks()
    verify_pipeline_checkpoint(m, opts, 'input')
    rebuild_use_lists(mut m)

    // 1. Build Control Flow Graph (Predecessors)
    cfg := build_cfg(mut m)
    verify_pipeline_checkpoint(m, opts, 'build_cfg')
    mut t1 := time.ticks()
    eprintln('  opt: build_cfg: ${t1 - t0}ms')
    t0 = t1

    // 2. Compute Dominator Tree (Lengauer-Tarjan)
    dom := compute_dominators(mut m, &cfg)
    verify_pipeline_checkpoint(m, opts, 'compute_dominators')
    t1 = time.ticks()
    // Count dom tree stats
    mut n_with_idom := 0
    mut n_in_dom_tree := 0
    mut n_no_idom := 0
    mut n_idom_self := 0
    mut n_idom_diff := 0
    for bi in 0 .. dom.idom.len {
        if dom.idom[bi] >= 0 {
            n_with_idom += 1
            if dom.idom[bi] == bi {
                n_idom_self += 1
            } else {
                n_idom_diff += 1
            }
        } else {
            n_no_idom += 1
        }
        if bi < dom.dom_tree.len {
            n_in_dom_tree += dom.dom_tree[bi].len
        }
    }
    eprintln('  opt: compute_dominators: ${t1 - t0}ms (with_idom=${n_with_idom} no_idom=${n_no_idom} self=${n_idom_self} diff=${n_idom_diff} dom_tree=${n_in_dom_tree})')
    t0 = t1

    // 3. Promote Memory to Register (Construct SSA / Phi Nodes)
    promote_memory_to_register(mut m, dom, &cfg)
    verify_pipeline_checkpoint(m, opts, 'promote_memory_to_register')
    rebuild_use_lists(mut m)
    t1 = time.ticks()
    // Count remaining allocas
    mut n_alloca := 0
    for vi in 0 .. m.values.len {
        val := m.values[vi]
        if val.kind == .instruction {
            idx := val.index
            if idx >= 0 && idx < m.instrs.len {
                if m.instrs[idx].op == .alloca {
                    n_alloca += 1
                }
            }
        }
    }
    eprintln('  opt: mem2reg: ${t1 - t0}ms (remaining allocas: ${n_alloca})')
    t0 = t1

    // 4. Simplify trivial Phi Nodes
    simplify_phi_nodes(mut m)
    verify_pipeline_checkpoint(m, opts, 'simplify_phi_nodes')
    rebuild_use_lists(mut m)
    t1 = time.ticks()
    eprintln('  opt: simplify_phi: ${t1 - t0}ms')
    t0 = t1

    // 5. Eliminate Phi Nodes (convert to copies in predecessor blocks)
    eliminate_phi_nodes(mut m)
    verify_pipeline_checkpoint(m, opts, 'eliminate_phi_nodes')
    rebuild_use_lists(mut m)
    t1 = time.ticks()
    eprintln('  opt: eliminate_phi: ${t1 - t0}ms')
    t0 = t1

    // 6. Remove dead instructions after structural SSA rewrites are complete.
    dead_code_elimination(mut m)
    verify_pipeline_checkpoint(m, opts, 'dead_code_elimination')
    t1 = time.ticks()
    eprintln('  opt: dce: ${t1 - t0}ms')
}

fn rebuild_use_lists(mut m ssa.Module) {
    for vi in 0 .. m.values.len {
        mut val := m.values[vi]
        val.uses = []
        m.values[vi] = val
    }
    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
            }
            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
                }
                instr_idx := m.values[val_id].index
                if instr_idx < 0 || instr_idx >= m.instrs.len {
                    continue
                }
                for op_id in m.instrs[instr_idx].operands {
                    if op_id >= 0 && op_id < m.values.len && val_id !in m.values[op_id].uses {
                        mut op_val := m.values[op_id]
                        op_val.uses << val_id
                        m.values[op_id] = op_val
                    }
                }
            }
        }
    }
}

fn verify_pipeline_checkpoint(m &ssa.Module, opts OptimizeOptions, pass_name string) {
    if opts.verify_each_pass || opts.strict_verify {
        verify_and_panic_with_options(m, pass_name, VerifyPanicOptions{
            allow_noncritical: !opts.strict_verify
        })
    }
}