/*
 * Andrea Di Biagio
 * Politecnico di Milano, 2007
 * 
 * axe_transform.c
 * Formal Languages & Compilers Machine, 2007/2008
 * 
 */

#include "axe_transform.h"
#include "symbol_table.h"
#include "cflow_constants.h"
#include "reg_alloc_constants.h"
#include "axe_errors.h"

extern int errorcode;
extern int cflow_errorcode;

typedef struct t_tempLabel
{
   t_axe_label *labelID;
   int regID;
} t_tempLabel;


static t_axe_instruction * _createUnary (t_program_infos *program
            , int reg, t_axe_label *label, int opcode);
static t_axe_instruction * createUnaryInstruction
               (t_program_infos *program, int reg, int opcode);
static t_list * updateNode(t_program_infos *program, t_reg_allocator *RA
            , t_cflow_Graph *graph, t_basic_block *block
            , t_cflow_Node *current_node, t_list *tempLabels);
static t_tempLabel * allocTempLabel(t_axe_label *labelID, int regID);
static void freeTempLabel(t_tempLabel *tempLabel);
static void finalizeListOfTempLabels(t_list *tempLabels);
static int compareTempLabels(void *valA, void *valB);
static t_list * performStore_Spill (t_program_infos *program, int regSource
            , t_cflow_Graph *graph, t_basic_block *block
            , t_cflow_Node *current_node, t_list *tempLabels);
static t_list * performLoad_Spill (t_program_infos *program, int regDest
            , int inst_reg, t_cflow_Graph *graph, t_basic_block *block
            , t_cflow_Node *current_node, t_list *tempLabels);
static void updateTheDataSegment
            (t_program_infos *program, t_list *tempLabels);
static void updateTheCodeSegment
            (t_program_infos *program, t_cflow_Graph *graph);
static t_list * _insertLoad(t_program_infos *program, t_cflow_Graph *graph
      , t_basic_block *block, t_cflow_Node *current_node
            , t_cflow_var *var, t_list *usedVars);
static t_list * _insertStore(t_program_infos *program, t_cflow_Graph *graph
      , t_basic_block *block, t_cflow_Node *current_node
            , t_cflow_var *var, t_list *usedVars);


t_list * _insertStore(t_program_infos *program, t_cflow_Graph *graph
      , t_basic_block *current_block, t_cflow_Node *current_node
            , t_cflow_var *var, t_list *usedVars)
{
   /* we have to insert a store instruction into the code */
   t_axe_instruction *storeInstr;
   t_cflow_Node *storeNode;

   /* create a load instruction */
   storeInstr = createStoreInstruction
            (program, var->ID);

   /* test if an error occurred */
   if (errorcode != AXE_OK) {
      free_Instruction(storeInstr);
      notifyError(errorcode);
   }
               
   if (storeInstr != NULL)
   {
      /* create a node for the store instruction */
      storeNode = allocNode (graph, storeInstr);
      if (cflow_errorcode != CFLOW_OK) {
         finalizeNode(storeNode);
         free_Instruction(storeInstr);
         return usedVars;
      }

      /* insert the node `loadNode' before `current_node' */
      insertNodeAfter(current_block, current_node, storeNode);
   
      /* update the list of usedVars */
      usedVars = removeElement(usedVars, storeNode->uses[0]);
   }

   return usedVars;
}

t_list * _insertLoad(t_program_infos *program, t_cflow_Graph *graph
      , t_basic_block *current_block, t_cflow_Node *current_node
            , t_cflow_var *var, t_list *usedVars)
{
   /* we have to insert a load instruction into the code */
   t_axe_instruction *current_instr;
   t_axe_instruction *loadInstr;
   t_cflow_Node *loadNode;

   /* retrieve the current instruction from the current node */
   current_instr = current_node->instr;

   /* create a load instruction */
   loadInstr = createLoadInstruction
         (program, var->ID);

   /* test if an error occurred */
   if (errorcode != AXE_OK) {
      free_Instruction(loadInstr);
      notifyError(errorcode);
   }

   if (loadInstr != NULL)
   {
      /* create a node for the load instruction */
      loadNode = allocNode (graph, loadInstr);

      /* test if an error occurred */
      if (cflow_errorcode != CFLOW_OK) {
         finalizeNode(loadNode);
         free_Instruction(loadInstr);
         return usedVars;
      }

      /* update the label informations */
      if (current_instr->labelID != NULL)
      {
         loadInstr->labelID = current_instr->labelID;
         current_instr->labelID = NULL;
      }
   
      /* insert the node `loadNode' before `current_node' */
      insertNodeBefore(current_block, current_node, loadNode);

      /* update the list of usedVars */
      usedVars = addElement(usedVars, loadNode->def, -1);
   }

   return usedVars;
}

