module ssa

import os
import v2.ast
import v2.parser
import v2.pref
import v2.token
import v2.transformer
import v2.types

fn module_storage_ssa_tmp_dir(label string) string {
	return os.join_path(os.temp_dir(), 'v2_module_storage_ssa_${label}_${os.getpid()}')
}

fn module_storage_ssa_builder_for_test_sources(label string, sources map[string]string) &Builder {
	tmp_dir := module_storage_ssa_tmp_dir(label)
	os.mkdir_all(tmp_dir) or { panic('cannot create ${tmp_dir}') }
	defer {
		os.rmdir_all(tmp_dir) or {}
	}
	mut paths := []string{}
	for rel_path, code in sources {
		path := os.join_path(tmp_dir, rel_path)
		os.mkdir_all(os.dir(path)) or { panic('cannot create ${os.dir(path)}') }
		os.write_file(path, code) or { panic('cannot write ${path}') }
		paths << path
	}
	prefs := &pref.Preferences{
		backend: .x64
		arch:    .x64
	}
	mut file_set := token.FileSet.new()
	mut par := parser.Parser.new(prefs)
	files := par.parse_files(paths, mut file_set)
	env := types.Environment.new()
	mut checker := types.Checker.new(prefs, file_set, env)
	checker.check_files(files)
	mut trans := transformer.Transformer.new_with_pref(env, prefs)
	trans.set_file_set(file_set)
	transformed_files := trans.transform_files(files)
	mut mod := Module.new('module_storage_test')
	mut builder := Builder.new_with_env(mod, env)
	builder.build_all(transformed_files)
	return builder
}

fn module_storage_ssa_for_test_sources(label string, sources map[string]string) &Module {
	builder := module_storage_ssa_builder_for_test_sources(label, sources)
	return builder.mod
}

fn module_storage_ssa_for_test_sources_flat(label string, sources map[string]string) &Module {
	tmp_dir := module_storage_ssa_tmp_dir(label)
	os.mkdir_all(tmp_dir) or { panic('cannot create ${tmp_dir}') }
	defer {
		os.rmdir_all(tmp_dir) or {}
	}
	mut paths := []string{}
	for rel_path, code in sources {
		path := os.join_path(tmp_dir, rel_path)
		os.mkdir_all(os.dir(path)) or { panic('cannot create ${os.dir(path)}') }
		os.write_file(path, code) or { panic('cannot write ${path}') }
		paths << path
	}
	prefs := &pref.Preferences{
		backend: .arm64
	}
	mut file_set := token.FileSet.new()
	mut par := parser.Parser.new(prefs)
	files := par.parse_files(paths, mut file_set)
	env := types.Environment.new()
	mut checker := types.Checker.new(prefs, file_set, env)
	checker.check_files(files)
	mut trans := transformer.Transformer.new_with_pref(env, prefs)
	trans.set_file_set(file_set)
	flat := ast.flatten_files(files)
	transformed_flat, _ := trans.transform_files_to_flat(&flat, files)
	mut mod := Module.new('module_storage_flat_test')
	mut builder := Builder.new_with_env(mod, env)
	builder.build_all_from_flat(&transformed_flat)
	return mod
}

fn module_storage_ssa_global_names(m &Module) []string {
	mut names := []string{}
	for global in m.globals {
		names << global.name
	}
	return names
}

fn module_storage_ssa_global_value_id(m &Module, name string) ?ValueID {
	for value in m.values {
		if value.kind == .global && value.name == name {
			return value.id
		}
	}
	return none
}

fn module_storage_ssa_has_store_to_global(m &Module, name string) bool {
	global_id := module_storage_ssa_global_value_id(m, name) or { return false }
	for instr in m.instrs {
		if instr.op == .store && instr.operands.len >= 2 && instr.operands[1] == global_id {
			return true
		}
	}
	return false
}

fn module_storage_ssa_func(m &Module, name string) ?Function {
	for func in m.funcs {
		if func.name == name {
			return func
		}
	}
	return none
}

fn module_storage_ssa_return_value_id(m &Module, func Function) ?ValueID {
	for blk_id in func.blocks {
		for val_id in m.blocks[blk_id].instrs {
			instr := m.instrs[m.values[val_id].index]
			if instr.op == .ret && instr.operands.len > 0 {
				return instr.operands[0]
			}
		}
	}
	return none
}

