#include "wasm_compile.h"
#include "wasm.h"
#include "stack_machine_instruction.h"

/* instruction structs are connected in a circular list */
void free_expr(struct instruction *expr)
{
	struct instruction *tmp;

	if (expr) {
		tmp = expr->next;
		while (tmp != expr) {
			tmp = tmp->next;
			free(tmp->prev);
		}

		free(expr);
	}
}

/* instructions are stored in a circular list */
int add_instruction(struct instruction **expr, uint16_t encoding,
		    struct instruction_data data)
{
	struct instruction *new;

	new = malloc(sizeof(struct instruction));

	if (!new) {
		PRERR(MSG_ALLOC_FAIL(sizeof(struct instruction)));
		return -1;
	}

	new->address_assigned = false;
	new->encoding = encoding;
	new->data = data;

	if (!*expr) {
		new->next = new;
		new->prev = new;
		*expr = new;
	} else {
		new->next = *expr;
		new->prev = (*expr)->prev;
		(*expr)->prev->next = new;
		(*expr)->prev = new;
	}

	return 0;
}

static uint8_t estimate_instruction_size(struct instruction *instruction)
{
	switch (instruction->data.info) {
	case DATA_NONE :
		return 2;
	case DATA_KNOWN :
	case DATA_INSTR_ADDR :
	case DATA_ADDR_AFTER :
		return 6;
	case DATA_KNOWN_21_BITS :
		return 4;
	default /* case DATA_KNOWN_6_BITS */ :
		return 2;
	}
}

static uint64_t estimate_expr_size(struct instruction *expr)
{
	struct instruction *tmp = expr;
	uint64_t max_expr_size = 0;

	do {
		max_expr_size += estimate_instruction_size(tmp);
		tmp = tmp->next;
	} while (tmp != expr);

	return max_expr_size;
}

static void assign_addresses_and_sizes(struct instruction *expr,
				       uint32_t address)
{
	struct instruction *tmp = expr;
	uint32_t target_addr;

	do {
		if ((tmp->data.info == DATA_INSTR_ADDR ||
		     tmp->data.info == DATA_ADDR_AFTER) &&
		    *tmp->data.data.ptr &&
		    (*tmp->data.data.ptr)->address_assigned) {
			target_addr = (*tmp->data.data.ptr)->address;

			if (tmp->data.info == DATA_ADDR_AFTER) {
				target_addr += estimate_instruction_size
					(*tmp->data.data.ptr);
			}

			tmp->data = im(target_addr);
		}

		tmp->address = address;
		tmp->address_assigned = true;

		address += estimate_instruction_size(tmp);
		tmp = tmp->next;
	} while (tmp != expr);
}

static uint16_t im_instruction(uint16_t payload)
{
	return payload | (((uint16_t) 1) << 15);
}

static void encode_instruction(struct instruction *instruction,
			       uint16_t *memory)
{
	uint32_t im = 0;
	uint16_t encoding = instruction->encoding;
	uint16_t *dest = memory + instruction->address / 2;

	if (instruction->data.info == DATA_INSTR_ADDR) {
		im = (*instruction->data.data.ptr)->address;
	} else if (instruction->data.info == DATA_ADDR_AFTER) {
		im = (*instruction->data.data.ptr)->address +
			estimate_instruction_size(*instruction->data.data.ptr);
	} else if (instruction->data.info != DATA_NONE) {
		im = instruction->data.data.im;
	}

	switch (instruction->data.info) {
	case DATA_INSTR_ADDR :
	case DATA_ADDR_AFTER :
	case DATA_KNOWN :
		*(dest++) = im_instruction(im >> 22);
	case DATA_KNOWN_21_BITS :
		*(dest++) = im_instruction(im >> 7);
	case DATA_KNOWN_6_BITS :
		encoding |= im & 0x7F;
	}

	*dest = encoding;
}

static void encode_expr(struct instruction *expr, uint16_t *memory)
{
	struct instruction *tmp = expr;

	do {
		encode_instruction(tmp, memory);
		tmp = tmp->next;
	} while (tmp != expr);
}

int assemble(uint32_t memory_size, uint16_t memory[memory_size],
	     struct module *module)
{
	uint32_t i;
	struct function *main_function = NULL;
	uint32_t current_address;
	uint64_t function_size;
	struct instruction *startup = NULL;
	unsigned short startup_size;
	int retval = -1;

	for (i = 0; i < module->exports_count; i++) {
		if (module->exports[i].desc == EXPORT_FUNCIDX &&
		    !strcmp(module->exports[i].name, "main")) {
			main_function = module->functions +
				module->exports[i].idx;
			break;
		}
	}

	if (!main_function) {
		PRERR("No 'main' function\n");
		goto fail;
	}

	/*
	 * We're first writing some value at STACK_FRAME_BACKUP_ADDR to be able
	 * to check, if the functions we call restore frame address properly
	 */
	if (i_const(im(0x23),                             &startup) ||
	    i_store(im(STACK_FRAME_BACKUP_ADDR),          &startup) ||
	    i_call (ptr(&main_function->translated_body), &startup) ||
	    i_halt (                                      &startup))
		goto fail;

	startup_size = estimate_expr_size(startup);

	i = module->functions_count;
	current_address = CODE_TOP_ADDR;

	while (i--) {
		function_size = estimate_expr_size(module->functions[i]
						   .translated_body);
		if (function_size > current_address - startup_size) {
			PRERR("Not enough space for code\n");
			goto fail;
		}

		current_address -= function_size;
		module->functions[i].start_addr = current_address;
	}

	i = module->functions_count;
	while (i--)
		assign_addresses_and_sizes(module->functions[i].translated_body,
					   module->functions[i].start_addr);

	assign_addresses_and_sizes(startup, 0x0);

	i = module->functions_count;
	while (i--)
		encode_expr(module->functions[i].translated_body, memory);

	encode_expr(startup, memory);

	retval = 0;

fail:
	if (retval)
		PRERR("Couldn't assemble code for stack machine\n");

	free_expr(startup);

	return retval;
}