void updateTheCodeSegment(t_program_infos *program, t_cflow_Graph *graph)
{
   t_list *current_bb_element;
   t_list *current_nd_element;
   t_basic_block *bblock;
   t_cflow_Node *node;
   
   /* preconditions */
   if (program == NULL)
      notifyError(AXE_PROGRAM_NOT_INITIALIZED);

   if (graph == NULL)
      notifyError(AXE_INVALID_CFLOW_GRAPH);

   current_bb_element = graph->blocks;
   while(current_bb_element != NULL)
   {
      bblock = (t_basic_block *) LDATA(current_bb_element);

      current_nd_element = bblock->nodes;
      while(current_nd_element != NULL)
      {
         node = (t_cflow_Node *) LDATA(current_nd_element);

         program->instructions =
               addElement(program->instructions, node->instr, -1);
         
         current_nd_element = LNEXT(current_nd_element);
      }

      current_bb_element = LNEXT(current_bb_element);
   }
}

void updateTheDataSegment(t_program_infos *program, t_list *tempLabels)
{
   t_list *current_element;
   t_tempLabel *current_TL;
   t_axe_data *new_data_info;

   /* preconditions */
   if (program == NULL) {
      finalizeListOfTempLabels(tempLabels);
      notifyError(AXE_PROGRAM_NOT_INITIALIZED);
   }

   /* initialize the value of `current_element' */
   current_element = tempLabels;
   while(current_element != NULL)
   {
      current_TL = (t_tempLabel *) LDATA(current_element);

      new_data_info = alloc_data (DIR_WORD, 0, current_TL->labelID);
         
      if (new_data_info == NULL){
         finalizeListOfTempLabels(tempLabels);
         notifyError(AXE_OUT_OF_MEMORY);
      }

      /* update the list of directives */
      program->data = addElement(program->data, new_data_info, -1);
      
      current_element = LNEXT(current_element);
   }
}

t_list * performLoad_Spill (t_program_infos *program, int regDest, int inst_reg
            , t_cflow_Graph *graph, t_basic_block *block
            , t_cflow_Node *current_node, t_list *tempLabels)
{
   t_axe_instruction *loadInstr;
   t_axe_register *axe_reg;
   t_axe_instruction *current_instr;
   t_cflow_Node *loadNode;

   /* initialize current_instr */
   current_instr = current_node->instr;

   /* intialize `axe_reg' */
   if (inst_reg == 0)
      axe_reg = current_instr->reg_1;
   else if (inst_reg == 1)
      axe_reg = current_instr->reg_2;
   else
      axe_reg = current_instr->reg_3;
   
   /* create a store instruction */
   loadInstr = createLoadInstruction (program, axe_reg->ID);
         
   /* test if an error occurred */
   if (errorcode != AXE_OK) {
      free_Instruction(loadInstr);
      return tempLabels;
   }

   /* the variable wasn't already defined */
   if (loadInstr == NULL)
   {
      t_list *elementFound;
      t_tempLabel pattern;
      t_tempLabel *tlabel;

      pattern.regID = axe_reg->ID;
      elementFound = CustomfindElement
            (tempLabels, &pattern, compareTempLabels);

      if (elementFound == NULL) {
         finalizeNode(loadNode);
         errorcode = AXE_TRANSFORM_ERROR;
         return tempLabels;
      }

      /* create a store instruction */
      loadInstr = _createUnary (program
            , tlabel->regID, tlabel->labelID, LOAD);

      /* test if an error occurred */
      if (errorcode != AXE_OK) {
         free_Instruction(loadInstr);
         finalizeNode(loadNode);
         return tempLabels;
      }
   }

   /* update the value of storeInstr and instr */
   (loadInstr->reg_1)->ID = regDest;
   axe_reg->ID = regDest;
            
   /* create a node for the load instruction */
   loadNode = allocNode (graph, loadInstr);
   if (cflow_errorcode != CFLOW_OK) {
      finalizeNode(loadNode);
      free_Instruction(loadInstr);
      errorcode = AXE_TRANSFORM_ERROR;
      return tempLabels;
   }

   if ((current_node->instr)->labelID != NULL)
   {
      /* modify the label informations */
      loadInstr->labelID = (current_node->instr)->labelID;
      (current_node->instr)->labelID = NULL;
   }
                  
   /* insert the node `loadNode' before `current_node' */
   insertNodeBefore(block, current_node, loadNode);

   return tempLabels;
}

