module optimize

import v3.ssa

// dead_code_elimination supports dead code elimination handling for optimize.
fn dead_code_elimination(mut m ssa.Module) {
	mut changed := true
	for changed {
		changed = false
		for fi in 0 .. m.funcs.len {
			dead_stores := find_dead_stores(m, m.funcs[fi])

			for blk_id in m.funcs[fi].blocks {
				mut dead_instrs := map[int]bool{}
				for val_id in m.blocks[blk_id].instrs {
					val := m.values[val_id]
					if val.kind != .instruction {
						continue
					}
					instr := m.instrs[val.index]
					if instr.op == .store && val_id in dead_stores {
						dead_instrs[val_id] = true
						continue
					}
					has_side_effects := instr.op in [.store, .call, .call_indirect, .call_sret,
						.ret, .br, .jmp, .switch_, .unreachable, .assign, .fence, .cmpxchg,
						.atomicrmw, .go_call, .spawn_call]
					if !has_side_effects && val.uses.len == 0 {
						dead_instrs[val_id] = true
					}
				}

				if dead_instrs.len > 0 {
					for val_id, _ in dead_instrs {
						instr := m.instrs[m.values[val_id].index]

						for op_id in instr.value_operands() {
							remove_use(mut m, op_id, val_id)
						}
					}
					mut new_instrs := []int{cap: m.blocks[blk_id].instrs.len}
					for val_id in m.blocks[blk_id].instrs {
						if val_id !in dead_instrs {
							new_instrs << val_id
						}
					}
					mut blk := m.blocks[blk_id]
					blk.instrs = new_instrs
					m.blocks[blk_id] = blk
					changed = true
				}
			}
		}
	}
}

// remove_use updates remove use state for optimize.
fn remove_use(mut m ssa.Module, val_id int, user_id int) {
	if val_id < 0 || val_id >= m.values.len {
		return
	}
	mut val := m.values[val_id]
	for i := val.uses.len - 1; i >= 0; i-- {
		if val.uses[i] == user_id {
			val.uses.delete(i)
		}
	}
	m.values[val_id] = val
}

// find_dead_stores resolves find dead stores information for optimize.
fn find_dead_stores(m &ssa.Module, func ssa.Function) map[int]bool {
	mut dead_stores := map[int]bool{}
	mut allocas := map[int]bool{}
	mut alloca_stores := map[int][]int{}
	mut alloca_escapes := map[int]bool{}

	for blk_id in func.blocks {
		for val_id in m.blocks[blk_id].instrs {
			val := m.values[val_id]
			if val.kind != .instruction {
				continue
			}
			op := m.instrs[val.index].op
			if op == .alloca || op == .heap_alloc {
				allocas[val_id] = true
			}
		}
	}

	for alloca_id, _ in allocas {
		for use_id in m.values[alloca_id].uses {
			if use_id < 0 || use_id >= m.values.len {
				continue
			}
			val := m.values[use_id]
			if val.kind != .instruction {
				alloca_escapes[alloca_id] = true
				continue
			}
			instr := m.instrs[val.index]
			match instr.op {
				.store {
					if instr.operands.len > 1 && instr.operands[1] == alloca_id {
						alloca_stores[alloca_id] << use_id
					} else {
						alloca_escapes[alloca_id] = true
					}
				}
				.load {}
				else {
					alloca_escapes[alloca_id] = true
				}
			}
		}
	}

	// An alloca with no loads and no escapes → all stores to it are dead
	for alloca_id, _ in allocas {
		if alloca_id in alloca_escapes {
			continue
		}
		mut has_load := false
		for use_id in m.values[alloca_id].uses {
			if use_id >= 0 && use_id < m.values.len && m.values[use_id].kind == .instruction {
				if m.instrs[m.values[use_id].index].op == .load {
					has_load = true
					break
				}
			}
		}
		if !has_load {
			if store_ids := alloca_stores[alloca_id] {
				for store_id in store_ids {
					dead_stores[store_id] = true
				}
			}
		}
	}

	return dead_stores
}