#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <stdarg.h>
#include "simple-reg.h"
#include "stack-man.h"

#define TRUE 1
#define FALSE 0

static uint8_t
find_reg (const reg_pool_t *pool, op_t reg)
{
    int i;

    for (i = 0; i < pool->n_registers; ++i)
    {
	if (pool->registers[i] == reg)
	    return i;
    }

    return -1;
}

static void
reg_context_init_internal (reg_context_t *ctx,
                           reg_context_t *parent,
                           reg_alloc_t   *allocator,
			   int		  n_preserved,
			   va_list	  list)
{
    const reg_pool_t *pool = allocator->registers;
    uint8_t preserved[MAX_REGISTERS] = { 0 };
    int i;

    for (i = 0; i < n_preserved; ++i)
    {
	reg_t reg = va_arg (list, reg_t);

	preserved[find_reg (pool, reg)] = TRUE;
    }

    ctx->parent = parent;
    ctx->allocator = allocator;

    for (i = 0; i < pool->n_registers; ++i)
    {
	if (!parent)
	{
	    if (preserved[i])
		ctx->info[i].state = IN_USE;
	    else
		ctx->info[i].state = UNUSED;

	    continue;
	}

	switch (parent->info[i].state)
	{
	case SPILLABLE:
	    if (preserved[i])
	    {
		/* Preserving a SPILLABLE register is nonsensical
		 * because SPILLABLE means the register was _never_
		 * in use by the parent, and therefore can't have
		 * a meaningful value.
		 */
		fprintf (stderr,
			 "Asking to preserve a register that "
			 "was previously spilled");
		abort();
	    }
	    ctx->info[i].state = SPILLABLE;
	    break;
            
	case SPILLED:
	case IN_USE:
	    if (preserved[i])
		ctx->info[i].state = IN_USE;
	    else
		ctx->info[i].state = SPILLABLE;
	    break;
            
	case CLOBBERED:
	case UNUSED:
	    if (preserved[i])
	    {
		/* Preserving an UNUSED or CLOBBERED
		 * register is not necessarily nonsensical.
		 * The user may know that the register in
		 * question was recently freed and not
		 * overwritten since.
		 */
		ctx->info[i].state = IN_USE;
	    }
	    else
	    {
		ctx->info[i].state = UNUSED;
	    }
	    break;
	}
    }
}

void
reg_context_init (reg_context_t *ctx, reg_context_t *parent,
                  int n_preserved, ...)
{
    va_list list;

    assert (parent);

    va_start (list, n_preserved);
    
    reg_context_init_internal (ctx, parent, parent->allocator,
			       n_preserved, list);

    va_end (list);
}

void
reg_alloc_init (reg_alloc_t *     allocator,
                const reg_pool_t *registers,
                stack_man_t *     stack_man,
                reg_context_t *   initial_context,
		int n_preserved,
		...)
{
    va_list list;

    allocator->registers = registers;
    allocator->stack_man = stack_man;

    va_start (list, n_preserved);

    reg_context_init_internal (initial_context, NULL, allocator, 
			       n_preserved, list);

    va_end (list);
}

op_t
reg_context_alloc (reg_context_t *ctx, fragment_t *frag)
{
    const reg_pool_t *pool = ctx->allocator->registers;
    stack_man_t *stack_man = ctx->allocator->stack_man;
    int i;

    /* First try to allocate an unused register */
    for (i = 0; i < pool->n_registers; ++i)
    {
        reg_info_t *info = &ctx->info[i];
	op_t reg = pool->registers[i];

        switch (info->state)
        {
        case UNUSED:
            info->state = IN_USE;
            return reg;
            break;

        case CLOBBERED:
            info->state = SPILLED;
	    return reg;
            break;

        case SPILLED:
        case IN_USE:
        case SPILLABLE:
            break;
        }
    }

    /* If that failed, find a spillable one */
    for (i = 0; i < pool->n_registers; ++i)
    {
	reg_info_t *info = &ctx->info[i];
	op_t reg = pool->registers[i];

	if (info->state == SPILLABLE)
	{
            int offset;

            if (!stack_manager_alloc (stack_man, pool->register_size, &offset))
                return (op_t)-1;

            info->state = SPILLED;
            info->spill_offset = offset;

	    BEGIN_ASM (frag)
		I_mov, BASE (rsp, offset), reg,
	    END_ASM ();

            return reg;
        }
    }

    /* If that failed, we can't allocate a register */
    return (op_t)0;
}