t_list * performStore_Spill (t_program_infos *program, int regSource
            , t_cflow_Graph *graph, t_basic_block *block
            , t_cflow_Node *current_node, t_list *tempLabels)
{
   t_axe_instruction *storeInstr;
   t_cflow_Node *storeNode;
   t_axe_instruction *current_instr;

   /* initialize current_instr */
   current_instr = current_node->instr;

   /* create a store instruction */
   storeInstr = createStoreInstruction
         (program, (current_instr->reg_1)->ID);
         
   /* test if an error occurred */
   if (errorcode != AXE_OK) {
      free_Instruction(storeInstr);
      return tempLabels;
   }

   /* the variable was already defined */
   if (storeInstr == NULL)
   {
      t_list *elementFound;
      t_tempLabel pattern;
      t_tempLabel *tlabel;

      pattern.regID = (current_instr->reg_1)->ID;
      elementFound = CustomfindElement
            (tempLabels, &pattern, compareTempLabels);

      if (elementFound == NULL)
      {
         t_axe_label *new_label;

         new_label = reserveLabel(program);
         if (new_label == NULL) {
            errorcode = AXE_OUT_OF_MEMORY;
            return tempLabels;
         }

         /* create a new tempLabel */
         tlabel = allocTempLabel(new_label, regSource);
         if (errorcode != AXE_OK)
            return tempLabels;

         /* update the list of tempLabels */
         tempLabels = addElement(tempLabels, tlabel, -1);
      }

      /* create a store instruction */
      storeInstr = _createUnary (program
            , tlabel->regID, tlabel->labelID, STORE);

      /* test if an error occurred */
      if (errorcode != AXE_OK) {
         return tempLabels;
         free_Instruction(storeInstr);
      }
   }

   /* update the value of storeInstr and instr */
   (storeInstr->reg_1)->ID = regSource;
   (current_instr->reg_1)->ID = regSource;
            
   /* create a node for the load instruction */
   storeNode = allocNode (graph, storeInstr);
   if (cflow_errorcode != CFLOW_OK) {
      finalizeNode(storeNode);
      free_Instruction(storeInstr);
      errorcode = AXE_TRANSFORM_ERROR;
      return tempLabels;
   }

   /* insert the node `loadNode' before `current_node' */
   insertNodeAfter(block, current_node, storeNode);

   return tempLabels;
}

t_axe_instruction * _createUnary (t_program_infos *program
            , int reg, t_axe_label *label, int opcode)
{
   t_axe_instruction *result;
   
   /* preconditions */
   if (program == NULL) {
      errorcode = AXE_PROGRAM_NOT_INITIALIZED;
      return NULL;
   }
   
   if (label == NULL) {
      errorcode = AXE_INVALID_LABEL;
      return NULL;
   }

   /* create an instance of `t_axe_instruction' */
   result = alloc_instruction(opcode);
   if (result == NULL) {
      errorcode = AXE_OUT_OF_MEMORY;
      return NULL;
   }

   result->reg_1 = alloc_register(reg, 0);
   if (result->reg_1 == NULL) {
      errorcode = AXE_OUT_OF_MEMORY;
      free_Instruction(result);
      return NULL;
   }

   /* initialize an address info */
   result->address = alloc_address(LABEL_TYPE, 0, label);
   if (result->address == NULL) {
      errorcode = AXE_OUT_OF_MEMORY;
      free_Instruction(result);
      return NULL;
   }

   return result;
}

