package lower

import "../hir"
import "../ir"
import "../types"
import "core:fmt"
import "core:mem"

State :: struct {
	hir_module:   ^hir.Module,
	instructions: [dynamic]ir.Instruction,
	local_values: []int,
	local_slots:  []int,
	allocator:    mem.Allocator,
}

append_instruction :: proc(state: ^State, instruction: ir.Instruction) -> int {
	id := len(state.instructions)
	append(&state.instructions, instruction)
	return id
}

clone_args :: proc(values: []int, allocator: mem.Allocator) -> []int {
	result := make([]int, len(values), allocator)
	copy(result, values)
	return result
}

sentinel :: proc(value_type: types.Type) -> i64 {
	switch value_type.bits {
	case 8:  return -86
	case 16: return -21846
	case 32: return -1431655766
	case:    return -6148914691236517206
	}
}

lower_expr :: proc(state: ^State, expr_id: int) -> int {
	if expr_id < 0 || expr_id >= len(state.hir_module.exprs) {
		trap := append_instruction(state, ir.Instruction{
			op=.Trap,
			type=types.VOID,
			target=-1,
			a=-1,
			b=-1,
			diagnostic=-1,
		})
		_ = trap
		return append_instruction(state, ir.Instruction{
			op=.Const,
			type=types.I64,
			integer=sentinel(types.I64),
			target=-1,
			a=-1,
			b=-1,
			diagnostic=-1,
		})
	}
	expr := state.hir_module.exprs[expr_id]
	switch expr.kind {
	case .Invalid:
		append_instruction(state, ir.Instruction{
			op=.Trap,
			span=expr.span,
			type=types.VOID,
			target=-1,
			a=-1,
			b=-1,
			diagnostic=expr.diagnostic,
		})
		fallback := expr.type
		if !types.is_concrete_integer(fallback) {
			fallback = types.I64
		}
		return append_instruction(state, ir.Instruction{
			op=.Const,
			span=expr.span,
			type=fallback,
			integer=sentinel(fallback),
			target=-1,
			a=-1,
			b=-1,
			diagnostic=-1,
		})
	case .Integer:
		return append_instruction(state, ir.Instruction{
			op=.Const,
			span=expr.span,
			type=expr.type,
			integer=expr.integer,
			target=-1,
			a=-1,
			b=-1,
			diagnostic=-1,
		})
	case .Local:
		if expr.target >= 0 && expr.target < len(state.local_slots) && state.local_slots[expr.target] >= 0 {
			return append_instruction(state, ir.Instruction{
				op=.Load,
				span=expr.span,
				type=expr.type,
				target=-1,
				a=state.local_slots[expr.target],
				b=-1,
				diagnostic=-1,
			})
		}
		if expr.target >= 0 && expr.target < len(state.local_values) {
			return state.local_values[expr.target]
		}
	case .Global:
		return append_instruction(state, ir.Instruction{
			op=.Load_Global,
			span=expr.span,
			type=expr.type,
			target=expr.target,
			a=-1,
			b=-1,
			diagnostic=-1,
		})
	case .Widen:
		value := lower_expr(state, expr.left)
		return append_instruction(state, ir.Instruction{
			op=.Widen,
			span=expr.span,
			type=expr.type,
			target=-1,
			a=value,
			b=-1,
			diagnostic=-1,
		})
	case .Add:
		left := lower_expr(state, expr.left)
		right := lower_expr(state, expr.right)
		return append_instruction(state, ir.Instruction{
			op=.Add_Checked,
			span=expr.span,
			type=expr.type,
			target=-1,
			a=left,
			b=right,
			diagnostic=-1,
		})
	case .Call:
		args := make([]int, len(expr.args), state.allocator)
		for arg, index in expr.args {
			args[index] = lower_expr(state, arg)
		}
		return append_instruction(state, ir.Instruction{
			op=.Call,
			span=expr.span,
			type=expr.type,
			target=expr.target,
			a=-1,
			b=-1,
			args=args,
			diagnostic=-1,
		})
	}
	append_instruction(state, ir.Instruction{
		op=.Trap,
		span=expr.span,
		type=types.VOID,
		target=-1,
		a=-1,
		b=-1,
		diagnostic=expr.diagnostic,
	})
	return append_instruction(state, ir.Instruction{
		op=.Const,
		span=expr.span,
		type=types.I64,
		integer=sentinel(types.I64),
		target=-1,
		a=-1,
		b=-1,
		diagnostic=-1,
	})
}

