#include "dead_code.h"

#include <map>
#include <string>
#include <vector>
#include <set>
#include <algorithm>

#include "../CFGraph/CFGraph.h"

using std::map;
using std::string;
using std::vector;
using std::set;

using std::remove_if;

#define PRINT_USELESS_STATEMENTS 1
#define DEBUG_IR 0

static void updateUseDefForInstruction(SAPFOR::BasicBlock* block, SAPFOR::Instruction* instr,
                                       set<SAPFOR::Argument*>& use, set<SAPFOR::Argument*>& def,
                                       vector<SAPFOR::Argument*>& formal_parameters,
                                       vector<SAPFOR::Argument*>& arg_stack, int& arg_stack_index, int& arg_stack_size,
                                       string& fName, const map<string, FuncInfo*>& funcByName,
                                       bool& useful, bool& last_stack_op_useful,
                                       set<SAPFOR::Argument*>& usedByThisBlock)
{
    set<SAPFOR::Argument*> res, args;
    bool last_stack_op;

    vector<LIVE_VARIABLES::LiveDeadVarsForCall> fcalls;

    getUseDefForInstruction(block, instr,
        args, res,
        formal_parameters, fcalls,
        arg_stack, arg_stack_index, arg_stack_size,
        last_stack_op,
        fName, funcByName,
        false
    );

    const auto instr_operation = instr->getOperation();

    if (!useful)
    {
        for (SAPFOR::Argument* r : res)
        {
            if (use.find(r) != use.end() || 
                r->getMemType() != SAPFOR::CFG_MEM_TYPE::LOCAL_ || 
                !(r->getType() == SAPFOR::CFG_ARG_TYPE::VAR || r->getType() == SAPFOR::CFG_ARG_TYPE::REG))
            {
                useful = true;
                break;
            }
        }
    }

    if (!useful)
    {
        static const set<SAPFOR::CFG_OP> always_useful =
        {
            SAPFOR::CFG_OP::POINTER_ASS,
            SAPFOR::CFG_OP::STORE,
            SAPFOR::CFG_OP::REC_REF_STORE,
            SAPFOR::CFG_OP::RANGE,
            SAPFOR::CFG_OP::ENTRY,
            SAPFOR::CFG_OP::EXIT,
            SAPFOR::CFG_OP::DVM_DIR,
            SAPFOR::CFG_OP::SPF_DIR,
            SAPFOR::CFG_OP::IO_PARAM
        };

        if (always_useful.find(instr_operation) != always_useful.end())
            useful = true;
        else if (instr_operation == SAPFOR::CFG_OP::F_CALL)
        {
            auto func_it = funcByName.find(instr->getArg1()->getValue());
            useful |= func_it == funcByName.end() || !(func_it->second->isPure);
        }
        else if (instr_operation == SAPFOR::CFG_OP::PARAM ||
                 instr_operation == SAPFOR::CFG_OP::REF)
            useful |= last_stack_op_useful;
    }

    if (useful)
    {
        if (instr_operation == SAPFOR::CFG_OP::F_CALL ||
            instr_operation == SAPFOR::CFG_OP::LOAD ||
            instr_operation == SAPFOR::CFG_OP::STORE)
            last_stack_op_useful = true;

        for (auto& e : res)
        {
            def.insert(e);
            use.erase(e);
        }

        for (auto& e : args)
        {
            use.insert(e);
            def.erase(e);
        }
        
        usedByThisBlock.insert(args.begin(), args.end());
    }

    if (last_stack_op)
        last_stack_op_useful = false;
}


//Build use and def sets of block. Result are stored in use and def
static void buildUseDef(SAPFOR::BasicBlock* block, set<SAPFOR::Argument*>& use, set<SAPFOR::Argument*>& def,
                        vector<SAPFOR::Argument*>& formal_parameters, vector<bool>& useful,
                        set<SAPFOR::Argument*>& usedByThisBlock, const map<string, FuncInfo*>& funcByName)
{
    set<SAPFOR::Argument*> use_with_regs = use, def_with_regs = def;

    vector<SAPFOR::Argument*> arg_stack;
    int arg_stack_index = 0, arg_stack_size = 0;
    string fName = "";
    bool last_fcall_useful = false;

    const auto& instructions = block->getInstructions();
    int instr_size = instructions.size();

    for (int i = instr_size - 1; i >= 0; i--)
    {
        bool u = useful[i];

        updateUseDefForInstruction(block, instructions[i]->getInstruction(),
            use_with_regs, def_with_regs,
            formal_parameters,
            arg_stack, arg_stack_index, arg_stack_size,
            fName, funcByName,
            u, last_fcall_useful,
            usedByThisBlock
        );

        useful[i] = u;
    }

    use.insert(use_with_regs.begin(), use_with_regs.end());
    def.insert(def_with_regs.begin(), def_with_regs.end());
}