int compareTempLabels(void *valA, void *valB)
{
   t_tempLabel *tlA;
   t_tempLabel *tlB;
   
   if (valA == NULL)
      return 0;
   if (valB == NULL)
      return 0;

   tlA = (t_tempLabel *) valA;
   tlB = (t_tempLabel *) valB;

   return (tlA->regID == tlB->regID);
}

t_tempLabel * allocTempLabel(t_axe_label *labelID, int regID)
{
   t_tempLabel *result;

   /* preconditions */
   if (labelID == NULL) {
      errorcode = AXE_INVALID_LABEL;
      return NULL;
   }

   /* create a new temp-label */
   result = _AXE_ALLOC_FUNCTION(sizeof(t_tempLabel));
   if (result == NULL) {
      errorcode = AXE_OUT_OF_MEMORY;
      return NULL;
   }
   
   /* initialize the temp label */
   result->labelID = labelID;
   result->regID = regID;

   return result;
}

void freeTempLabel(t_tempLabel *tempLabel)
{
   if (tempLabel != NULL)
      _AXE_FREE_FUNCTION(tempLabel);
}

void finalizeListOfTempLabels(t_list *tempLabels)
{
   t_list *current_element;
   t_tempLabel *tempLabel;

   /* test the preconditions */
   if (tempLabels == NULL)
      return;

   /* assign to current_element the address of the first
    * element of the list `tempLabels' */
   current_element = tempLabels;
   
   while(current_element != NULL)
   {
      tempLabel = (t_tempLabel *) LDATA(current_element);
      if (tempLabel != NULL)
         freeTempLabel(tempLabel);
      
      current_element = LNEXT(current_element);
   }

   /* finalize the list of elements of type `t_tempLabel' */
   freeList(tempLabels);
}

t_list * updateNode(t_program_infos *program, t_reg_allocator *RA
            , t_cflow_Graph *graph, t_basic_block *block
            , t_cflow_Node *current_node, t_list *tempLabels)
{
   t_axe_instruction *instr;
   int bindings[3];

   /* initialize bindings */
   bindings[0] = RA_REGISTER_INVALID;
   bindings[1] = RA_REGISTER_INVALID;
   bindings[2] = RA_REGISTER_INVALID;

   /* initialize `instr' */
   instr = current_node->instr;
   
   if (instr->reg_1 != NULL)
      bindings[0] = RA->bindings[(instr->reg_1)->ID];
   if (instr->reg_2 != NULL)
      bindings[1] = RA->bindings[(instr->reg_2)->ID];
   if (instr->reg_3 != NULL)
      bindings[2] = RA->bindings[(instr->reg_3)->ID];

   /* update bindings */
   if (bindings[0] == RA_SPILL_REQUIRED)
      bindings[0] = RA->regNum + 1;
   if (bindings[1] == RA_SPILL_REQUIRED)
      bindings[1] = RA->regNum + 2;
   if (bindings[2] == RA_SPILL_REQUIRED)
      bindings[2] = RA->regNum + 3;

   instr = current_node->instr;
   if ((instr->reg_1 != NULL) && (bindings[0] != RA_REGISTER_INVALID))
   {
      if ( RA->bindings[(instr->reg_1)->ID] == RA_SPILL_REQUIRED)
      {
         if ((instr->reg_1)->indirect)
         {
            tempLabels = performLoad_Spill (program, bindings[0], 0
                        , graph, block, current_node, tempLabels);
         
            /* test if an error occurred */
            if (errorcode != AXE_OK)
               return tempLabels;
         }
         else
         {
            tempLabels = performStore_Spill (program, bindings[0]
                        , graph, block, current_node, tempLabels);

            /* test if an error occurred */
            if (errorcode != AXE_OK)
               return tempLabels;
         }
      } 
      else
         (instr->reg_1)->ID = RA->bindings[(instr->reg_1)->ID];
   }
   if ((instr->reg_2 != NULL) && (bindings[1] != RA_REGISTER_INVALID))
   {
      if (RA->bindings[(instr->reg_2)->ID] == RA_SPILL_REQUIRED)
      {
         tempLabels = performLoad_Spill (program, bindings[1], 1
                     , graph, block, current_node, tempLabels);
         
         /* test if an error occurred */
         if (errorcode != AXE_OK)
            return tempLabels;
      }
      else
         (instr->reg_2)->ID = RA->bindings[(instr->reg_2)->ID];
   }
   if ((instr->reg_3 != NULL) && (bindings[2] != RA_REGISTER_INVALID))
   {
      if (RA->bindings[(instr->reg_3)->ID] == RA_SPILL_REQUIRED)
      {
         tempLabels = performLoad_Spill (program, bindings[2], 2
                     , graph, block, current_node, tempLabels);
         
         /* test if an error occurred */
         if (errorcode != AXE_OK)
            return tempLabels;
      }
      else
         (instr->reg_3)->ID = RA->bindings[(instr->reg_3)->ID];
   }

   return tempLabels;
}