fn module_storage_ssa_return_value_ids(m &Module, func Function) []ValueID {
	mut ret_vals := []ValueID{}
	for blk_id in func.blocks {
		for val_id in m.blocks[blk_id].instrs {
			instr := m.instrs[m.values[val_id].index]
			if instr.op == .ret && instr.operands.len > 0 {
				ret_vals << instr.operands[0]
			}
		}
	}
	return ret_vals
}

fn module_storage_ssa_call_callees(m &Module, func Function) []string {
	mut callees := []string{}
	for blk_id in func.blocks {
		for val_id in m.blocks[blk_id].instrs {
			value := m.values[val_id]
			if value.kind != .instruction {
				continue
			}
			instr := m.instrs[value.index]
			if instr.op != .call || instr.operands.len == 0 {
				continue
			}
			callee_id := instr.operands[0]
			if callee_id <= 0 || callee_id >= m.values.len {
				continue
			}
			callee := m.values[callee_id]
			if callee.kind == .func_ref {
				callees << callee.name
			}
		}
	}
	return callees
}

fn test_module_storage_import_alias_resolves_to_declaring_module_in_ssa() {
	m := module_storage_ssa_for_test_sources('alias', {
		'report/report.v': 'module report

pub __global mut errors = 0
'
		'main.v':          'module main

import report as r

fn main() {
	r.errors += 1
}
'
	})
	names := module_storage_ssa_global_names(m)
	assert 'report__errors' in names
	assert 'r__errors' !in names
	assert module_storage_ssa_has_store_to_global(m, 'report__errors')
}

fn module_storage_selective_import_nested_sources() map[string]string {
	return {
		'main.v':                              'module main

import mymodules.submodule { sub_xy }

fn main() {
	value := sub_xy(10, 7)
	_ = value
}
'
		'mymodules/submodule/sub_functions.v': 'module submodule

pub fn sub_xy(x int, y int) int {
	return x - y
}
'
	}
}

fn module_storage_assert_nested_selective_import_callee(m &Module, label string) {
	main_func := module_storage_ssa_func(m, 'main') or { panic('${label}: missing main') }
	callees := module_storage_ssa_call_callees(m, main_func)
	assert 'submodule__sub_xy' in callees, '${label}: missing submodule__sub_xy in ${callees}'
	assert 'sub_xy' !in callees, '${label}: bare sub_xy callee leaked into ${callees}'
	assert 'mymodules_submodule__sub_xy' !in callees, '${label}: import path callee leaked into ${callees}'
}

fn test_module_storage_legacy_selective_import_nested_module_uses_declared_leaf_module() {
	m := module_storage_ssa_for_test_sources('selective_nested_legacy',
		module_storage_selective_import_nested_sources())
	module_storage_assert_nested_selective_import_callee(m, 'legacy')
}

fn test_module_storage_flat_selective_import_nested_module_uses_declared_leaf_module() {
	m := module_storage_ssa_for_test_sources_flat('selective_nested_flat',
		module_storage_selective_import_nested_sources())
	module_storage_assert_nested_selective_import_callee(m, 'flat')
}

fn module_storage_direct_import_nested_sources() map[string]string {
	return {
		'main.v':                              'module main

import mymodules.submodule

fn main() {
	value := submodule.sub_xy(10, 7)
	_ = value
}
'
		'mymodules/submodule/sub_functions.v': 'module submodule

pub fn sub_xy(x int, y int) int {
	return x - y
}
'
	}
}

fn test_module_storage_legacy_direct_import_nested_module_selector_uses_declared_leaf_module() {
	m := module_storage_ssa_for_test_sources('direct_nested_legacy',
		module_storage_direct_import_nested_sources())
	module_storage_assert_nested_selective_import_callee(m, 'legacy direct')
}

fn test_module_storage_flat_direct_import_nested_module_selector_uses_declared_leaf_module() {
	m := module_storage_ssa_for_test_sources_flat('direct_nested_flat',
		module_storage_direct_import_nested_sources())
	module_storage_assert_nested_selective_import_callee(m, 'flat direct')
}