class DeadCodeAnalysisNode : public DataFlowAnalysisNode<map<SAPFOR::Argument*, vector<SAPFOR::BasicBlock*>>>
{
private:
    vector<bool> useful;
    bool useful_block = false;
    set<DeadCodeAnalysisNode*> next_notempty_in, next_notempty_out;
    bool useful_jump = false;

    vector<SAPFOR::Argument*>& formal_parameters;
    const map<string, FuncInfo*>& funcByName;
public:
    bool updateJumpStatus()
    {
        bool res = false;
        const auto& prev = getPrevBlocks();
        if (prev.size() > 1 && !useful.back() &&
            getBlock()->getInstructions().back()->getInstruction()->getOperation() == SAPFOR::CFG_OP::JUMP_IF)
        {
            const auto& example = dynamic_cast<DeadCodeAnalysisNode*>(*prev.begin())->next_notempty_out;
            
            for (auto& p : prev)
            {
                DeadCodeAnalysisNode* next = dynamic_cast<DeadCodeAnalysisNode*>(p);

                if (next->next_notempty_out != example)
                {
                    useful.back() = true;
                    res = true;
                    break;
                }
            }
        }

        return res;
    }

    bool updateNextNotEmpty()
    {
        bool updated = false;

        if(!useful_block)
        {
            set<DeadCodeAnalysisNode*> current = {};

            for (auto& nextP : getPrevBlocks())
            {
                DeadCodeAnalysisNode* next = dynamic_cast<DeadCodeAnalysisNode*>(nextP);

                if(!next)
                    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

                current.insert(next->next_notempty_out.begin(), next->next_notempty_out.end());
            }

            if (current != next_notempty_in)
            {
                next_notempty_in = current;
                updated = true;
            }
        }

        return updated;
    }

    map<SAPFOR::Argument*, vector<SAPFOR::BasicBlock*>> getIn()
    {
        return getBlock()->getLiveOut();
    }

    map<SAPFOR::Argument*, vector<SAPFOR::BasicBlock*>> getOut()
    {
        return getBlock()->getLiveIn();
    }

    bool addIn(const map<SAPFOR::Argument*, vector<SAPFOR::BasicBlock*>>& data)
    {
        return getBlock()->addLiveOut(data);
    }

    bool addOut(const map<SAPFOR::Argument*, vector<SAPFOR::BasicBlock*>>& data)
    {
        return getBlock()->addLiveIn(data);
    }

    bool updateState()
    {
        bool updated = false;

        if (!useful_block)
            updated |= updateNextNotEmpty();

        if (!useful_block)
        {
            if (next_notempty_in != next_notempty_out)
            {
                updated = true;
                next_notempty_out = next_notempty_in;
            }
        }

        updated |= updateJumpStatus();

        if(updated)
            this->forwardData({ });

        return updated;
    }

    DATA_FLOW_UPD_STATUS forwardData(const map<SAPFOR::Argument*, vector<SAPFOR::BasicBlock*>>& data)
    {
        bool inserted_prop = false, inserted_gen = false;

        SAPFOR::BasicBlock* bb = getBlock();

        set<SAPFOR::Argument*> use, def;

        for (const auto& byArg : data)
            use.insert(byArg.first);

        set<SAPFOR::Argument*> usedByThisBlock;
        buildUseDef(bb, use, def, this->formal_parameters, useful, usedByThisBlock, funcByName);

        auto in = bb->getLiveIn(), out = bb->getLiveOut();

        for (SAPFOR::Argument* arg : use)
        {
            if(arg && (arg->getType() == SAPFOR::CFG_ARG_TYPE::VAR || arg->getType() == SAPFOR::CFG_ARG_TYPE::REG))
            {
                bool this_block = usedByThisBlock.find(arg) != usedByThisBlock.end();

                if (!this_block)
                {
                    auto data_it = data.find(arg);

                    if (data_it == data.end())
                        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

                    inserted_prop |= bb->addLiveIn({ *data_it });
                }
                else
                {
                    auto in_it = in.find(arg);
                    bool skip = false;
                    if (in_it != in.end())
                    {
                        if (in_it->second.size() == 1 && *(in_it->second.begin()) == bb) 
                            skip = true; // nothing to do, inserted = false
                        else
                            bb->removeLiveIn(arg);
                    }
                    if(!skip)
                        inserted_gen |= bb->addLiveIn({ { arg, { bb } } });
                }
            }
        }

        if(!useful_block)
        {
            for(bool status : useful)
            {
                if (status)
                {
                    useful_block = true;

                    inserted_gen = true;
                    next_notempty_out = { this };
                    break;
                }
            }
        }

        if(inserted_gen)
            return DATA_FLOW_UPD_STATUS::GENERATED;
        else if(inserted_prop)
            return DATA_FLOW_UPD_STATUS::PROPAGATED;

        return DATA_FLOW_UPD_STATUS::NO_CHANGE;
    }

