mie/mie/rewrite/rewriter.c

#include <mie/ctx.h>
#include <mie/ir/block.h>
#include <mie/ir/op-definition.h>
#include <mie/ir/op.h>
#include <mie/ir/region.h>
#include <mie/rewrite/rewriter.h>
#include <string.h>

void mie_rewriter_init(struct mie_rewriter *rw, struct mie_ctx *ctx)
{
	memset(rw, 0x0, sizeof *rw);
	rw->r_ctx = ctx;

	rw->r_base.e_get_ctx = (mie_emit_get_ctx)mie_rewriter_get_ctx;
	rw->r_base.e_put_op = (mie_emit_put_op)mie_rewriter_put_op;
	rw->r_base.e_put_name = (mie_emit_put_name)mie_rewriter_put_name;
	rw->r_base.e_put_block = NULL;
}

void mie_rewriter_cleanup(struct mie_rewriter *rw)
{
	/* TODO */
}

struct mie_ctx *mie_rewriter_get_ctx(const struct mie_rewriter *rw)
{
	return rw->r_ctx;
}

struct mie_block *mie_rewriter_get_insertion_block(struct mie_rewriter *rw)
{
	return rw->r_insert_block;
}

struct mie_op *mie_rewriter_get_insertion_point(struct mie_rewriter *rw)
{
	return rw->r_insert_point;
}

struct mie_block *mie_rewriter_set_insertion_block(
	struct mie_rewriter *rw, struct mie_block *block)
{
	rw->r_insert_block = block;
	rw->r_insert_point = NULL;
	return block;
}

struct mie_op *mie_rewriter_set_insertion_point(
	struct mie_rewriter *rw, struct mie_op *insert_point)
{
	if (!insert_point && rw->r_insert_point) {
		insert_point = mie_block_get_last_op(rw->r_insert_block);
	}

	rw->r_insert_point = insert_point;
	return rw->r_insert_point;
}

static struct mie_name_map *get_name_map_for_region(struct mie_region *region)
{
	while (region) {
		if (region->r_names) {
			return region->r_names;
		}

		struct mie_op *op = region->r_parent;
		if (!op) {
			break;
		}

		struct mie_block *block = op->op_container;
		if (!block) {
			break;
		}

		region = block->b_parent;
	}

	return NULL;
}

struct mie_block *mie_rewriter_split_block(
	struct mie_rewriter *rw, struct mie_block *block, struct mie_op *before,
	const char *name)
{
	if (before->op_container != block) {
		return NULL;
	}

	struct mie_region *region = block->b_parent;
	struct mie_name_map *names = get_name_map_for_region(region);

	struct mie_block *new_block = mie_region_add_block_after(region, block);
	mie_name_map_put(names, &new_block->b_name, name, 0);

	fx_queue_entry *cur = &before->op_entry;
	while (cur) {
		struct mie_op *cur_op = fx_unbox(struct mie_op, cur, op_entry);
		fx_queue_entry *next = fx_queue_next(cur);
		fx_queue_delete(&block->b_ops, cur);
		fx_queue_push_back(&new_block->b_ops, cur);

		cur_op->op_container = new_block;
		cur = next;
	}

	return new_block;
}

struct mie_block *mie_rewriter_create_block(
	struct mie_rewriter *rw, struct mie_block *insert_before, const char *name)
{
	struct mie_region *parent = insert_before->b_parent;
	struct mie_block *block
		= mie_region_add_block_before(parent, insert_before);
	mie_rewriter_put_name(rw, &block->b_name, name);
	return block;
}

struct mie_register *mie_rewriter_add_block_parameter(
	struct mie_rewriter *rw, struct mie_block *block, const char *name,
	const struct mie_type *type)
{
	struct mie_register *reg = mie_block_add_param(block);

	if (mie_rewriter_put_name(rw, &reg->reg_name, name) != MIE_SUCCESS) {
		return NULL;
	}

	reg->reg_flags = MIE_REGISTER_F_VIRTUAL | MIE_REGISTER_F_BLOCK_PARAM;
	reg->reg_block = block;
	reg->reg_type = type;

	return reg;
}

enum mie_status mie_rewriter_move_block_to_start(
	struct mie_rewriter *rw, struct mie_block *block,
	struct mie_region *from, struct mie_region *to)
{
	if (block->b_parent != from) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	fx_queue_delete(&from->r_blocks, &block->b_entry);
	fx_queue_push_front(&to->r_blocks, &block->b_entry);

	block->b_parent = to;

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_move_block_to_end(
	struct mie_rewriter *rw, struct mie_block *block,
	struct mie_region *from, struct mie_region *to)
{
	if (block->b_parent != from) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	fx_queue_delete(&from->r_blocks, &block->b_entry);
	fx_queue_push_back(&to->r_blocks, &block->b_entry);

