module optimize

import v3.ssa

// constant_fold supports constant fold handling for optimize.
fn constant_fold(mut m ssa.Module) bool {
	mut changed := false
	for fi in 0 .. m.funcs.len {
		for blk_id in m.funcs[fi].blocks {
			for val_id in m.blocks[blk_id].instrs {
				if m.values[val_id].kind != .instruction {
					continue
				}
				instr := m.instrs[m.values[val_id].index]
				if instr.operands.len != 2 {
					continue
				}
				lhs := m.values[instr.operands[0]]
				rhs := m.values[instr.operands[1]]

				// Skip undef values - can't fold/simplify with undefined.
				if lhs.kind == .constant && lhs.name == 'undef' {
					continue
				}
				if rhs.kind == .constant && rhs.name == 'undef' {
					continue
				}

				// Algebraic simplifications first (work even with non-constant operands):
				// x+0, x*1, x*0, x-x, x^x, x&x, x|x, x<<0, x*2 -> x<<1, etc.
				repl, needs_zero := try_algebraic_simplify(m, val_id, instr, lhs, rhs)
				if repl >= 0 || repl == -2 {
					if repl == -2 {
						// x * 2 -> x << 1
						other_id := if lhs.kind == .constant {
							instr.operands[1]
						} else {
							instr.operands[0]
						}
						typ := m.values[val_id].typ
						one_val := m.get_or_add_const(typ, '1')
						mut shl_instr := m.instrs[m.values[val_id].index]
						shl_instr.op = .shl
						shl_instr.operands = [other_id, one_val]
						m.instrs[m.values[val_id].index] = shl_instr
						changed = true
						continue
					} else if needs_zero {
						// x * 0, x & 0, x - x, x ^ x -> 0
						typ := m.values[val_id].typ
						zero_val := m.get_or_add_const(typ, '0')
						m.replace_uses(val_id, zero_val)
					} else {
						m.replace_uses(val_id, repl)
					}
					changed = true
					continue
				}

				if lhs.kind != .constant || rhs.kind != .constant {
					continue
				}
				l_int := lhs.name.i64()
				r_int := rhs.name.i64()

				mut result := i64(0)
				mut folded := false

				match instr.op {
					.add {
						result = l_int + r_int
						folded = true
					}
					.sub {
						result = l_int - r_int
						folded = true
					}
					.mul {
						result = l_int * r_int
						folded = true
					}
					.sdiv {
						if r_int != 0 {
							result = l_int / r_int
							folded = true
						}
					}
					.srem {
						if r_int != 0 {
							result = l_int % r_int
							folded = true
						}
					}
					.udiv {
						if r_int != 0 {
							result = i64(u64(l_int) / u64(r_int))
							folded = true
						}
					}
					.urem {
						if r_int != 0 {
							result = i64(u64(l_int) % u64(r_int))
							folded = true
						}
					}
					.and_ {
						result = l_int & r_int
						folded = true
					}
					.or_ {
						result = l_int | r_int
						folded = true
					}
					.xor {
						result = l_int ^ r_int
						folded = true
					}
					.shl {
						if r_int >= 0 && r_int < 64 {
							result = i64(u64(l_int) << u64(r_int))
							folded = true
						}
					}
					.ashr {
						if r_int >= 0 && r_int < 64 {
							result = l_int >> u64(r_int)
							folded = true
						}
					}
					.lshr {
						if r_int >= 0 && r_int < 64 {
							result = i64(u64(l_int) >> u64(r_int))
							folded = true
						}
					}
					.eq {
						result = if l_int == r_int { 1 } else { 0 }
						folded = true
					}
					.ne {
						result = if l_int != r_int { 1 } else { 0 }
						folded = true
					}
					.lt {
						result = if l_int < r_int { 1 } else { 0 }
						folded = true
					}
					.gt {
						result = if l_int > r_int { 1 } else { 0 }
						folded = true
					}
					.le {
						result = if l_int <= r_int { 1 } else { 0 }
						folded = true
					}
					.ge {
						result = if l_int >= r_int { 1 } else { 0 }
						folded = true
					}
					.ult {
						result = if u64(l_int) < u64(r_int) { 1 } else { 0 }
						folded = true
					}
					.ugt {
						result = if u64(l_int) > u64(r_int) { 1 } else { 0 }
						folded = true
					}
					.ule {
						result = if u64(l_int) <= u64(r_int) { 1 } else { 0 }
						folded = true
					}
					.uge {
						result = if u64(l_int) >= u64(r_int) { 1 } else { 0 }
						folded = true
					}
					else {}
				}

				if folded {
					typ := m.values[val_id].typ
					const_val := m.get_or_add_const(typ, '${result}')
					m.replace_uses(val_id, const_val)
					changed = true
				}
			}
		}
	}
	return changed
}

