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

fn remove_unreachable_blocks(mut m ssa.Module) {
	// Re-build CFG first
	build_cfg(mut m)
	for fi in 0 .. m.funcs.len {
		if m.funcs[fi].blocks.len == 0 {
			continue
		}
		// BFS/DFS from entry
		mut reachable := map[int]bool{}
		mut q := [m.funcs[fi].blocks[0]]
		reachable[m.funcs[fi].blocks[0]] = true

		for q.len > 0 {
			curr := q[q.len - 1]
			q.delete_last()
			for succ in m.blocks[curr].succs {
				if !reachable[succ] {
					reachable[succ] = true
					q << succ
				}
			}
		}

		mut new_blocks := []int{}
		for blk in m.funcs[fi].blocks {
			if reachable[blk] {
				new_blocks << blk
			}
		}
		// Avoid m.funcs[fi].blocks = ... -- chained field assign broken in ARM64 self-hosted
		mut func := m.funcs[fi]
		func.blocks = new_blocks
		m.funcs[fi] = func
	}
}

fn merge_blocks(mut m ssa.Module) {
	// If Block A jumps unconditionally to B, and B has only A as predecessor:
	// 1. Move instructions from B to A
	// 2. Update A's successors to B's successors
	// 3. Remove B

	// We need to be careful about iteration while modifying.
	// Loop until no changes.
	mut changed := true
	mut first_iter := true
	for changed {
		changed = false
		// Only rebuild CFG on first iteration or after actual changes
		if first_iter {
			build_cfg(mut m)
			first_iter = false
		}

		for fi in 0 .. m.funcs.len {
			// We iterate through blocks.
			// If we merge A->B, we can't merge B->C in same pass easily.
			mut merged := map[int]bool{}

			for blk_id in m.funcs[fi].blocks {
				if merged[blk_id] {
					continue
				}

				// Check if unconditional jump
				if m.blocks[blk_id].instrs.len > 0 {
					last_val := m.blocks[blk_id].instrs[m.blocks[blk_id].instrs.len - 1]
					last_instr := m.instrs[m.values[last_val].index]

					if last_instr.op == .jmp {
						target_val := last_instr.operands[0]
						target_id := m.get_block_from_val(target_val)

						// Check if target has phi nodes - don't merge if it does
						mut has_phi := false
						for vid in m.blocks[target_id].instrs {
							if m.values[vid].kind == .instruction {
								if m.instrs[m.values[vid].index].op == .phi {
									has_phi = true
									break
								}
							}
						}

						// Candidate: target_id (only if no phi nodes)
						if target_id != blk_id && m.blocks[target_id].preds.len == 1
							&& m.blocks[target_id].preds[0] == blk_id && !has_phi {
							// MERGE: Remove JMP from A, then append B's instrs to A
							// Read whole struct, modify, write back (chained broken in ARM64)
							mut merge_blk := m.blocks[blk_id]
							merge_blk.instrs.delete_last()
							merge_blk.instrs << m.blocks[target_id].instrs
							m.blocks[blk_id] = merge_blk

							// Update instructions in B to point to A (for their 'block' field)?
							// Not strictly needed if we just use the list.
							// But we need to update Phis in successors of B?
							// If B has successors, their Phis might refer to B.
							// Since B is gone, they now refer to A.
							for succ_id in m.blocks[target_id].succs {
								n_succ_instrs := m.blocks[succ_id].instrs.len
								for ivi in 0 .. n_succ_instrs {
									iv := m.blocks[succ_id].instrs[ivi]
									v := m.values[iv]
									if v.kind != .instruction {
										continue
									}
									ins := m.instrs[v.index]
									if ins.op == .phi {
										// Replace all occurrences (defensive - handles edge cases)
										// i=1,3,5... are block references in phi [val0, blk0, val1, blk1, ...]
										// Avoid m.instrs[X].operands[i] = ... -- chained broken in ARM64 self-hosted
										mut phi_ops := ins.operands.clone()
										mut phi_modified := false
										for i := 1; i < phi_ops.len; i += 2 {
											if phi_ops[i] == m.blocks[target_id].val_id {
												phi_ops[i] = m.blocks[blk_id].val_id
												phi_modified = true
											}
										}
										if phi_modified {
											mut phi_ins := m.instrs[v.index]
											phi_ins.operands = phi_ops
											m.instrs[v.index] = phi_ins
										}
									}
								}
							}

							// Remove B from func
							merged[target_id] = true
							changed = true
						}
					}
				}
			}

			// Filter out merged blocks
			if merged.len > 0 {
				mut new_blks := []int{}
				for b in m.funcs[fi].blocks {
					if !merged[b] {
						new_blks << b
					}
				}
				// Avoid m.funcs[fi].blocks = ... -- chained field assign broken in ARM64 self-hosted
				mut func2 := m.funcs[fi]
				func2.blocks = new_blks
				m.funcs[fi] = func2
			}
		}
		// Rebuild CFG for next iteration if we made changes
		if changed {
			build_cfg(mut m)
		}
	}
}