	block->b_parent = to;

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_move_block_before(
	struct mie_rewriter *rw, struct mie_block *block,
	struct mie_region *from, struct mie_region *to, struct mie_block *before)
{
	if (block->b_parent != from) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	if (before->b_parent != to) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	fx_queue_delete(&from->r_blocks, &block->b_entry);
	fx_queue_insert_before(&to->r_blocks, &block->b_entry, &before->b_entry);

	block->b_parent = to;

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_move_block_after(
	struct mie_rewriter *rw, struct mie_block *block,
	struct mie_region *from, struct mie_region *to, struct mie_block *after)
{
	if (block->b_parent != from) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	if (after->b_parent != to) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	fx_queue_delete(&from->r_blocks, &block->b_entry);
	fx_queue_insert_after(&to->r_blocks, &block->b_entry, &after->b_entry);

	block->b_parent = to;

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_rename_block(
	struct mie_rewriter *rw, struct mie_block *block, const char *name)
{
	struct mie_region *region = block->b_parent;
	struct mie_name_map *names = get_name_map_for_region(region);

	mie_name_destroy(&block->b_name);
	struct mie_name *result = mie_name_map_put(names, &block->b_name, name, 0);

	return result ? MIE_SUCCESS : MIE_ERR_NO_MEMORY;
}

enum mie_status mie_rewriter_move_blocks_to_start(
	struct mie_rewriter *rw, struct mie_region *from, struct mie_region *to)
{
	fx_queue_entry *cur = fx_queue_first(&from->r_blocks);
	fx_queue_entry *insert_point = NULL;
	while (cur) {
		fx_queue_entry *next = fx_queue_next(cur);

		struct mie_block *block = fx_unbox(struct mie_block, cur, b_entry);
		fx_queue_delete(&from->r_blocks, cur);

		if (insert_point) {
			fx_queue_insert_after(&to->r_blocks, cur, insert_point);
		} else {
			fx_queue_push_front(&to->r_blocks, cur);
		}

		block->b_parent = to;
		insert_point = cur;
		cur = next;
	}

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_move_blocks_to_end(
	struct mie_rewriter *rw, struct mie_region *from, struct mie_region *to)
{
	fx_queue_entry *cur = fx_queue_first(&from->r_blocks);
	while (cur) {
		fx_queue_entry *next = fx_queue_next(cur);

		struct mie_block *block = fx_unbox(struct mie_block, cur, b_entry);
		fx_queue_delete(&from->r_blocks, cur);
		fx_queue_push_back(&to->r_blocks, cur);

		block->b_parent = to;
		cur = next;
	}

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_move_blocks_before(
	struct mie_rewriter *rw, struct mie_region *from, struct mie_region *to,
	struct mie_block *before)
{
	if (before->b_parent != to) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	fx_queue_entry *cur = fx_queue_last(&from->r_blocks);
	while (cur) {
		fx_queue_entry *next = fx_queue_prev(cur);

		struct mie_block *block = fx_unbox(struct mie_block, cur, b_entry);
		fx_queue_delete(&from->r_blocks, cur);
		fx_queue_insert_before(&to->r_blocks, cur, &before->b_entry);

		block->b_parent = to;
		cur = next;
	}

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_move_blocks_after(
	struct mie_rewriter *rw, struct mie_region *from, struct mie_region *to,
	struct mie_block *after)
{
	if (after->b_parent != to) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	fx_queue_entry *cur = fx_queue_first(&from->r_blocks);
	fx_queue_entry *insert_point = &after->b_entry;

	while (cur) {
		fx_queue_entry *next = fx_queue_next(cur);

		struct mie_block *block = fx_unbox(struct mie_block, cur, b_entry);
		fx_queue_delete(&from->r_blocks, cur);
		fx_queue_insert_after(&to->r_blocks, cur, insert_point);

		block->b_parent = to;
		insert_point = cur;
		cur = next;
	}

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_move_region_to_start(
	struct mie_rewriter *rw, struct mie_region *region, struct mie_op *from,
	struct mie_op *to)
{
	if (region->r_parent != from) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	fx_queue_delete(&from->op_regions, &region->r_entry);
	fx_queue_push_front(&to->op_regions, &region->r_entry);

	region->r_parent = to;

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_move_region_to_end(
	struct mie_rewriter *rw, struct mie_region *region, struct mie_op *from,
	struct mie_op *to)
{
	if (region->r_parent != from) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	fx_queue_delete(&from->op_regions, &region->r_entry);
	fx_queue_push_back(&to->op_regions, &region->r_entry);