    DeadCodeAnalysisNode(SAPFOR::BasicBlock* block,
                         vector<SAPFOR::Argument*>& formal_parameters,
                         const map<string, FuncInfo*>& funcByName)
    :
        formal_parameters(formal_parameters),
        funcByName(funcByName)
    {
        setBlock(block);
        useful.resize(block->getInstructions().size(), false);
        this->forwardData({ });
    }
    
    const vector<bool>& getResult() { return useful; }
};

class DeadCodeAnalysis : public BackwardDataFlowAnalysis<DeadCodeAnalysisNode>
{
protected:
    vector<SAPFOR::Argument*>& formal_parameters;
    const map<string, FuncInfo*>& funcByName;

    DeadCodeAnalysisNode* createNode(SAPFOR::BasicBlock* block) override
    {
        return new DeadCodeAnalysisNode(block, formal_parameters, funcByName);
    }
public:
    DeadCodeAnalysis(vector<SAPFOR::Argument*>& formal_parameters, 
                     const map<string, FuncInfo*>& funcByName)
    :
        formal_parameters(formal_parameters),
        funcByName(funcByName)
    { }
};

static bool hasCalls(SgExpression* ex)
{
    if (ex)
    {
        if (ex->variant() == FUNC_CALL)
            return true;
        return hasCalls(ex->lhs()) || hasCalls(ex->rhs());
    }
    return false;
}

//TODO: add check for side effects for function calls
static bool hasCalls(SgStatement* stat)
{
    if (stat->variant() == FOR_NODE)
    {
        auto forSt = isSgForStmt(stat);
        return hasCalls(forSt->start()) || hasCalls(forSt->end()) || hasCalls(forSt->step());
    }
    else
        return hasCalls(stat->expr(0)) || hasCalls(stat->expr(1)) || hasCalls(stat->expr(2));
}

static void moveComment(SgStatement* stat)
{
    const char* comm = stat->comments();
    if (comm)
    {
        SgStatement* next = stat->lastNodeOfStmt()->lexNext();
        if (next)
        {
            const char* commNext = next->comments();
            if (commNext)
            {
                string newComm(comm);
                newComm += commNext;

                next->delComments();
                next->addComment(newComm.c_str());
            }
            else
                next->addComment(comm);
        }
    }
}