// try_algebraic_simplify implements identity/strength-reduction simplifications:
// x+0=x, x*1=x, x*0=0, x-x=0, x^x=0, x&x=x, x|x=x, x<<0=x, x*2=x<<1, x/1=x.
// Returns (replacement_id, needs_zero):
//   repl >= 0          -> replace all uses of val_id with `repl`
//   repl == -2         -> rewrite as `x << 1` (caller builds the shift)
//   needs_zero == true -> replace all uses with a fresh zero constant
//   repl == -1         -> no simplification
fn try_algebraic_simplify(_m &ssa.Module, val_id int, instr ssa.Instruction, lhs ssa.Value, rhs ssa.Value) (int, bool) {
	lhs_id := instr.operands[0]
	rhs_id := instr.operands[1]

	// Same-operand simplifications (x op x). Integer ops only.
	if lhs_id == rhs_id {
		match instr.op {
			.sub, .xor {
				return val_id, true // x - x = 0, x ^ x = 0
			}
			.and_, .or_ {
				return lhs_id, false // x & x = x, x | x = x
			}
			else {}
		}
	}

	mut const_val := i64(0)
	mut const_is_rhs := false
	mut other_id := 0
	if lhs.kind == .constant && lhs.name != 'undef' {
		const_val = lhs.name.i64()
		const_is_rhs = false
		other_id = instr.operands[1]
	} else if rhs.kind == .constant && rhs.name != 'undef' {
		const_val = rhs.name.i64()
		const_is_rhs = true
		other_id = instr.operands[0]
	} else {
		return -1, false
	}

	match instr.op {
		.add {
			if const_val == 0 {
				return other_id, false // x + 0 = x
			}
		}
		.sub {
			if const_is_rhs && const_val == 0 {
				return other_id, false // x - 0 = x
			}
		}
		.mul {
			if const_val == 0 {
				return val_id, true // x * 0 = 0
			}
			if const_val == 1 {
				return other_id, false // x * 1 = x
			}
			if const_val == 2 {
				return -2, false // x * 2 = x << 1
			}
		}
		.sdiv, .udiv {
			if const_is_rhs && const_val == 1 {
				return other_id, false // x / 1 = x
			}
		}
		.and_ {
			if const_val == 0 {
				return val_id, true // x & 0 = 0
			}
		}
		.or_ {
			if const_val == 0 {
				return other_id, false // x | 0 = x
			}
		}
		.xor {
			if const_val == 0 {
				return other_id, false // x ^ 0 = x
			}
		}
		.shl, .ashr, .lshr {
			if const_is_rhs && const_val == 0 {
				return other_id, false // x << 0 = x, x >> 0 = x
			}
		}
		else {}
	}

	return -1, false
}

// branch_fold supports branch fold handling for optimize.
fn branch_fold(mut m ssa.Module) bool {
	mut changed := false
	for fi in 0 .. m.funcs.len {
		for blk_id in m.funcs[fi].blocks {
			if m.blocks[blk_id].instrs.len == 0 {
				continue
			}
			term_val_id := m.blocks[blk_id].instrs.last()
			term := m.instrs[m.values[term_val_id].index]
			if term.op == .br && term.operands.len >= 3 {
				cond_val := m.values[term.operands[0]]
				if cond_val.kind == .constant {
					cond_int := cond_val.name.i64()
					target := if cond_int != 0 { term.operands[1] } else { term.operands[2] }
					mut jmp_instr := m.instrs[m.values[term_val_id].index]
					jmp_instr.op = .jmp
					mut new_ops := []ssa.ValueID{}
					new_ops << target
					jmp_instr.operands = new_ops
					m.instrs[m.values[term_val_id].index] = jmp_instr
					changed = true
				}
			}
		}
	}
	return changed
}