module optimize

import v3.ssa

// verify_ssa validates verify ssa state for optimize.
fn verify_ssa(m &ssa.Module, stage string) {
	verify_module_indexes(m, stage)
	verify_functions(m, stage)
	verify_uses(m, stage)
}

// verify_module_indexes validates verify module indexes state for optimize.
fn verify_module_indexes(m &ssa.Module, stage string) {
	for vi in 0 .. m.values.len {
		if m.values[vi].id != vi {
			verify_fail(stage, 'value ${vi} has id ${m.values[vi].id}')
		}
		if m.values[vi].kind == .instruction {
			instr_idx := m.values[vi].index
			if instr_idx < 0 || instr_idx >= m.instrs.len {
				verify_fail(stage,
					'instruction value ${vi} has invalid instruction index ${instr_idx}')
			}
		}
	}
	for bi in 0 .. m.blocks.len {
		if m.blocks[bi].id != bi {
			verify_fail(stage, 'block ${bi} has id ${m.blocks[bi].id}')
		}
	}
	for fi in 0 .. m.funcs.len {
		if m.funcs[fi].id != fi {
			verify_fail(stage, 'function ${fi} has id ${m.funcs[fi].id}')
		}
	}
}

// verify_functions validates verify functions state for optimize.
fn verify_functions(m &ssa.Module, stage string) {
	for fi in 0 .. m.funcs.len {
		func := m.funcs[fi]
		mut func_blocks := map[int]bool{}
		for blk_id in func.blocks {
			if !verify_block_id(m, stage, blk_id, fi, 'function block list') {
				continue
			}
			if func_blocks[blk_id] {
				verify_fail(stage, 'function ${fi} contains duplicate block ${blk_id}')
			}
			func_blocks[blk_id] = true
		}
		for blk_id in func.blocks {
			if !verify_block_id(m, stage, blk_id, fi, 'function block list') {
				continue
			}
			blk := m.blocks[blk_id]
			if blk.parent != fi {
				verify_fail(stage, 'block ${blk_id} parent is ${blk.parent}, expected ${fi}')
			}
			verify_block_instrs(m, stage, fi, blk_id, func_blocks)
			verify_cfg_edges(m, stage, fi, blk_id, func_blocks)
		}
	}
}

// verify_block_instrs validates verify block instrs state for optimize.
fn verify_block_instrs(m &ssa.Module, stage string, fi int, blk_id int, func_blocks map[int]bool) {
	blk := m.blocks[blk_id]
	if blk.instrs.len == 0 {
		return
	}
	for ii, val_id in blk.instrs {
		if val_id <= 0 || val_id >= m.values.len {
			verify_fail(stage, 'block ${blk_id} contains invalid instruction value ${val_id}')
		}
		val := m.values[val_id]
		if val.kind != .instruction {
			verify_fail(stage, 'block ${blk_id} contains non-instruction value ${val_id}')
		}
		if val.index < 0 || val.index >= m.instrs.len {
			verify_fail(stage,
				'instruction value ${val_id} has invalid instruction index ${val.index}')
		}
		instr := m.instrs[val.index]
		if instr.block != blk_id {
			verify_fail(stage,
				'instruction value ${val_id} is in block ${blk_id} but records block ${instr.block}')
		}
		verify_instruction_operands(m, stage, fi, blk_id, val_id, instr, func_blocks,
			ii == blk.instrs.len - 1)
	}
}

// verify_instruction_operands validates verify instruction operands state for optimize.
fn verify_instruction_operands(m &ssa.Module, stage string, fi int, blk_id int, val_id int, instr ssa.Instruction, func_blocks map[int]bool, is_final bool) {
	for op_id in instr.value_operands() {
		if op_id <= 0 || op_id >= m.values.len {
			verify_fail(stage,
				'instruction value ${val_id} in block ${blk_id} uses invalid value ${op_id}')
		}
		if m.values[op_id].kind == .unknown {
			verify_fail(stage,
				'instruction value ${val_id} in block ${blk_id} uses unknown value ${op_id}')
		}
	}
	match instr.op {
		.ret {
			if instr.operands.len > 1 {
				verify_fail(stage,
					'ret value ${val_id} in block ${blk_id} has ${instr.operands.len} operands')
			}
		}
		.jmp {
			if instr.operands.len != 1 {
				verify_fail(stage,
					'jmp value ${val_id} in block ${blk_id} has ${instr.operands.len} operands')
			} else if is_final {
				verify_target_block(m, stage, fi, blk_id, int(instr.operands[0]), func_blocks,
					'jmp')
			}
		}
		.unreachable {
			if instr.operands.len != 0 {
				verify_fail(stage,
					'unreachable value ${val_id} in block ${blk_id} has ${instr.operands.len} operands')
			}
		}
		.br {
			if instr.operands.len != 3 {
				verify_fail(stage,
					'br value ${val_id} in block ${blk_id} has ${instr.operands.len} operands')
			} else if is_final {
				verify_target_block(m, stage, fi, blk_id, int(instr.operands[1]), func_blocks,
					'br true')
				verify_target_block(m, stage, fi, blk_id, int(instr.operands[2]), func_blocks,
					'br false')
			}
		}
		.phi {
			if instr.operands.len == 0 || instr.operands.len % 2 != 0 {
				verify_fail(stage,
					'phi value ${val_id} in block ${blk_id} has ${instr.operands.len} operands')
			}
			for oi := 1; oi < instr.operands.len; oi += 2 {
				verify_target_block(m, stage, fi, blk_id, int(instr.operands[oi]), func_blocks,
					'phi predecessor')
			}
		}
		else {}
	}
}