	region->r_parent = to;

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_move_region_before(
	struct mie_rewriter *rw, struct mie_region *region, struct mie_op *from,
	struct mie_op *to, struct mie_region *before)
{
	if (region->r_parent != from) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	if (before->r_parent != to) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	fx_queue_delete(&from->op_regions, &region->r_entry);
	fx_queue_insert_before(&to->op_regions, &region->r_entry, &before->r_entry);

	region->r_parent = to;

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_move_region_after(
	struct mie_rewriter *rw, struct mie_region *region, struct mie_op *from,
	struct mie_op *to, struct mie_region *after)
{
	if (region->r_parent != from) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	if (after->r_parent != to) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	fx_queue_delete(&from->op_regions, &region->r_entry);
	fx_queue_insert_after(&to->op_regions, &region->r_entry, &after->r_entry);

	region->r_parent = to;

	return MIE_SUCCESS;
}

struct mie_op *mie_rewriter_put_op(
	struct mie_rewriter *rw, const char *dialect_name, const char *op_name,
	struct mie_register **args, size_t nr_args)
{
	if (!rw->r_insert_block) {
		return NULL;
	}

	const struct mie_op_definition *op_def
		= mie_ctx_get_op_definition(rw->r_ctx, dialect_name, op_name);
	if (!op_def) {
		return NULL;
	}

	struct mie_op *op
		= mie_block_add_op_after(rw->r_insert_block, rw->r_insert_point);
	if (!op) {
		return NULL;
	}

	op->op_flags = MIE_OP_F_OP_RESOLVED;
	op->op_info = op_def;

	for (size_t i = 0; i < nr_args; i++) {
		struct mie_op_arg *arg = mie_op_add_arg(op);
		arg->arg_flags = MIE_OP_F_ARG_RESOLVED;
		arg->arg_value.u_reg = args[i];
		arg->arg_value.u_user = op;
		fx_queue_push_back(&args[i]->reg_use, &arg->arg_value.u_entry);
	}

	return op;
}

struct mie_op *mie_rewriter_replace_op(
	struct mie_rewriter *rw, struct mie_op *op, const char *dialect_name,
	const char *op_name)
{
	const struct mie_op_definition *op_def
		= mie_ctx_get_op_definition(rw->r_ctx, dialect_name, op_name);

	if (!op_def) {
		return NULL;
	}

	op->op_info = op_def;

	return op;
}

struct mie_op_arg *mie_rewriter_add_op_arg(
	struct mie_rewriter *rw, struct mie_op *op, struct mie_register *value)
{
	struct mie_op_arg *arg = mie_op_add_arg(op);
	arg->arg_flags = MIE_OP_F_ARG_RESOLVED;
	arg->arg_value.u_reg = value;
	arg->arg_value.u_user = op;
	fx_queue_push_back(&value->reg_use, &arg->arg_value.u_entry);

	return arg;
}

MIE_API struct mie_op_arg *mie_rewriter_add_op_successor_arg(
	struct mie_rewriter *rw, struct mie_op *op, struct mie_op_successor *s,
	struct mie_register *value)
{
	struct mie_op_arg *arg = mie_vector_emplace_back(s->s_args, NULL);
	arg->arg_flags = MIE_OP_F_ARG_RESOLVED;
	arg->arg_value.u_reg = value;
	arg->arg_value.u_user = op;
	fx_queue_push_back(&value->reg_use, &arg->arg_value.u_entry);

	return arg;
}

struct mie_op_successor *mie_rewriter_add_op_successor(
	struct mie_rewriter *rw, struct mie_op *op, struct mie_block *dest,
	struct mie_register **args, size_t nr_args)
{
	return mie_op_add_successor(op, dest, args, nr_args);
}

enum mie_status mie_rewriter_erase_op(struct mie_rewriter *rw, struct mie_op *op)
{
	struct mie_block *parent = op->op_container;
	if (!parent) {
		return MIE_ERR_BAD_STATE;
	}

	fx_queue_delete(&parent->b_ops, &op->op_entry);

	for (size_t i = 0; i < MIE_VECTOR_COUNT(op->op_successors); i++) {
		struct mie_op_successor *s = &op->op_successors.items[i];
		for (size_t k = 0; k < MIE_VECTOR_COUNT(s->s_args); k++) {
			struct mie_op_arg *arg = &s->s_args.items[k];
			if (!(arg->arg_flags & MIE_OP_F_ARG_RESOLVED)) {
				continue;
			}

			fx_queue_delete(
				&arg->arg_value.u_reg->reg_use,
				&arg->arg_value.u_entry);
		}
	}

	for (size_t k = 0; k < MIE_VECTOR_COUNT(op->op_args); k++) {
		struct mie_op_arg *arg = &op->op_args.items[k];
		if (!(arg->arg_flags & MIE_OP_F_ARG_RESOLVED)) {
			continue;
		}

		fx_queue_delete(
			&arg->arg_value.u_reg->reg_use, &arg->arg_value.u_entry);
	}

