//===----------- FernandoLiveStack.cpp - Printing prototype ---------------===//
//
//              Register Allocation via Coloring of Chordal Graphs 
//
// This file was developed by Fernando Magno Quintao Pereira, fernando at ucla
// dot edu, and it holds the variables alive at certain point during a scan of
// the control flow graph.
//
// Date: May 28th, 2006
//
//===----------------------------------------------------------------------===//
//
// This pass collects the liveness information of the control flow graph.
// For each pair of basic blocks u and v, it records the set of variables
// alive between u and v. The requirements in terms of space is O(N*N*T),
// where N is the number of basic blocks, and T is the total number of
// temporaries in the control flow graph. Notice that this program only
// records information about virtual registers. Physical registers are
// not supposed to be alive between basic blocks.
//
//===----------------------------------------------------------------------===//

#include <vector>
#include <sstream>                               // ostringstream
#include <iostream>                              // cerr

#include "llvm/CodeGen/Passes.h"                 // EdgeLivenessFernandoPassID
#include "llvm/CodeGen/SSARegMap.h"
#include "llvm/Instructions.h"                   // MachineOperand
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/CodeGen/MachineFunctionPass.h"    // MachineBasicBlock
#include "llvm/Target/MRegisterInfo.h"
#include "llvm/Target/TargetInstrInfo.h"
#include <llvm/Target/TargetMachine.h>
#include "llvm/ADT/DepthFirstIterator.h"         // df_ext_iterator

#include "llvm/fernando/CriticalEdgeRemoval_Fer.h"
#include "llvm/fernando/DefUseSites_Fer.h"
#include "llvm/fernando/EdgeLiveness_Fernando.h" // class definition

using namespace llvm;


static RegisterAnalysis<EdgeLiveness_Fernando>
X("edge_liveness", "Edge liveness analysis");


//===--------------------------------------------------------------------------
// This method fills a vector with all the virtual registers alive between the
// origin and destiny blocks.
//===--------------------------------------------------------------------------
void EdgeLiveness_Fernando::to_vector(
    const MachineBasicBlock & org,
    const MachineBasicBlock & dst,
    std::vector<unsigned> & ans
) {
    int row = org.getNumber();
    int col = dst.getNumber();
    std::vector<bool> & live_variables = this->live_sets[row][col];
    ans.clear();

    for (unsigned i = 0, e = live_variables.size(); i != e; ++i) {
        if(live_variables[i]) {
            ans.push_back(i + MRegisterInfo::FirstVirtualRegister);
        }
    }
}


//===--------------------------------------------------------------------------
// returns a string in the format (t1, t2, ..., tn), where ti is one of the
// temporaries alive in the edge that connects the machine basic block
// "origin" to "destiny".
//===--------------------------------------------------------------------------
std::string EdgeLiveness_Fernando::to_string(
    const MachineBasicBlock & org,
    const MachineBasicBlock & dst
) {
    int row = org.getNumber();
    int col = dst.getNumber();
    std::vector<bool> live_variables = this->live_sets[row][col];
    bool isFirst = true;
    std::ostringstream aux;
    aux << "(";
    for (unsigned i = 0, e = live_variables.size(); i != e; ++i) {
        if(live_variables[i]) {
            if(isFirst) {
                isFirst = false;
                aux << "t" << i;
            } else {
                aux << ", " << "t" << i;
            }
        }
    }
    aux << ")";
    return aux.str();
}


//===--------------------------------------------------------------------------
// Removes temporary from the live set of the edge connecting the block origin
// to the block destiny.
//===--------------------------------------------------------------------------
void EdgeLiveness_Fernando::kill (
                                   const MachineBasicBlock & org,
                                   const MachineBasicBlock & dst,
                                   const MachineOperand & temporary
                                 )
{
    int row = org.getNumber();
    int col = dst.getNumber();

    if (temporary.getReg()) {
        unsigned reg_num = temporary.getReg();
        if(MRegisterInfo::isVirtualRegister(reg_num)) {
            unsigned base_reg = reg_num - MRegisterInfo::FirstVirtualRegister;
            std::vector<bool> & edge = this->live_sets[row][col];
            assert(edge.size() > base_reg &&
                   "EdgeLive_Fer: Index out of range in set alive.");
           edge[base_reg] = false;
        }
    }
}

//===--------------------------------------------------------------------------
// Removes the v_reg from the live set of all the blocks that can be
// reached from mbb.
//===--------------------------------------------------------------------------
void EdgeLiveness_Fernando::kill_forward
                   (const MachineOperand & mo, const MachineBasicBlock & mbb) {
    for(MachineBasicBlock::const_succ_iterator succ_aux = mbb.succ_begin();
        succ_aux != mbb.succ_end(); succ_aux++) {
        const MachineBasicBlock* succ_mbb = *succ_aux;
        if( is_alive(mbb, *succ_mbb, mo) ) {
            kill(mbb, *succ_mbb, mo);
            kill_forward(mo, *succ_mbb);
        }
    }
}