void updateProgramInfos(t_program_infos *program
         , t_cflow_Graph *graph, t_reg_allocator *RA)
{
   t_list *current_bb_element;
   t_basic_block *current_bblock;
   t_list *current_nd_element;
   t_cflow_Node *current_node;
   t_list *tempLabels;

   /* preconditions */
   assert(RA != NULL);
   assert(program != NULL);
   assert(graph != NULL);

   /* initialize tempLabels */
   tempLabels = NULL;
   current_bb_element = graph->blocks;
   while(current_bb_element != NULL)
   {
      current_bblock = (t_basic_block *) LDATA(current_bb_element);
      current_nd_element = current_bblock->nodes;
      while(current_nd_element != NULL)
      {
         current_node = (t_cflow_Node *) LDATA(current_nd_element);

         /* update the cflow informations for this node */
         updateNode(program, RA, graph, current_bblock, current_node, tempLabels);
         if (errorcode != AXE_OK){
            finalizeListOfTempLabels(tempLabels);
            notifyError(errorcode);
         }

         current_nd_element = LNEXT(current_nd_element);
      }

      current_bb_element = LNEXT(current_bb_element);
   }

   /* update the data segment informations */
   updateTheDataSegment(program, tempLabels);

   /* erase the old code segment */
   freeList(program->instructions);
   program->instructions = NULL;

   /* update the code segment informations */
   updateTheCodeSegment(program, graph);

   /* finalize the list of tempLabels */
   finalizeListOfTempLabels(tempLabels);
}

t_axe_instruction * createUnaryInstruction
               (t_program_infos *program, int reg, int opcode)
{
   char *varID;
   int sy_errorcode;
   t_axe_variable *axe_var;

   /* test the preconditions */
   if ((reg == REG_INVALID) || (reg == REG_0)) {
      errorcode = AXE_INVALID_REGISTER_INFO;
      return NULL;
   }

   if (program == NULL || program->sy_table == NULL) {
      errorcode = AXE_PROGRAM_NOT_INITIALIZED;
      return NULL;
   }

   /* initialize the value of errorcode */
   sy_errorcode = SY_TABLE_OK;
   varID = getIDfromLocation(program->sy_table, reg, &sy_errorcode);
   if (varID == NULL)
      return NULL;

   /* retrieve the variable associated with `varID' */
   axe_var = getVariable (program, varID);

   return _createUnary (program, reg, axe_var->labelID, opcode);
}

t_axe_instruction * createStoreInstruction
                  (t_program_infos *program, int reg)
{
   return createUnaryInstruction(program, reg, STORE);
}

t_axe_instruction * createLoadInstruction
                  (t_program_infos *program, int reg)
{
   return createUnaryInstruction(program, reg, LOAD);
}