/* When a register is freed, it doesn't become invalid
 * until next time a register is allocated. A register
 * known to be free-but-valid can be resurrected by
 * listing it in a call to reg_context_init().
 */
void
reg_context_free (reg_context_t *ctx, op_t op)
{
    const reg_pool_t *pool = ctx->allocator->registers;
    uint8_t idx = find_reg (pool, op);
    reg_info_t *info = &ctx->info[idx];

    switch (info->state)
    {
    case SPILLED:
        info->state = CLOBBERED;
        break;

    case IN_USE:
        info->state = UNUSED;
        break;

    case SPILLABLE:
    case CLOBBERED:
    case UNUSED:
        abort();
        break;
    }
}

/* Frees all registers allocated, except those listed, which
 * are reallocated in the parent allocator. Note that the
 * list is of pointers to op_t's. That's because the reallocation
 * may case the variables to move to other registers
 *
 * If it is not possible to reallocate all preserved registers,
 * FALSE is returned.
 */
bool_t
reg_context_fini (reg_context_t *ctx, fragment_t *frag,
                  int n_preserved,
                  /* op_t *var, */...)
{
    const reg_pool_t *pool = ctx->allocator->registers;
    stack_man_t *stack_man = ctx->allocator->stack_man;
    uint8_t preserved[MAX_REGISTERS] = { 0 };
    op_t *locations[MAX_REGISTERS] = { NULL };
    reg_context_t *parent = ctx->parent;
    va_list list;
    int i;
	
    /* If there is no parent, then noone cares whether any registers
     * are preserved.
     */
    if (!ctx->parent)
        return TRUE;

    va_start (list, n_preserved);
    for (i = 0; i < n_preserved; ++i)
    {
	op_t *reg = va_arg (list, op_t *);
        uint8_t idx = find_reg (pool, *reg);

        preserved[idx] = TRUE;
        locations[idx] = reg;
    }
    va_end (list);

    for (i = 0; i < pool->n_registers; ++i)
    {
        reg_info_t *info = &ctx->info[i];

        switch (info->state)
        {
        case UNUSED:
        case SPILLABLE:
        case CLOBBERED:
            if (preserved[i])
            {
                fprintf (
                    stderr, "Attempted to preserve an unallocated register\n");
                abort();
            }
            if (info->state == CLOBBERED)
                goto restore;
            break;
            
        case IN_USE:
            switch (ctx->parent->info[i].state)
            {
            case SPILLED:
            case IN_USE:
            case SPILLABLE:
                abort();
                break;
            case CLOBBERED:
		if (preserved[i])
		    parent->info[i].state = SPILLED;
                break;
            case UNUSED:
		if (preserved[i])
		    parent->info[i].state = IN_USE;
                break;
            }
            break;

        case SPILLED:
            if (preserved[i])
            {
                /* This is the complicated case. We have to preserve the
                 * register, but the parent needs it too. That means
                 * it must be allocated somewhere else.
                 *
                 * We will deal with this in a separate pass.
                 */
            }
            else
            {
            restore:
                /* Restore parent's value */

		BEGIN_ASM (frag)
		    I_mov, pool->registers[i], BASE (rsp, info->spill_offset),
		END_ASM ();

                stack_manager_free (stack_man, info->spill_offset);
            }
            break;
        }
    }

    /* Now deal with registers that are both spilled and preserved */
    for (i = 0; i < pool->n_registers; ++i)
    {
        reg_info_t *info = &ctx->info[i];
	op_t reg = pool->registers[i];

        if (info->state == SPILLED && preserved[i])
        {
            op_t new_location;

            new_location = reg_context_alloc (parent, frag);
            if (new_location == (op_t)0)
                return FALSE;

	    BEGIN_ASM (frag)
                I_mov, new_location, reg,
                I_mov, reg, BASE (rsp, info->spill_offset),
	    END_ASM ();

            *(locations[i]) = new_location;

            stack_manager_free (stack_man, info->spill_offset);
        }
    }
    
    return TRUE;
}