mie/mie/ir/walk.c

#include <mie/ir/block.h>
#include <mie/ir/op.h>
#include <mie/ir/region.h>
#include <mie/ir/walk.h>

#define ITEM_TYPE(f) ((f) & 0x0Fu)

enum walk_schedule_item_flags {
	SCHED_ITEM_F_NONE = 0,
	SCHED_ITEM_F_OP = 0x01u,
	SCHED_ITEM_F_BLOCK = 0x02u,
	SCHED_ITEM_F_REGION = 0x04u,

	SCHED_ITEM_F_VISITED = 0x10u,
	SCHED_ITEM_F_CHILDREN_SCHEDULED = 0x20u,
};

struct walk_schedule_item {
	enum walk_schedule_item_flags i_flags;
	b_queue_entry i_entry;
	size_t i_depth;

	union {
		struct mie_op *i_op;
		struct mie_block *i_block;
		struct mie_region *i_region;
	};
};

struct mie_walker {
	enum mie_walker_flags w_flags;
	struct mie_op *w_root;
	struct mie_walker_item w_cur;
	b_queue w_sched;
};

static struct walk_schedule_item *op_schedule_item_create(struct mie_op *op)
{
	struct walk_schedule_item *out = malloc(sizeof *out);
	if (!out) {
		return NULL;
	}

	memset(out, 0x0, sizeof *out);

	out->i_flags = SCHED_ITEM_F_OP;
	out->i_op = op;

	return out;
}

static struct walk_schedule_item *block_schedule_item_create(struct mie_block *block)
{
	struct walk_schedule_item *out = malloc(sizeof *out);
	if (!out) {
		return NULL;
	}

	memset(out, 0x0, sizeof *out);

	out->i_flags = SCHED_ITEM_F_BLOCK;
	out->i_block = block;

	return out;
}

static struct walk_schedule_item *region_schedule_item_create(
	struct mie_region *region)
{
	struct walk_schedule_item *out = malloc(sizeof *out);
	if (!out) {
		return NULL;
	}

	memset(out, 0x0, sizeof *out);

	out->i_flags = SCHED_ITEM_F_REGION;
	out->i_region = region;

	return out;
}

static bool should_ignore_item(
	const struct mie_walker *walker, const struct walk_schedule_item *item)
{
	switch (ITEM_TYPE(item->i_flags)) {
	case SCHED_ITEM_F_OP:
		return (walker->w_flags & MIE_WALKER_F_INCLUDE_OPS) == 0;
	case SCHED_ITEM_F_BLOCK:
		return (walker->w_flags & MIE_WALKER_F_INCLUDE_BLOCKS) == 0;
	case SCHED_ITEM_F_REGION:
		return (walker->w_flags & MIE_WALKER_F_INCLUDE_REGIONS) == 0;
	default:
		return true;
	}
}

static void schedule_child(
	struct mie_walker *walker, struct walk_schedule_item *parent,
	struct walk_schedule_item *child, b_queue_entry **ep)
{
#define REVERSE 0x04u

	enum {
		PREORDER = 0x01u,
		PREORDER_REVERSE = PREORDER | REVERSE,
		POSTORDER = 0x02u,
		POSTORDER_REVERSE = POSTORDER | REVERSE,
	} mode;

	if (walker->w_flags & MIE_WALKER_F_POSTORDER) {
		mode = POSTORDER;
	} else {
		mode = PREORDER;
	}

	if (walker->w_flags & MIE_WALKER_F_BACKWARD) {
		mode |= REVERSE;
	}

	child->i_depth = parent->i_depth;
	if (!should_ignore_item(walker, child)) {
		child->i_depth++;
	}

	switch (mode) {
	case PREORDER:
		b_queue_insert_after(&walker->w_sched, &child->i_entry, *ep);
		*ep = &child->i_entry;
		break;
	case PREORDER_REVERSE:
		b_queue_insert_after(
			&walker->w_sched, &child->i_entry, &parent->i_entry);
		*ep = &parent->i_entry;
		break;
	case POSTORDER:
		if (*ep == &parent->i_entry) {
			b_queue_insert_before(
				&walker->w_sched, &child->i_entry, *ep);
		} else {
			b_queue_insert_after(
				&walker->w_sched, &child->i_entry, *ep);
		}

		*ep = &child->i_entry;
		break;
	case POSTORDER_REVERSE:
		b_queue_insert_before(&walker->w_sched, &child->i_entry, *ep);
		*ep = &child->i_entry;
		break;
	default:
		return;
	}

#undef REVERSE
}

static enum mie_status schedule_children_of_region(
	struct mie_walker *walker, struct walk_schedule_item *item)
{
	struct mie_region *region = item->i_region;

	b_queue_entry *tmp = &item->i_entry;

	for (size_t i = 0; i < MIE_VECTOR_COUNT(region->r_blocks); i++) {
		struct mie_block *block = region->r_blocks.items[i];
		struct walk_schedule_item *child
			= block_schedule_item_create(block);

		schedule_child(walker, item, child, &tmp);
	}