t_cflow_Graph * insertLoadAndStoreInstr
         (t_program_infos *program, t_cflow_Graph *graph)
{
   t_list *current_bb_element;
   t_basic_block *current_block;
   t_list *current_nd_element;
   t_cflow_Node *current_node;
   t_list *usedVars;

   /* assertions */
   assert(program != NULL);
   assert(graph != NULL);
   
   /* initialize the list of variables used by this basic block */
   usedVars = NULL;
   current_bb_element = graph->blocks;
   while (current_bb_element != NULL)
   {
      current_block = (t_basic_block *) LDATA(current_bb_element);

      /* retrieve the list of nodes for the current basic block */
      current_nd_element = current_block->nodes;
      while(current_nd_element != NULL)
      {
         current_node = (t_cflow_Node *) LDATA(current_nd_element);

         /* test if we have to insert a store */
         if (  (current_node->def != NULL)
               && ((current_node->def)->ID != REG_0) )
         {
            if (findElement(usedVars, current_node->def) == NULL)
            {
               usedVars = addElement(usedVars, current_node->def, -1);
            }
         }

         /* test if we have to insert a load */
         if (  (current_node->uses[0] != NULL)
               && ((current_node->uses[0])->ID != REG_0) )
         {
            if (findElement(usedVars, current_node->uses[0]) == NULL)
            {
               usedVars = _insertLoad(program, graph, current_block
                     , current_node, current_node->uses[0], usedVars);
            }
         }
         if (  (current_node->uses[1] != NULL)
               && ((current_node->uses[1])->ID != REG_0) )
         {
            if (findElement(usedVars, current_node->uses[1]) == NULL)
            {
               usedVars = _insertLoad(program, graph, current_block
                     , current_node, current_node->uses[1], usedVars);
            }
         }
         if (  (current_node->uses[2] != NULL)
               && ((current_node->uses[2])->ID != REG_0) )
         {
            if (findElement(usedVars, current_node->uses[2]) == NULL)
            {
               usedVars = _insertLoad(program, graph, current_block
                     , current_node, current_node->uses[2], usedVars);
            }
         }

         /* retrieve the next element */
         current_nd_element = LNEXT(current_nd_element);
      }

      current_nd_element = getLastElement(current_block->nodes);
      while((current_nd_element != NULL) && (usedVars != NULL))
      {
         current_node = (t_cflow_Node *) LDATA(current_nd_element);

         /* test if we have to insert a store */
         if (  (current_node->uses[0] != NULL)
               && ((current_node->uses[0])->ID != REG_0) )
         {
            if (findElement(usedVars, current_node->uses[0]) != NULL)
            {
               usedVars = _insertStore(program, graph, current_block
                     , current_node, current_node->uses[0], usedVars);
            }
         }
         /* test if we have to insert a store */
         if (  (current_node->uses[1] != NULL)
               && ((current_node->uses[1])->ID != REG_0) )
         {
            if (findElement(usedVars, current_node->uses[1]) != NULL)
            {
               usedVars = _insertStore(program, graph, current_block
                     , current_node, current_node->uses[1], usedVars);
            }
         }
         /* test if we have to insert a store */
         if (  (current_node->uses[2] != NULL)
               && ((current_node->uses[2])->ID != REG_0) )
         {
            if (findElement(usedVars, current_node->uses[2]) != NULL)
            {
               usedVars = _insertStore(program, graph, current_block
                     , current_node, current_node->uses[2], usedVars);
            }
         }
         /* test if we have to insert a store */
         if (  (current_node->def != NULL)
               && ((current_node->def)->ID != REG_0) )
         {
            if (findElement(usedVars, current_node->def) != NULL)
            {
               usedVars = _insertStore(program, graph, current_block
                     , current_node, current_node->def, usedVars);
            }
         }
         
         /* retrieve the previous element */
         current_nd_element = LPREV(current_nd_element);
      }

      current_bb_element = LNEXT(current_bb_element);
   }

   /* free the list `usedVars' */
   freeList(usedVars);
   
   return graph;
}