//===--------------------------------------------------------------------------
// Removes the v_reg from the live set of all the blocks that have a
// path leading to mbb.
//===--------------------------------------------------------------------------
void EdgeLiveness_Fernando::kill_backward
                   (const MachineOperand & mo, const MachineBasicBlock & mbb) {
    for(MachineBasicBlock::const_pred_iterator pred_aux = mbb.pred_begin();
                                      pred_aux != mbb.pred_end(); pred_aux++) {
        MachineBasicBlock* pred_mbb = *pred_aux;
        if( is_alive(*pred_mbb, mbb, mo) ) {
            kill(*pred_mbb, mbb, mo);
            kill_backward(mo, *pred_mbb);
        }
    }
}


//===--------------------------------------------------------------------------
// This method removes the live range of v_sp. This must be done after spill
// is performed. The placement of load and store instructions must be
// performed in such a way that is always safe to remove the live range of the
// spilled register.
//===--------------------------------------------------------------------------
void EdgeLiveness_Fernando::kill_live_range
                   (const MachineOperand & mo, const MachineBasicBlock & mbb) {
    kill_forward(mo, mbb);
    kill_backward(mo, mbb);
}


//===--------------------------------------------------------------------------
// This method returns a vector containing the registers alive in the edge
// between the given basic blocks. Notice that each element in the returned
// vector is not the LLVM representation of the virtual register. To get the
// virtual register number, add to the MRegisterInfo::FirstVirtualRegister to
// the index in the returned vector.
//===--------------------------------------------------------------------------
const std::vector<bool> &
EdgeLiveness_Fernando::get_live_variables
         (const MachineBasicBlock & org, const MachineBasicBlock & dst) const {
    int row = org.getNumber();
    int col = dst.getNumber();
    return this->live_sets[row][col];
}


//===--------------------------------------------------------------------------
// Every time a use is found, the graph is traversed backwards, and all the
// edges in which the used temporary is alive are marked.
//===--------------------------------------------------------------------------
void EdgeLiveness_Fernando::handle_use (
     const MachineOperand & m_op,
     const MachineBasicBlock & mbb)
{
    int block_number = mbb.getNumber();
    if(this->def_use_sites->get_def_block(m_op) != block_number) {
        // it means that the register is not defined in this very block; thus,
        // it is alive in all the edges that reach this block, if it is not
        // used only in a PHI function.
        for (MachineBasicBlock::const_pred_iterator pred = mbb.pred_begin(),
                                       E = mbb.pred_end(); pred != E; ++pred) {
            const MachineBasicBlock *pred_block = *pred;
            if( ! is_alive(*pred_block, mbb, m_op) ) {
                set_alive(*pred_block, mbb, m_op);
                handle_use(m_op, *pred_block);
            }
        }
    }
}


//===--------------------------------------------------------------------------
// This function scans every instruction of the basic block, and compute
// liveness information in a backward fashion. Every time a use is found, the
// graph is scanned backwards, and all the edges in which the used temporary is
// alive are marked. However, this does not happen if the use was found in a
// PHI instruction. In this case, only the edge from where comes the parameter
// of the PHI instruction is inspected. As a remark, PHI instruction, in LLVM,
// contain a list of operands like (v1, b1, v2, b2, ... v_n, bn), where v_i is
// a virtual register, and b_i is the block number from where the virtual
// register comes.
//===--------------------------------------------------------------------------
void EdgeLiveness_Fernando::compute_liveness(const MachineBasicBlock & block) {
    MachineBasicBlock::const_iterator mbbi = block.begin();
    while (mbbi != block.end()) {
        const MachineInstr *mi = mbbi++;
        for (unsigned i = 0; i != mi->getNumOperands(); ++i) {
            const MachineOperand & machine_operand = mi->getOperand(i);
            if(machine_operand.isRegister()) {
                if(machine_operand.getReg()) {
                    unsigned reg = machine_operand.getReg();
                    if(MRegisterInfo::isVirtualRegister(reg)) {
                        if (machine_operand.isUse()) {
                            if(mi->getOpcode() == TargetInstrInfo::PHI) {
                                const MachineBasicBlock * sbck =
                                    mi->getOperand(i+1).getMachineBasicBlock();
                                set_alive(*sbck, block, machine_operand);
                                handle_use(machine_operand, *sbck);
                            } else {
                                handle_use(machine_operand, block);
                            }
                        }
                    }
                }
            }
        }
    }
}


