module optimize

import v3.ssa

// CfgData holds successor/predecessor lists indexed by block id. v3 branch and
// phi operands store raw block ids, so this is built straight from the CFG that
// build_cfg() already wrote into m.blocks.
struct CfgData {
mut:
	succs [][]int
	preds [][]int
}

// cfg_data_from_module snapshots the successor/predecessor sets from m.blocks.
// build_cfg(mut m) must have been called first.
fn cfg_data_from_module(m &ssa.Module) CfgData {
	n_blocks := m.blocks.len
	mut cfg := CfgData{
		succs: [][]int{len: n_blocks}
		preds: [][]int{len: n_blocks}
	}
	for bi in 0 .. n_blocks {
		mut s := []int{cap: m.blocks[bi].succs.len}
		for x in m.blocks[bi].succs {
			s << int(x)
		}
		cfg.succs[bi] = s
		mut p := []int{cap: m.blocks[bi].preds.len}
		for x in m.blocks[bi].preds {
			p << int(x)
		}
		cfg.preds[bi] = p
	}
	return cfg
}

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

// LTContext stores ltcontext state used by optimize.
struct LTContext {
mut:
	parent   []int
	semi     []int
	vertex   []int
	bucket   [][]int
	dfnum    []int
	ancestor []int
	label    []int
	idom     []int
	dom_tree [][]int
	n        int
}

// DfsFrame represents dfs frame data used by optimize.
struct DfsFrame {
mut:
	node     int
	succ_idx int
}

// DomInfo holds dominator results in flat arrays indexed by block id.
struct DomInfo {
mut:
	idom     []int
	dom_tree [][]int
}

// compute_dominators runs Lengauer-Tarjan for every function and also writes the
// idom / dom_tree results back into m.blocks for the verifier's dominance checks.
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
		}

		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
			}
			for s in cfg.succs[blk_id] {
				if s < 0 || s >= max_id {
					valid = false
					break
				}
			}
			if !valid {
				break
			}
			for p in cfg.preds[blk_id] {
				if p < 0 || p >= max_id {
					valid = false
					break
				}
			}
			if !valid {
				break
			}
		}
		if !valid {
			continue
		}

		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] = []
			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)

		for i := ctx.n; i >= 2; i-- {
			if i >= ctx.vertex.len {
				continue
			}
			w := ctx.vertex[i]
			if w < 0 || w >= max_id {
				continue
			}

			for p in cfg.preds[w] {
				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]
				}
			}

			semi_w2 := ctx.semi[w]
			if semi_w2 >= 0 && semi_w2 < max_id {
				mut b := ctx.bucket[semi_w2]
				b << w
				ctx.bucket[semi_w2] = b
			}

			ctx.link(ctx.parent[w], w)

			parent_w := ctx.parent[w]
			if parent_w < 0 || parent_w >= max_id {
				continue
			}
			for v in ctx.bucket[parent_w] {
				if v < 0 || v >= max_id {
					continue
				}
				u := ctx.eval(v)
				if u < 0 || u >= max_id {
					continue
				}
				if ctx.semi[u] == ctx.semi[v] {
					ctx.idom[v] = parent_w
				} else {
					ctx.idom[v] = u
				}
			}
			ctx.bucket[parent_w] = []
		}

		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

		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 {
					mut t := ctx.dom_tree[idom]
					t << blk_id
					ctx.dom_tree[idom] = t
				}
			}
		}
	}

	// Write idom / dom_tree back into the blocks for verifier dominance checks.
	for bi in 0 .. max_id {
		mut blk := m.blocks[bi]
		blk.idom = ctx.idom[bi]
		blk.dom_tree = []ssa.BlockID{cap: ctx.dom_tree[bi].len}
		for c in ctx.dom_tree[bi] {
			blk.dom_tree << ssa.BlockID(c)
		}
		m.blocks[bi] = blk
	}

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

// lt_dfs supports lt dfs handling for optimize.
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{}
	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
		}
		n_succs := cfg.succs[node].len
		si := stack[top].succ_idx
		if si < n_succs {
			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()
		}
	}
}

// compress supports compress handling for LTContext.
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
		}
	}
}

// eval supports eval handling for LTContext.
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
}

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