#include "array_index_rewrite.h"
#include "assert.h"
#include "dbug.h"
#include "memory.h"
#include "symbol_table.h"
#include "traverse.h"
#include "tree_basic.h"
#include "mymacros.h"
#include "free.h"
#include <string.h>

/* Generates a new, single-dimensional index expression. */
node* generate_new_index_expr(int relframeno, node* index_exprs, struct st_entry_arr* array_ent)
{
    int dims = 0;
    node* new_index_expr = TBmakeInt(0);
    for (node* exprs = index_exprs; exprs != NULL; exprs = EXPRS_NEXT(exprs), dims++)
    {
        node* multiplied_idx = EXPRS_EXPR(exprs);
        for (int i = dims + 1; i < array_ent->n_dims; ++i)
        {
            struct st_entry* dim_var = array_ent->dim_vars[i];
            node* var = TBmakeVar(strdup(dim_var->name), NULL);
            VAR_STLINK(var) = dim_var;
            // dimensions are stored in the same frame that the original array variable is.
            VAR_RELFRAMENO(var) = relframeno;

            multiplied_idx = TBmakeBinop(BO_mul, multiplied_idx, var);
            BINOP_TYPE(multiplied_idx) = BINOP_EXPRTYPE(multiplied_idx) = TY_int;
        }
        new_index_expr = TBmakeBinop(BO_add, new_index_expr, multiplied_idx);
        BINOP_TYPE(new_index_expr) = BINOP_EXPRTYPE(new_index_expr) = TY_int;
        EXPRS_EXPR(exprs) = NULL;
    }

    return new_index_expr;
}

node* AIRvar(node* arg_node, info* arg_info)
{
    if (!VAR_IS_ARRAYIDX(arg_node))
    {
        return arg_node;
    }

    struct st_entry_arr* array_ent = (struct st_entry_arr*)VAR_STLINK(arg_node);

    node* index_exprs = VAR_INDICES(arg_node);

    /* We don't need to rewrite single-dimensional array indices. */
    if (EXPRS_NEXT(index_exprs) == NULL)
    {
        return arg_node;
    }

    node* new_index_expr = generate_new_index_expr(VAR_RELFRAMENO(arg_node), index_exprs, array_ent);
    FREEdoFreeTree(index_exprs);

    node* exprs = TBmakeExprs(new_index_expr, NULL);

    VAR_INDICES(arg_node) = exprs;
    return arg_node;
}

node* AIRvarlet(node* arg_node, info* arg_info)
{
    if (!VARLET_IS_ARRAYIDX(arg_node))
    {
        return arg_node;
    }

    struct st_entry_arr* array_ent = (struct st_entry_arr*)VARLET_STLINK(arg_node);

    node* index_exprs = VARLET_INDICES(arg_node);

    /* We don't need to rewrite single-dimensional array indices. */
    if (EXPRS_NEXT(index_exprs) == NULL)
    {
        return arg_node;
    }

    node* new_index_expr = generate_new_index_expr(VARLET_RELFRAMENO(arg_node), index_exprs, array_ent);
    FREEdoFreeTree(index_exprs);

    node* exprs = TBmakeExprs(new_index_expr, NULL);

    VARLET_INDICES(arg_node) = exprs;
    return arg_node;
}


node* AIRdoRewrite(node* syntaxtree)
{
    DBUG_ENTER("AIRdoRewrite");

    TRAVpush(TR_air);
    syntaxtree = TRAVdo(syntaxtree, NULL);
    TRAVpop();

    DBUG_RETURN(syntaxtree);
}