	for (size_t i = 0; i < MIE_VECTOR_COUNT(op->op_result); i++) {
		struct mie_register *reg = &op->op_result.items[i];

		fx_queue_entry *cur = fx_queue_first(&reg->reg_use);
		while (cur) {
			fx_queue_entry *next = fx_queue_next(cur);

			struct mie_register_use *use
				= fx_unbox(struct mie_register_use, cur, u_entry);

			fx_queue_delete(&reg->reg_use, &use->u_entry);
			use->u_reg = NULL;
			cur = next;
		}
	}

	mie_attribute_map_cleanup(&op->op_attrib);

	free(op);
	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_move_op_args_to_successor(
	struct mie_rewriter *rw, struct mie_op *op, struct mie_op_successor *s)
{
	for (size_t i = 0; i < MIE_VECTOR_COUNT(op->op_args); i++) {
		struct mie_op_arg *src = &op->op_args.items[i];
		struct mie_op_arg *dest = mie_vector_emplace_back(s->s_args, NULL);

		struct mie_register *reg = NULL;
		if (src->arg_flags & MIE_OP_F_ARG_RESOLVED) {
			reg = src->arg_value.u_reg;
			fx_queue_delete(&reg->reg_use, &src->arg_value.u_entry);
		}

		memcpy(dest, src, sizeof *src);
		if (reg) {
			memset(&dest->arg_value.u_entry, 0x0,
			       sizeof dest->arg_value.u_entry);
			fx_queue_push_back(&reg->reg_use, &dest->arg_value.u_entry);
		}
	}

	mie_vector_destroy(op->op_args, NULL);

	return MIE_SUCCESS;
}

static enum mie_status replace_register_usage(
	struct mie_op *op, struct mie_register *old, struct mie_register *new)
{
	for (size_t i = 0; i < MIE_VECTOR_COUNT(op->op_args); i++) {
		struct mie_op_arg *arg = &op->op_args.items[i];
		if (!(arg->arg_flags & MIE_OP_F_ARG_RESOLVED)) {
			continue;
		}

		if (arg->arg_value.u_reg != old) {
			continue;
		}

		fx_queue_delete(&old->reg_use, &arg->arg_value.u_entry);
		arg->arg_value.u_reg = new;
		fx_queue_push_back(&new->reg_use, &arg->arg_value.u_entry);
	}

	for (size_t i = 0; i < MIE_VECTOR_COUNT(op->op_successors); i++) {
		struct mie_op_successor *s = &op->op_successors.items[i];
		for (size_t k = 0; k < MIE_VECTOR_COUNT(s->s_args); k++) {
			struct mie_op_arg *arg = &s->s_args.items[i];
			if (!(arg->arg_flags & MIE_OP_F_ARG_RESOLVED)) {
				continue;
			}

			if (arg->arg_value.u_reg != old) {
				continue;
			}

			fx_queue_delete(&old->reg_use, &arg->arg_value.u_entry);
			arg->arg_value.u_reg = new;
			fx_queue_push_back(&new->reg_use, &arg->arg_value.u_entry);
		}
	}

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_replace_register(
	struct mie_rewriter *rw, struct mie_register *old, struct mie_register *new)
{
	fx_queue_entry *cur = fx_queue_first(&old->reg_use);
	while (cur) {
		struct mie_register_use *use
			= fx_unbox(struct mie_register_use, cur, u_entry);
		struct mie_op *op = use->u_user;

		if (use->u_reg == old) {
			replace_register_usage(op, old, new);
		}

		cur = fx_queue_next(cur);
	}

	return MIE_SUCCESS;
}

enum mie_status mie_rewriter_rename_register(
	struct mie_rewriter *rw, struct mie_register *reg, const char *name)
{
	if (!reg->reg_block) {
		return MIE_ERR_INVALID_ARGUMENT;
	}

	struct mie_region *region = reg->reg_block->b_parent;
	struct mie_name_map *names = get_name_map_for_region(region);

	mie_name_destroy(&reg->reg_name);
	struct mie_name *result = mie_name_map_put(names, &reg->reg_name, name, 0);

	return result ? MIE_SUCCESS : MIE_ERR_NO_MEMORY;
}

enum mie_status mie_rewriter_put_name(
	struct mie_rewriter *rw, struct mie_name *name, const char *hint)
{
	if (!rw->r_insert_block) {
		return MIE_ERR_BAD_STATE;
	}

	struct mie_region *region = rw->r_insert_block->b_parent;
	struct mie_name_map *names = get_name_map_for_region(region);

	mie_name_destroy(name);
	struct mie_name *result = mie_name_map_put(names, name, hint, 0);

	return result ? MIE_SUCCESS : MIE_ERR_NO_MEMORY;
}