module mir

import v3.ssa

// TargetArch lists target arch values used by mir.
pub enum TargetArch {
	unknown
	arm64
}

// AbiKind lists abi kind values used by mir.
pub enum AbiKind {
	unknown
	aapcs64
}

// Target represents target data used by mir.
pub struct Target {
pub:
	arch         TargetArch
	abi          AbiKind
	pointer_size int
	int_arg_regs int
	ret_regs     int
	stack_align  int
}

// AbiLocationKind lists abi location kind values used by mir.
pub enum AbiLocationKind {
	none
	register
	stack
	indirect
}

// AbiLocation represents abi location data used by mir.
pub struct AbiLocation {
pub:
	kind           AbiLocationKind
	register_index int
	register_count int
	stack_offset   int
	size           int
}

// FunctionAbi represents function abi data used by mir.
pub struct FunctionAbi {
pub:
	params          []AbiLocation
	ret             AbiLocation
	stack_args_size int
	stack_align     int
}

// Value represents value data used by mir.
pub struct Value {
pub mut:
	id    int
	kind  ssa.ValueKind
	typ   ssa.TypeID
	name  string
	index int
	uses  []ssa.ValueID
}

// Instruction represents instruction data used by mir.
pub struct Instruction {
pub mut:
	op       ssa.OpCode
	operands []ssa.ValueID
	typ      ssa.TypeID
	block    int
}

// BasicBlock represents basic block data used by mir.
pub struct BasicBlock {
pub mut:
	id     int
	name   string
	parent int
	instrs []ssa.ValueID
	preds  []ssa.BlockID
	succs  []ssa.BlockID
}

// Function represents function data used by mir.
pub struct Function {
pub mut:
	id          int
	name        string
	typ         ssa.TypeID
	blocks      []ssa.BlockID
	params      []ssa.ValueID
	is_c_extern bool
	abi         FunctionAbi
}

// Module represents module data used by mir.
@[heap]
pub struct Module {
pub mut:
	target     Target
	type_store ssa.TypeStore
	values     []Value
	instrs     []Instruction
	blocks     []BasicBlock
	funcs      []Function
	globals    []ssa.GlobalVar
}

// default_target returns a conservative target descriptor for target-neutral MIR.
pub fn default_target() Target {
	return Target{
		arch:         .unknown
		abi:          .unknown
		pointer_size: 8
		stack_align:  8
	}
}

// arm64_target returns the target descriptor used by the native ARM64 pipeline.
pub fn arm64_target() Target {
	return Target{
		arch:         .arm64
		abi:          .aapcs64
		pointer_size: 8
		int_arg_regs: 8
		ret_regs:     8
		stack_align:  16
	}
}

// lower_from_ssa lowers an SSA module into target-neutral MIR.
pub fn lower_from_ssa(m &ssa.Module) Module {
	return lower_from_ssa_for_target(m, default_target())
}

// lower_from_ssa_for_target lowers an SSA module into MIR with target ABI metadata.
pub fn lower_from_ssa_for_target(m &ssa.Module, target Target) Module {
	mut mod := Module{
		target:     target
		type_store: m.type_store
		values:     []Value{len: m.values.len}
		instrs:     []Instruction{len: m.instrs.len}
		blocks:     []BasicBlock{len: m.blocks.len}
		funcs:      []Function{len: m.funcs.len}
		globals:    m.globals
	}

	for i, val in m.values {
		mod.values[i] = Value{
			id:    val.id
			kind:  val.kind
			typ:   val.typ
			name:  val.name
			index: val.index
			uses:  val.uses
		}
	}

	for i, instr in m.instrs {
		mod.instrs[i] = Instruction{
			op:       instr.op
			operands: instr.operands
			typ:      instr.typ
			block:    instr.block
		}
	}

	for i, blk in m.blocks {
		mod.blocks[i] = BasicBlock{
			id:     blk.id
			name:   blk.name
			parent: blk.parent
			instrs: blk.instrs
			preds:  blk.preds
			succs:  blk.succs
		}
	}

	for i, f in m.funcs {
		mod.funcs[i] = Function{
			id:          f.id
			name:        f.name
			typ:         f.typ
			blocks:      f.blocks
			params:      f.params
			is_c_extern: f.is_c_extern
		}
	}

	for i, f in mod.funcs {
		mut func := f
		func.abi = mod.build_function_abi(func)
		mod.funcs[i] = func
	}

	return mod
}

// type_size returns type size data for Module.
pub fn (m &Module) type_size(typ_id ssa.TypeID) int {
	return m.type_size_inner(typ_id, 0)
}