	return MIE_SUCCESS;
}

static enum mie_status schedule_children_of_block(
	struct mie_walker *walker, struct walk_schedule_item *item)
{
	struct mie_block *block = item->i_block;

	b_queue_entry *tmp = &item->i_entry;

	for (size_t i = 0; i < MIE_VECTOR_COUNT(block->b_ops); i++) {
		struct mie_op *op = &block->b_ops.items[i];
		struct walk_schedule_item *child = op_schedule_item_create(op);

		schedule_child(walker, item, child, &tmp);
	}

	return MIE_SUCCESS;
}

static enum mie_status schedule_children_of_op(
	struct mie_walker *walker, struct walk_schedule_item *item)
{
	struct mie_op *op = item->i_op;

	b_queue_entry *tmp = &item->i_entry;

	for (size_t i = 0; i < MIE_VECTOR_COUNT(op->op_regions); i++) {
		struct mie_region *region = &op->op_regions.items[i];
		struct walk_schedule_item *child
			= region_schedule_item_create(region);

		schedule_child(walker, item, child, &tmp);
	}

	return MIE_SUCCESS;
}

static enum mie_status schedule_children(
	struct mie_walker *walker, struct walk_schedule_item *item)
{
	if (item->i_flags & SCHED_ITEM_F_CHILDREN_SCHEDULED) {
		return MIE_SUCCESS;
	}

	enum mie_status status;

	switch (ITEM_TYPE(item->i_flags)) {
	case SCHED_ITEM_F_BLOCK:
		status = schedule_children_of_block(walker, item);
		break;
	case SCHED_ITEM_F_REGION:
		status = schedule_children_of_region(walker, item);
		break;
	case SCHED_ITEM_F_OP:
		status = schedule_children_of_op(walker, item);
		break;
	default:
		status = MIE_ERR_BAD_STATE;
		break;
	}

	item->i_flags |= SCHED_ITEM_F_CHILDREN_SCHEDULED;
	return MIE_SUCCESS;
}

static enum mie_status prepare_postorder_traversal(struct mie_walker *walker)
{
	b_queue_entry *cur = b_queue_last(&walker->w_sched);
	while (cur) {
		struct walk_schedule_item *item
			= b_unbox(struct walk_schedule_item, cur, i_entry);
		schedule_children(walker, item);

		cur = b_queue_prev(cur);
	}

	return MIE_SUCCESS;
}

struct mie_walker *mie_walker_begin(struct mie_op *op, enum mie_walker_flags flags)
{
	struct mie_walker *out = malloc(sizeof *out);
	if (!out) {
		return NULL;
	}

	memset(out, 0x0, sizeof *out);

	out->w_flags = flags;
	out->w_root = op;

	struct walk_schedule_item *item = op_schedule_item_create(op);
	b_queue_push_back(&out->w_sched, &item->i_entry);

	if (flags & MIE_WALKER_F_POSTORDER) {
		prepare_postorder_traversal(out);
	}

	return out;
}

void mie_walker_end(struct mie_walker *walker)
{
}

static struct walk_schedule_item *current_item(struct mie_walker *walker)
{
	b_queue_entry *top = b_queue_first(&walker->w_sched);
	if (!top) {
		return NULL;
	}

	return b_unbox(struct walk_schedule_item, top, i_entry);
}

static enum mie_status step(struct mie_walker *walker)
{
	struct walk_schedule_item *cur = current_item(walker);
	if (!cur) {
		return MIE_ERR_NO_DATA;
	}

	schedule_children(walker, cur);

	b_queue_delete(&walker->w_sched, &cur->i_entry);
	free(cur);

	if (!current_item(walker)) {
		return MIE_ERR_NO_DATA;
	}

	return MIE_SUCCESS;
}

enum mie_status mie_walker_step(struct mie_walker *walker)
{
	struct walk_schedule_item *cur = current_item(walker);
	do {
		step(walker);
		cur = current_item(walker);
	} while (cur && should_ignore_item(walker, cur));

	return cur ? MIE_SUCCESS : MIE_ERR_NO_DATA;
}

struct mie_walker_item *mie_walker_get(struct mie_walker *walker)
{
	struct walk_schedule_item *cur = current_item(walker);
	if (!cur) {
		return NULL;
	}

	memset(&walker->w_cur, 0x0, sizeof walker->w_cur);

	walker->w_cur.i_depth = cur->i_depth;

	switch (ITEM_TYPE(cur->i_flags)) {
	case SCHED_ITEM_F_OP:
		walker->w_cur.i_op = cur->i_op;
		break;
	case SCHED_ITEM_F_BLOCK:
		walker->w_cur.i_block = cur->i_block;
		break;
	case SCHED_ITEM_F_REGION:
		walker->w_cur.i_region = cur->i_region;
		break;
	default:
		return NULL;
	}

	return &walker->w_cur;
}