int removeDeadCode(SgStatement* func,
                   const map<string, vector<FuncInfo*>>& allFuncs,
                   const map<string, CommonBlock*>& commonBlocks,
                   SgStatement* intervalDelStart, SgStatement* intervalDelEnd)
{
    int countOfTransform = 0;
    if (intervalDelStart && !intervalDelEnd || !intervalDelStart && intervalDelEnd)
        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

    SgProgHedrStmt* prog = isSgProgHedrStmt(func);

    if (intervalDelStart)
        if (intervalDelStart->lineNumber() < prog->lineNumber() || intervalDelStart->lineNumber() > prog->lastNodeOfStmt()->lineNumber())
            printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

    if (intervalDelEnd)
        if (intervalDelEnd->lineNumber() < prog->lineNumber() || intervalDelEnd->lineNumber() > prog->lastNodeOfStmt()->lineNumber())
            printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

    auto cfg = buildCFGforCurrentFunc(func, SAPFOR::CFG_Settings(true, false, false, false, false, false, false), commonBlocks, allFuncs);

    if(cfg.size() != 1)
        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

    auto& cfg_pair = *(cfg.begin());
    removedUnreachableBlocks(cfg_pair.second);

#if DEBUG_IR
    dumpCFG({ cfg_pair }, false);
#endif
    // detect useless code
    vector<SAPFOR::Argument*> func_parameters(cfg_pair.first->funcParams.countOfPars, NULL);

    map<string, FuncInfo*> funcByName;

    for (auto& byFile : allFuncs)
        for (auto& byFunc : byFile.second)
            funcByName[byFunc->funcName] = byFunc;

    DeadCodeAnalysis analysis_object(func_parameters, funcByName);

    analysis_object.fit(cfg_pair.second);
    analysis_object.analyze();

    // detect dead statements

    set<SgStatement*> useful;

    for (DeadCodeAnalysisNode* byNode : analysis_object.getNodes())
    {
        const auto& instructions = byNode->getBlock()->getInstructions();
        const auto& statuses = byNode->getResult();

        int size = instructions.size();

        if(size != statuses.size())
            printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

        for (int i = 0; i < size; i++) 
        {
            if(statuses[i])
            {
                SgStatement* stmt = instructions[i]->getInstruction()->getOperator();

                while(stmt && useful.insert(stmt).second)
                    stmt = stmt->controlParent();
            }
        }
    }

    // remove dead statements
    static const set<int> removable =
    {
        ASSIGN_STAT,
        PROC_STAT,
        WRITE_STAT,
        READ_STAT
    };

    vector<SgStatement*> remove;
    SgStatement* start = intervalDelStart ? intervalDelStart : func;
    SgStatement* end = intervalDelEnd ? intervalDelEnd : func->lastNodeOfStmt();

    for (auto st = start; st != end; st = st->lexNext())
    {
        const int var = st->variant();
        if (removable.find(var) != removable.end() && useful.find(st) == useful.end())
        {            
            remove.push_back(st);
            st = st->lastNodeOfStmt();
        }

        if (var == CONTAINS_STMT)
            break;
    }

    for (auto& rem : remove)
    {
        __spf_print(PRINT_USELESS_STATEMENTS, "[Useless statement on line %d and file %s]\n", rem->lineNumber(), rem->fileName());

        auto cp = rem->controlParent();
        if (cp->variant() == LOGIF_NODE)
        {
            if (hasCalls(cp))
            {
                ((SgLogIfStmt*)cp)->convertLogicIf();
                rem->deleteStmt();
            }
            else
            {
                moveLabelBefore(cp);
                moveComment(cp);
                cp->deleteStmt();
            }
        }
        else
        {
            moveLabelBefore(rem);
            moveComment(rem);
            rem->deleteStmt();
        }
    }
    countOfTransform += remove.size();

    //remove empty blocks
    do
    {
        remove.clear();
        for (auto st = start; st != end; st = st->lexNext())
        {
            const int var = st->variant();
            if ((var == FOR_NODE || var == WHILE_NODE || var == SWITCH_NODE) &&
                st->lexNext()->variant() == CONTROL_END)
            {
                if (!hasCalls(st))
                    remove.push_back(st);
            }
            else if (var == IF_NODE)
            {
                if (!hasCalls(st))
                {
                    bool hasCalls_ = false;
                    SgStatement* ifS = st;
                    while (ifS->lexNext()->variant() == ELSEIF_NODE)
                    {
                        hasCalls_ |= hasCalls(ifS->lexNext());
                        ifS = ifS->lexNext();
                    }

                    SgStatement* lastNode = ifS->lastNodeOfStmt();
                    ifS = ifS->lexNext();

                    while (ifS->variant() == CONTROL_END && ifS != lastNode)
                        ifS = ifS->lexNext();

                    if (ifS == lastNode && !hasCalls_)
                        remove.push_back(st);
                }
            }

            //TODO: other block statements

            if (var == CONTAINS_STMT)
                break;
        }

        bool mainRemoved = false;
        for (auto& rem : remove)
        {
            __spf_print(PRINT_USELESS_STATEMENTS, "[Useless block statement on line %d and file %s]\n", rem->lineNumber(), rem->fileName());
            moveLabelBefore(rem);
            moveComment(rem);
            rem->deleteStmt();
            if (rem == start)
                mainRemoved = true;
        }
        countOfTransform += remove.size();
        if (mainRemoved)
            break;
    } while (remove.size());

    deleteCFG(cfg);
    return countOfTransform;
}