// type_size_inner returns type size inner data for Module.
fn (m &Module) type_size_inner(typ_id ssa.TypeID, depth int) int {
	if typ_id <= 0 || typ_id >= m.type_store.types.len {
		return 0
	}
	if depth > 32 {
		return 8
	}
	typ := m.type_store.types[typ_id]
	if typ.width > 0 {
		return (typ.width + 7) / 8
	}
	if typ.elem_type > 0 && typ.fields.len == 0 {
		return 8
	}
	if typ.fields.len == 0 {
		if typ.params.len > 0 || typ.ret_type > 0 {
			return 8
		}
		return 0
	}
	if typ.fields.len > 256 {
		return 8
	}
	mut offset := 0
	mut max_align := 1
	for i in 0 .. typ.fields.len {
		field_typ := typ.fields[i]
		align := m.type_align_inner(field_typ, depth + 1)
		if align > max_align {
			max_align = align
		}
		if align > 1 && offset % align != 0 {
			offset = (offset + align - 1) & ~(align - 1)
		}
		offset += m.type_size_inner(field_typ, depth + 1)
	}
	total := if max_align > 1 && offset % max_align != 0 {
		(offset + max_align - 1) & ~(max_align - 1)
	} else {
		offset
	}
	if total > 0 {
		return total
	}
	return 8
}

// type_align returns type align data for Module.
pub fn (m &Module) type_align(typ_id ssa.TypeID) int {
	return m.type_align_inner(typ_id, 0)
}

// type_align_inner returns type align inner data for Module.
fn (m &Module) type_align_inner(typ_id ssa.TypeID, depth int) int {
	if typ_id <= 0 || typ_id >= m.type_store.types.len {
		return 1
	}
	if depth > 32 {
		return 8
	}
	typ := m.type_store.types[typ_id]
	if typ.width > 0 {
		size := (typ.width + 7) / 8
		if size >= 8 {
			return 8
		}
		if size >= 4 {
			return 4
		}
		return 1
	}
	if typ.elem_type > 0 && typ.fields.len == 0 {
		return 8
	}
	if typ.fields.len > 0 {
		if typ.fields.len > 256 {
			return 8
		}
		mut max_align := 1
		for i in 0 .. typ.fields.len {
			field_typ := typ.fields[i]
			a := m.type_align_inner(field_typ, depth + 1)
			if a > max_align {
				max_align = a
			}
		}
		return max_align
	}
	if typ.params.len > 0 || typ.ret_type > 0 {
		return 8
	}
	size := m.type_size_inner(typ_id, depth + 1)
	if size >= 8 {
		return 8
	}
	if size >= 4 {
		return 4
	}
	return 1
}

// build_function_abi builds function abi data for mir.
fn (m &Module) build_function_abi(f Function) FunctionAbi {
	mut params := []AbiLocation{}
	mut next_reg := 0
	mut next_stack := 0
	word_size := m.target_word_size()
	for pid in f.params {
		size := m.value_size(pid)
		n_words := words_for(size, word_size)
		if m.target.int_arg_regs > 0 && next_reg + n_words <= m.target.int_arg_regs {
			params << AbiLocation{
				kind:           .register
				register_index: next_reg
				register_count: n_words
				size:           size
			}
			next_reg += n_words
			continue
		}
		stack_offset := align_to(next_stack, word_size)
		params << AbiLocation{
			kind:           .stack
			register_count: n_words
			stack_offset:   stack_offset
			size:           size
		}
		next_stack = stack_offset + align_to(size, word_size)
	}

	ret_size := m.type_size(f.typ)
	ret_words := words_for(ret_size, word_size)
	ret := if ret_size == 0 {
		AbiLocation{
			kind: .none
		}
	} else if m.target.ret_regs > 0 && ret_words <= m.target.ret_regs {
		AbiLocation{
			kind:           .register
			register_index: 0
			register_count: ret_words
			size:           ret_size
		}
	} else {
		AbiLocation{
			kind:           .indirect
			register_index: 0
			register_count: 1
			size:           ret_size
		}
	}
	return FunctionAbi{
		params:          params
		ret:             ret
		stack_args_size: align_to(next_stack, m.target_stack_align())
		stack_align:     m.target_stack_align()
	}
}

// value_size returns value size data for Module.
fn (m &Module) value_size(val_id ssa.ValueID) int {
	if val_id <= 0 || val_id >= m.values.len {
		return m.target_word_size()
	}
	size := m.type_size(m.values[val_id].typ)
	if size > 0 {
		return size
	}
	return m.target_word_size()
}

// target_word_size supports target word size handling for Module.
fn (m &Module) target_word_size() int {
	if m.target.pointer_size > 0 {
		return m.target.pointer_size
	}
	return 8
}

// target_stack_align supports target stack align handling for Module.
fn (m &Module) target_stack_align() int {
	if m.target.stack_align > 0 {
		return m.target.stack_align
	}
	return m.target_word_size()
}

// words_for supports words for handling for mir.
fn words_for(size int, word_size int) int {
	if size <= 0 {
		return 0
	}
	return (size + word_size - 1) / word_size
}

// align_to supports align to handling for mir.
fn align_to(value int, alignment int) int {
	if alignment <= 1 {
		return value
	}
	return (value + alignment - 1) & ~(alignment - 1)
}