#include "builder.h"

#include <mie/ctx.h>
#include <mie/ir/branch.h>
#include <mie/select/graph.h>
#include <mie/select/node.h>

static enum mie_status get_block_node(
	struct mie_select_builder *builder, struct mie_block *block,
	struct mie_select_value *out)
{
	struct mie_select_node *block_node = NULL;
	struct mie_select_graph *graph = mie_select_builder_get_graph(builder);
	struct mie_type *ptr_type = mie_ctx_get_type(
		mie_select_builder_get_ctx(builder), MIE_TYPE_PTR);

	enum mie_status status = mie_select_graph_get_node(
		graph, mie_target_builtin(), MIE_SELECT_OP_BLOCK, NULL, 0,
		&ptr_type, 1, &block_node);
	if (status != MIE_SUCCESS) {
		return status;
	}

	block_node->n_value.v = MIE_VALUE(block);
	block_node->n_flags |= MIE_SELECT_NODE_F_PVALUE;

	return mie_select_node_get_value(block_node, ptr_type, 0, out);
}

static enum mie_status create_br_cond(
	struct mie_select_builder *builder, struct mie_value *cond,
	struct mie_select_value *incoming_chain, struct mie_block *dest_block,
	struct mie_select_value *out)
{
	struct mie_select_graph *graph = mie_select_builder_get_graph(builder);

	if (!incoming_chain || !incoming_chain->v_node) {
		incoming_chain = mie_select_builder_get_mem_access(builder, cond);
	}

	if (!incoming_chain || !incoming_chain->v_node) {
		incoming_chain = &graph->g_entry;
	}

	struct mie_select_value dest;
	enum mie_status status = get_block_node(builder, dest_block, &dest);
	if (status != MIE_SUCCESS) {
		return status;
	}

	struct mie_select_value *operands[] = {
		incoming_chain,
		mie_select_builder_get_value(builder, cond),
		&dest,
	};
	size_t nr_operands = sizeof operands / sizeof operands[0];

	struct mie_type *chain_type = mie_ctx_get_type(
		mie_select_builder_get_ctx(builder), MIE_TYPE_OTHER);

	struct mie_select_node *br_node = NULL;

	status = mie_select_graph_get_node(
		graph, mie_target_builtin(), MIE_SELECT_OP_BR_COND, operands,
		nr_operands, &chain_type, 1, &br_node);
	if (status != MIE_SUCCESS) {
		return status;
	}

	mie_select_node_get_value(br_node, chain_type, 0, out);

	return MIE_SUCCESS;
}

static enum mie_status create_br(
	struct mie_select_builder *builder, struct mie_block *dest_block,
	struct mie_select_value *incoming_chain, struct mie_select_value *out)
{
	struct mie_select_graph *graph = mie_select_builder_get_graph(builder);
	if (!incoming_chain || !incoming_chain->v_node) {
		incoming_chain = &graph->g_entry;
	}

	struct mie_select_value dest;
	enum mie_status status = get_block_node(builder, dest_block, &dest);
	if (status != MIE_SUCCESS) {
		return status;
	}

	struct mie_select_value *operands[] = {
		incoming_chain,
		&dest,
	};
	size_t nr_operands = sizeof operands / sizeof operands[0];

	struct mie_type *chain_type = mie_ctx_get_type(
		mie_select_builder_get_ctx(builder), MIE_TYPE_OTHER);

	struct mie_select_node *br_node = NULL;

	status = mie_select_graph_get_node(
		graph, mie_target_builtin(), MIE_SELECT_OP_BR, operands,
		nr_operands, &chain_type, 1, &br_node);
	if (status != MIE_SUCCESS) {
		return status;
	}

	mie_select_node_get_value(br_node, chain_type, 0, out);

	return MIE_SUCCESS;
}

static enum mie_status push_br(
	struct mie_select_builder *builder, struct mie_instr *instr)
{
	struct mie_branch *br = (struct mie_branch *)instr;

	struct mie_select_value incoming_chain = {};
	enum mie_status status = mie_select_builder_collapse_chain_ends(
		builder, &incoming_chain);
	if (status != MIE_SUCCESS) {
		return status;
	}

	struct mie_select_value br_result;
	status = create_br(builder, br->b_dest, &incoming_chain, &br_result);
	if (status != MIE_SUCCESS) {
		return status;
	}

	return mie_select_builder_set_value(builder, MIE_VALUE(instr), &br_result);
}

static enum mie_status push_br_if(
	struct mie_select_builder *builder, struct mie_instr *instr)
{
	struct mie_branch_if *br = (struct mie_branch_if *)instr;

	struct mie_select_value incoming_chain = {};
	enum mie_status status = mie_select_builder_collapse_chain_ends(
		builder, &incoming_chain);
	if (status != MIE_SUCCESS) {
		return status;
	}

	struct mie_select_value br_true_result, br_false_result;
	status = create_br_cond(
		builder, br->b_cond, &incoming_chain, br->b_true_block,
		&br_true_result);
	if (status != MIE_SUCCESS) {
		return status;
	}

	status = create_br(
		builder, br->b_false_block, &br_true_result, &br_false_result);

	if (status != MIE_SUCCESS) {
		return status;
	}

	return mie_select_builder_set_value(
		builder, MIE_VALUE(instr), &br_false_result);
}

struct select_instr_type select_br = {
	.i_push = push_br,
};

struct select_instr_type select_br_if = {
	.i_push = push_br_if,
};