lower_body :: proc(hir_module: ^hir.Module, function: hir.Function, allocator: mem.Allocator) -> []ir.Instruction {
	state := State{
		hir_module=hir_module,
		allocator=allocator,
		local_values=make([]int, len(function.locals), allocator),
		local_slots=make([]int, len(function.locals), allocator),
	}
	state.instructions.allocator = allocator
	defer {
		delete(state.local_values, allocator)
		delete(state.local_slots, allocator)
	}
	for _, index in state.local_values {
		state.local_values[index] = -1
		state.local_slots[index] = -1
	}
	for local_id in function.params {
		param := append_instruction(&state, ir.Instruction{
			op=.Param,
			type=function.locals[local_id].type,
			target=local_id,
			a=-1,
			b=-1,
			diagnostic=-1,
		})
		state.local_values[local_id] = param
	}

	for statement_id in function.body {
		statement := hir_module.statements[statement_id]
		switch statement.kind {
		case .Declaration:
			value := lower_expr(&state, statement.expr)
			local := function.locals[statement.local]
			if local.mutable {
				slot := append_instruction(&state, ir.Instruction{
					op=.Alloca,
					span=statement.span,
					type=local.type,
					target=statement.local,
					a=-1,
					b=-1,
					diagnostic=-1,
				})
				state.local_slots[statement.local] = slot
				append_instruction(&state, ir.Instruction{
					op=.Store,
					span=statement.span,
					type=local.type,
					target=-1,
					a=slot,
					b=value,
					diagnostic=-1,
				})
			} else {
				state.local_values[statement.local] = value
			}
		case .Assignment:
			value := lower_expr(&state, statement.expr)
			slot := -1
			if statement.local >= 0 && statement.local < len(state.local_slots) {
				slot = state.local_slots[statement.local]
			}
			append_instruction(&state, ir.Instruction{
				op=.Store,
				span=statement.span,
				type=function.locals[statement.local].type,
				target=-1,
				a=slot,
				b=value,
				diagnostic=-1,
			})
		case .Return:
			if statement.expr < 0 {
				append_instruction(&state, ir.Instruction{
					op=.Return_Void,
					span=statement.span,
					type=types.VOID,
					target=-1,
					a=-1,
					b=-1,
					diagnostic=-1,
				})
			} else {
				value := lower_expr(&state, statement.expr)
				append_instruction(&state, ir.Instruction{
					op=.Return,
					span=statement.span,
					type=function.result,
					target=-1,
					a=value,
					b=-1,
					diagnostic=-1,
				})
			}
		case .Expression, .Sink:
			_ = lower_expr(&state, statement.expr)
		case .Trap:
			append_instruction(&state, ir.Instruction{
				op=.Trap,
				span=statement.span,
				type=types.VOID,
				target=-1,
				a=-1,
				b=-1,
				diagnostic=statement.diagnostic,
			})
		}
	}
	if len(state.instructions) == 0 ||
	   (state.instructions[len(state.instructions)-1].op != .Return &&
	    state.instructions[len(state.instructions)-1].op != .Return_Void) {
		if function.result.kind == .Void {
			append_instruction(&state, ir.Instruction{op=.Return_Void, type=types.VOID, target=-1, a=-1, b=-1, diagnostic=-1})
		} else {
			value := append_instruction(&state, ir.Instruction{
				op=.Const,
				type=function.result,
				integer=sentinel(function.result),
				target=-1,
				a=-1,
				b=-1,
				diagnostic=-1,
			})
			append_instruction(&state, ir.Instruction{op=.Return, type=function.result, target=-1, a=value, b=-1, diagnostic=-1})
		}
	}
	return state.instructions[:]
}

lower_global_initializer :: proc(hir_module: ^hir.Module, global: hir.Global, allocator: mem.Allocator) -> []ir.Instruction {
	state := State{hir_module=hir_module, allocator=allocator}
	state.instructions.allocator = allocator
	value := lower_expr(&state, global.expr)
	append_instruction(&state, ir.Instruction{
		op=.Return,
		type=global.type,
		target=-1,
		a=value,
		b=-1,
		diagnostic=-1,
	})
	return state.instructions[:]
}

lower :: proc(hir_module: ^hir.Module, allocator := context.allocator) -> ir.Module {
	module := ir.init_module(allocator)
	for global in hir_module.globals {
		append(&module.globals, ir.Global{
			name=global.name,
			type=global.type,
			is_static=global.is_static,
			static_value=global.static_value,
			initializer=nil if global.is_static else lower_global_initializer(hir_module, global, allocator),
			problematic=global.problematic,
			diagnostic=global.diagnostic,
		})
	}
	for function in hir_module.functions {
		param_types := make([]types.Type, len(function.params), allocator)
		for local_id, index in function.params {
			param_types[index] = function.locals[local_id].type
		}
		append(&module.functions, ir.Function{
			link_name=fmt.aprintf("%s", function.link_name, allocator=allocator),
			c_abi=function.c_abi,
			is_main=function.is_main,
			param_types=param_types,
			result=function.result,
			instructions=lower_body(hir_module, function, allocator),
			problematic=function.problematic,
		})
	}
	return module
}