fn test_module_storage_legacy_direct_import_module_selector_one_arg_stays_call() {
	m := module_storage_ssa_for_test_sources('module_selector_one_arg_call', {
		'main.v':          'module main

import left as l
import right
import helper

fn main() {
	first := l.state == 10 && right.state == 200 && l.score() == 1001 && right.score() == 20005
	direct := l.bump(3)
	transitive := helper.bump_right(5)
	second := direct == 132 && transitive == 2057 && l.state == 13 && right.state == 205
	_ = first
	_ = second
}
'
		'left/left.v':     'module left

pub __global mut state = 10
__global mut private_hits = 1

pub fn bump(delta int) int {
	state += delta
	private_hits += 1
	return state * 10 + private_hits
}

pub fn score() int {
	return state * 100 + private_hits
}
'
		'right/right.v':   'module right

pub __global mut state = 200
__global mut private_hits = 5

pub fn bump(delta int) int {
	state += delta
	private_hits += 2
	return state * 10 + private_hits
}

pub fn score() int {
	return state * 100 + private_hits
}
'
		'helper/helper.v': 'module helper

import right as r

pub fn bump_right(delta int) int {
	return r.bump(delta)
}

pub fn right_score() int {
	return r.score()
}
'
	})
	main_func := module_storage_ssa_func(m, 'main') or { panic('missing main') }
	callees := module_storage_ssa_call_callees(m, main_func)
	assert 'helper__bump_right' in callees
}

fn test_module_storage_legacy_call_or_cast_module_selector_prefers_registered_call() {
	mut builder := module_storage_ssa_builder_for_test_sources('module_selector_call_or_cast_decision', {
		'helper/helper.v': 'module helper

pub fn bump_right(delta int) int {
	return delta + 1
}
'
		'main.v':          'module main

import helper

fn main() {}
'
	})
	sel := ast.SelectorExpr{
		lhs: ast.Expr(ast.Ident{
			name: 'helper'
		})
		rhs: ast.Ident{
			name: 'bump_right'
		}
	}
	assert !builder.call_or_cast_selector_should_remain_cast(sel, 'helper__bump_right')
}

fn test_module_qualified_alias_array_field_index_has_alias_element_type() {
	m := module_storage_ssa_for_test_sources('selector_alias_array_elem', {
		'v2/ssa/value.v': 'module ssa

pub type ValueID = int
'
		'main.v':         'module main

import v2.ssa

struct Holder {
	items []ssa.ValueID
}

fn get_item(h Holder, i int) ssa.ValueID {
	return h.items[i]
}
'
	})
	get_item := module_storage_ssa_func(m, 'main__get_item') or {
		module_storage_ssa_func(m, 'get_item') or {
			module_storage_ssa_func(m, 'main.get_item') or { panic('missing get_item') }
		}
	}
	mut ret_val := ValueID(0)
	for blk_id in get_item.blocks {
		for val_id in m.blocks[blk_id].instrs {
			instr := m.instrs[m.values[val_id].index]
			if instr.op == .ret && instr.operands.len > 0 {
				ret_val = instr.operands[0]
			}
		}
	}
	assert ret_val != 0
	ret_typ := m.type_store.types[m.values[ret_val].typ]
	assert ret_typ.kind == .int_t
	assert ret_typ.width == 32
}

fn test_map_array_alias_value_index_has_alias_element_type() {
	m := module_storage_ssa_for_test_sources('map_array_alias_value_index', {
		'main.v': 'module main

type TypeID = int

fn read_elem(values map[string][]TypeID, key string, idx int) TypeID {
	param_elem_types := values[key] or { [TypeID(0)] }
	return param_elem_types[idx]
}
'
	})
	read_elem := module_storage_ssa_func(m, 'main__read_elem') or {
		module_storage_ssa_func(m, 'read_elem') or { panic('missing read_elem') }
	}
	ret_val := module_storage_ssa_return_value_id(m, read_elem) or {
		panic('missing read_elem return value')
	}
	ret_typ := m.type_store.types[m.values[ret_val].typ]
	assert ret_typ.kind == .int_t
	assert ret_typ.width == 32
}

fn test_method_map_array_alias_value_index_has_alias_element_type() {
	m := module_storage_ssa_for_test_sources('method_map_array_alias_value_index', {
		'main.v': 'module main

type TypeID = int

struct Builder {
mut:
	fn_param_array_elem_types map[string][]TypeID
}

fn (mut b Builder) call_param_array_elem_type(fn_name string, param_idx int) TypeID {
	param_elem_types := b.fn_param_array_elem_types[fn_name] or { return TypeID(-1) }
	if param_idx >= param_elem_types.len {
		return TypeID(-2)
	}
	elem_type := param_elem_types[param_idx]
	if elem_type == 0 {
		return TypeID(-3)
	}
	return elem_type
}
'
	})
	func := module_storage_ssa_func(m, 'Builder__call_param_array_elem_type') or {
		panic('missing method function')
	}
	ret_vals := module_storage_ssa_return_value_ids(m, func)
	assert ret_vals.len > 0
	for ret_val in ret_vals {
		ret_typ := m.type_store.types[m.values[ret_val].typ]
		assert ret_typ.kind == .int_t
		assert ret_typ.width == 32
	}
}