//===--------------------------------------------------------------------------
// adds the temporary to the live set of the edge that goes from machine basic
// block "origin" to "destiny". Notice that only virtual registers are added.
// Physical registers are not supposed to be alive across basic blocks.
//===--------------------------------------------------------------------------
void EdgeLiveness_Fernando::set_alive(
                                       const MachineBasicBlock & org,
                                       const MachineBasicBlock & dst,
                                       const MachineOperand & temporary
                                     )
{
    int row = org.getNumber();
    int col = dst.getNumber();

    if (temporary.getReg()) {
        unsigned reg_num = temporary.getReg();
        if(MRegisterInfo::isVirtualRegister(reg_num)) {
            unsigned base_reg = reg_num - MRegisterInfo::FirstVirtualRegister;
            std::vector<bool> & edge = this->live_sets[row][col];
            assert(this->num_vreg > base_reg &&
                             "EdgeLive_Fer: Index out of range in set alive.");
            if(edge.size() == 0) {
                edge.assign(this->num_vreg, false);
            }
            edge[base_reg] = true;
        }
    }
}


//===--------------------------------------------------------------------------
// This method informs if the given temporary is alive in the edge between the
// blocks "origin" and "destiny".
//===--------------------------------------------------------------------------
bool EdgeLiveness_Fernando::is_alive(
                                       const MachineBasicBlock & org,
                                       const MachineBasicBlock & dst,
                                       const MachineOperand & temporary
                                     )
{
    int row = org.getNumber();
    int col = dst.getNumber();
    unsigned reg_num = temporary.getReg();
    unsigned base_reg = reg_num - MRegisterInfo::FirstVirtualRegister;
    std::vector<bool> & edge = this->live_sets[row][col];
    if(edge.size() == 0) {
        return false;
    }
    return edge[base_reg];
}


//===--------------------------------------------------------------------------
// This method allocates memory for the data structures used in this analysis.
// Ihese data structures comprise the matrix (O(T*N*N)) that holds liveness
// information, and the vector that maps registers to the blocks where they
// where defined (O(T)).
//===--------------------------------------------------------------------------
void EdgeLiveness_Fernando::initialize_data_structures(MachineFunction & mf) {
    unsigned n_of_blocks = mf.size();
    unsigned i, j;
    SSARegMap * srm = mf.getSSARegMap();
    unsigned last_v_reg = srm->getLastVirtReg() + 1;
    this->num_vreg = last_v_reg - MRegisterInfo::FirstVirtualRegister;
    // initialize the matrix that will hold data flow information.
    this->live_sets.resize(n_of_blocks);
    for(unsigned u = 0; u < n_of_blocks; u++) {
        this->live_sets[u].resize(n_of_blocks);
    }
    // gets the list of def/use sites
    this->def_use_sites = & getAnalysis<DefUseSites_Fer>();
}


//===--------------------------------------------------------------------------
// Free the memory, so the operating system can use it in other applications.
//===--------------------------------------------------------------------------
void EdgeLiveness_Fernando::releaseMemory() {
    unsigned i, j;
    for(i = 0; i < this->live_sets.size(); i++) {
        for(j = 0; j < this->live_sets[i].size(); j++) {
            this->live_sets[i][j].clear();
        }
    }
}


//===--------------------------------------------------------------------------
// The method prints all the results obtained during the liveness analysis.
// It only print liveness analysis results. Although this class also contains
// a data structure that marks the block where each temporary is defined,
// the dump method does not print this information.
//===--------------------------------------------------------------------------
void EdgeLiveness_Fernando::dump() {
    for (MachineFunction::iterator mbb = this->machine_function->begin(),
         mbbe = this->machine_function->end();
         mbb != mbbe; ++mbb) {
        for (MachineBasicBlock::const_pred_iterator pred = mbb->pred_begin(),
             E = mbb->pred_end(); pred != E; ++pred) {
            const MachineBasicBlock *pred_block = *pred;
            int row = pred_block->getNumber();
            int col = mbb->getNumber();
            std::cerr << "(" << row << ", " << col << ") "
            << to_string(*pred_block, *mbb) << "\n";
        }
    }
}


//===--------------------------------------------------------------------------
// Collects liveness information about each edge of the control flow graph.
//===--------------------------------------------------------------------------
bool EdgeLiveness_Fernando::runOnMachineFunction(MachineFunction & mf) {
    this->machine_function = & mf;
    this->reg_info = mf.getTarget().getRegisterInfo();

    initialize_data_structures(mf);

    // traverse the control flow graph in depth first order pinpointing the
    // registers alive in each edge:
    for (MachineFunction::iterator mbb = mf.begin(), mbbe = mf.end();
                                                          mbb != mbbe; ++mbb) {
        compute_liveness(*mbb);
    }

    return false;
}