// verify_target_block validates verify target block state for optimize.
fn verify_target_block(m &ssa.Module, stage string, fi int, blk_id int, target_id int, func_blocks map[int]bool, label string) {
	if !verify_block_id(m, stage, target_id, fi, '${label} target from block ${blk_id}') {
		return
	}
	if !func_blocks[target_id] {
		verify_fail(stage,
			'${label} target ${target_id} from block ${blk_id} is outside function ${fi}')
	}
}

// verify_cfg_edges validates verify cfg edges state for optimize.
fn verify_cfg_edges(m &ssa.Module, stage string, fi int, blk_id int, func_blocks map[int]bool) {
	blk := m.blocks[blk_id]
	mut expected_succs := []int{}
	if blk.instrs.len > 0 {
		term_val_id := blk.instrs.last()
		term := m.instrs[m.values[term_val_id].index]
		match term.op {
			.jmp {
				if term.operands.len == 1 {
					expected_succs << int(term.operands[0])
				}
			}
			.br {
				if term.operands.len == 3 {
					expected_succs << int(term.operands[1])
					if term.operands[2] != term.operands[1] {
						expected_succs << int(term.operands[2])
					}
				}
			}
			else {}
		}
	}
	verify_same_blocks(stage, 'successors for block ${blk_id}', expected_succs, blk.succs)

	for succ in blk.succs {
		if !verify_block_id(m, stage, succ, fi, 'successor of block ${blk_id}') {
			continue
		}
		if !func_blocks[succ] {
			verify_fail(stage, 'successor ${succ} from block ${blk_id} is outside function ${fi}')
		}
		if !arr_contains(m.blocks[succ].preds, blk_id) {
			verify_fail(stage, 'successor ${succ} is missing predecessor ${blk_id}')
		}
	}
	for pred in blk.preds {
		if !verify_block_id(m, stage, pred, fi, 'predecessor of block ${blk_id}') {
			continue
		}
		if !func_blocks[pred] {
			verify_fail(stage, 'predecessor ${pred} of block ${blk_id} is outside function ${fi}')
		}
		if !arr_contains(m.blocks[pred].succs, blk_id) {
			verify_fail(stage, 'predecessor ${pred} is missing successor ${blk_id}')
		}
	}
}

// verify_uses validates verify uses state for optimize.
fn verify_uses(m &ssa.Module, stage string) {
	mut expected_uses := map[int][]int{}
	for func in m.funcs {
		for blk_id in func.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 := m.instrs[m.values[val_id].index]
				for op_id in instr.value_operands() {
					if op_id > 0 && op_id < m.values.len {
						if !arr_contains(expected_uses[op_id], val_id) {
							expected_uses[op_id] << val_id
						}
					}
				}
			}
		}
	}
	for vi in 0 .. m.values.len {
		verify_same_blocks(stage, 'uses for value ${vi}', expected_uses[vi], m.values[vi].uses)
	}
}

// verify_block_id validates verify block id state for optimize.
fn verify_block_id(m &ssa.Module, stage string, blk_id int, fi int, context string) bool {
	if blk_id < 0 || blk_id >= m.blocks.len {
		verify_fail(stage, '${context} references invalid block ${blk_id} in function ${fi}')
		return false
	}
	return true
}

// verify_same_blocks validates verify same blocks state for optimize.
fn verify_same_blocks(stage string, label string, expected []int, actual []int) {
	if expected.len != actual.len {
		verify_fail(stage, '${label}: expected ${expected}, got ${actual}')
	}
	for v in expected {
		if !arr_contains(actual, v) {
			verify_fail(stage, '${label}: expected ${expected}, got ${actual}')
		}
	}
	for v in actual {
		if !arr_contains(expected, v) {
			verify_fail(stage, '${label}: expected ${expected}, got ${actual}')
		}
	}
}

// verify_fail validates verify fail state for optimize.
fn verify_fail(stage string, msg string) {
	panic('SSA verifier failed after ${stage}: ${msg}')
}