fn test_flat_method_map_array_alias_value_index_has_alias_element_type() {
	m := module_storage_ssa_for_test_sources_flat('flat_method_map_array_alias_value_index', {
		'main.v': 'module main

type TypeID = int

struct Builder {
mut:
	fn_param_array_elem_types map[string][]TypeID
}

fn (mut b Builder) call_param_array_elem_type(fn_name string, param_idx int) TypeID {
	param_elem_types := b.fn_param_array_elem_types[fn_name] or { return TypeID(-1) }
	if param_idx >= param_elem_types.len {
		return TypeID(-2)
	}
	elem_type := param_elem_types[param_idx]
	if elem_type == 0 {
		return TypeID(-3)
	}
	return elem_type
}
'
	})
	func := module_storage_ssa_func(m, 'Builder__call_param_array_elem_type') or {
		panic('missing method function')
	}
	ret_vals := module_storage_ssa_return_value_ids(m, func)
	assert ret_vals.len > 0
	for ret_val in ret_vals {
		ret_typ := m.type_store.types[m.values[ret_val].typ]
		assert ret_typ.kind == .int_t
		assert ret_typ.width == 32
	}
}

fn test_module_storage_same_short_names_stay_distinct_in_ssa() {
	m := module_storage_ssa_for_test_sources('collisions', {
		'a/a.v':  'module a

pub __global state = 1
'
		'b/b.v':  'module b

pub __global state = 2
'
		'main.v': 'module main

import a
import b

fn main() {
	_ = a.state + b.state
}
'
	})
	names := module_storage_ssa_global_names(m)
	assert 'a__state' in names
	assert 'b__state' in names
	assert 'state' !in names
	a_id := module_storage_ssa_global_value_id(m, 'a__state') or {
		panic('missing a__state global')
	}
	b_id := module_storage_ssa_global_value_id(m, 'b__state') or {
		panic('missing b__state global')
	}
	assert a_id != b_id
}

fn test_module_storage_bare_dunder_name_resolves_to_current_module_in_ssa() {
	m := module_storage_ssa_for_test_sources('bare_dunder', {
		'a/a.v':  'module a

pub __global mut state__x = 1

pub fn bump() int {
	state__x += 1
	return state__x
}
'
		'b/b.v':  'module b

pub __global mut state__x = 2

pub fn bump() int {
	state__x += 1
	return state__x
}
'
		'main.v': 'module main

import a
import b

fn main() {
	_ = a.bump() + b.bump()
	_ = a.state__x + b.state__x
}
'
	})
	names := module_storage_ssa_global_names(m)
	assert 'a__state__x' in names
	assert 'b__state__x' in names
	assert 'state__x' !in names
	assert module_storage_ssa_has_store_to_global(m, 'a__state__x')
	assert module_storage_ssa_has_store_to_global(m, 'b__state__x')
	a_id := module_storage_ssa_global_value_id(m, 'a__state__x') or {
		panic('missing a__state__x global')
	}
	b_id := module_storage_ssa_global_value_id(m, 'b__state__x') or {
		panic('missing b__state__x global')
	}
	assert a_id != b_id
}

fn test_module_storage_c_extern_raw_global_uses_c_symbol_in_ssa() {
	m := module_storage_ssa_for_test_sources('c_extern_raw', {
		'ext/ext.v': 'module ext

@[c_extern]
pub __global raw_status int

pub fn read_status() int {
	return raw_status
}
'
		'main.v':    'module main

import ext

fn main() {
	_ = ext.raw_status
	_ = ext.read_status()
}
'
	})
	names := module_storage_ssa_global_names(m)
	assert 'raw_status' in names
	assert 'ext__raw_status' !in names
	_ := module_storage_ssa_global_value_id(m, 'raw_status') or {
		panic('missing raw_status C extern global')
	}
	assert module_storage_ssa_global_value_id(m, 'ext__raw_status') == none
}