SAPFOR/src/LoopAnalyzer/implicit_loops_analyzer.cpp

#include <map>
#include <set>
#include <vector>
#include <string>
#include <iostream>
#include <stack>

#include "../CFGraph/IR.h"
#include "GraphCall/graph_calls.h"
#include "implicit_loops_analyzer.h"

using std::map;
using std::set;
using std::vector;
using std::pair;
using std::string;
using std::cout;
using std::endl;
using std::make_pair;
using std::to_string;

enum VisitState { UNVISITED = 0, VISITING = 1, VISITED = 2 };

static void dfs(SAPFOR::BasicBlock* block, map<int, int>& visit, vector<pair<SAPFOR::BasicBlock*, SAPFOR::BasicBlock*>>& startAndEnd, SAPFOR::BasicBlock* prev) {
    if (!block) 
        return;

    if (visit[block->getNumber()] == VISITED) 
    {
        cout << "error";
        return;
    }

    if (visit[block->getNumber()] == VISITING) 
    {
        visit[block->getNumber()] = VISITED;
        startAndEnd.push_back(make_pair(prev, block));
        return;
    }

    visit[block->getNumber()] = VISITING;
    for (auto i : block->getNext())
        dfs(i, visit, startAndEnd, block);
}

static void printBlock(SAPFOR::BasicBlock* block) {
    cout << "block - " << block->getNumber() << endl;
    cout << "next -";
    for (auto i : block->getNext())
        cout << " " << i->getNumber();

    cout << endl << "prev -";
    for (auto i : block->getPrev())
        cout << " " << i->getNumber();
    cout << endl;

    for (auto i : block->getInstructions())
    {
        string resValue = "";
        string arg1Value = "";
        string arg2Value = "";
        if (i->getInstruction()->getResult() != NULL && i->getInstruction()->getResult()->getType() == SAPFOR::CFG_ARG_TYPE::VAR) {
            resValue = i->getInstruction()->getResult()->getValue();
            i->getInstruction()->getResult()->setValue(i->getInstruction()->getResult()->getValue() + to_string(i->getInstruction()->getResult()->getNumber()));
        }

        if (i->getInstruction()->getArg1() != NULL && i->getInstruction()->getArg1()->getType() == SAPFOR::CFG_ARG_TYPE::VAR) {
            arg1Value = i->getInstruction()->getArg1()->getValue();
            i->getInstruction()->getArg1()->setValue(i->getInstruction()->getArg1()->getValue() + to_string(i->getInstruction()->getArg1()->getNumber()));
        }

        if (i->getInstruction()->getArg2() != NULL && i->getInstruction()->getArg2()->getType() == SAPFOR::CFG_ARG_TYPE::VAR) {
            arg2Value = i->getInstruction()->getArg2()->getValue();
            i->getInstruction()->getArg2()->setValue(i->getInstruction()->getArg2()->getValue() + to_string(i->getInstruction()->getArg2()->getNumber()));
        }

        cout << i->getNumber() << " " << i->getInstruction()->dump() << endl;

        if (i->getInstruction()->getResult() != NULL && i->getInstruction()->getResult()->getType() == SAPFOR::CFG_ARG_TYPE::VAR) 
            i->getInstruction()->getResult()->setValue(resValue);        

        if (i->getInstruction()->getArg1() != NULL && i->getInstruction()->getArg1()->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
            i->getInstruction()->getArg1()->setValue(arg1Value);

        if (i->getInstruction()->getArg2() != NULL && i->getInstruction()->getArg2()->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
            i->getInstruction()->getArg2()->setValue(arg2Value);
    }

    cout << endl;
}

static void getLoopBody(SAPFOR::BasicBlock* loopHeader, const set<SAPFOR::BasicBlock*>& loopExits, std::vector<SAPFOR::BasicBlock*>& loopBody) 
{
    set<SAPFOR::BasicBlock*> visited;
    std::stack<SAPFOR::BasicBlock*> stack;

    stack.push(loopHeader);

    while (!stack.empty()) 
    {
        auto block = stack.top();
        stack.pop();

        if (visited.count(block)) 
            continue;
        visited.insert(block);

        for (auto succ : block->getNext()) 
        {
            if (loopExits.count(succ)) 
                continue;

            if (!visited.count(succ))
                stack.push(succ);            
        }
    }

    set<SAPFOR::BasicBlock*> backReachable;
    std::stack<SAPFOR::BasicBlock*> reverseStack;
    reverseStack.push(loopHeader);

    while (!reverseStack.empty()) 
    {
        auto block = reverseStack.top();
        reverseStack.pop();

        if (backReachable.count(block)) 
            continue;
        backReachable.insert(block);

        for (auto pred : block->getPrev())
            if (visited.count(pred) && !backReachable.count(pred)) 
                reverseStack.push(pred);
    }

    for (auto block : visited) 
        if (backReachable.count(block))
            loopBody.push_back(block);
}


static set<SAPFOR::Argument*> findRegisterSourceVariables(const std::vector<SAPFOR::BasicBlock*>& blocks, SAPFOR::Argument* var)
{
    set<SAPFOR::Argument*> result;
    set<SAPFOR::Argument*> visited;
    std::stack<SAPFOR::Argument*> workStack;
    workStack.push(var);

    auto isBinaryOp = [](SAPFOR::CFG_OP op) {
        return op == SAPFOR::CFG_OP::ADD || op == SAPFOR::CFG_OP::SUBT ||
            op == SAPFOR::CFG_OP::MULT || op == SAPFOR::CFG_OP::DIV ||
            op == SAPFOR::CFG_OP::POW ||
            op == SAPFOR::CFG_OP::GE || op == SAPFOR::CFG_OP::LE ||
            op == SAPFOR::CFG_OP::GT || op == SAPFOR::CFG_OP::LT ||
            op == SAPFOR::CFG_OP::EQ || op == SAPFOR::CFG_OP::NEQV ||
            op == SAPFOR::CFG_OP::EQV || op == SAPFOR::CFG_OP::EMPTY ||
            op == SAPFOR::CFG_OP::OR || op == SAPFOR::CFG_OP::AND;
        };

    auto isUnaryOp = [](SAPFOR::CFG_OP op) {
        return op == SAPFOR::CFG_OP::UN_ADD || op == SAPFOR::CFG_OP::UN_MINUS ||
            op == SAPFOR::CFG_OP::NOT || op == SAPFOR::CFG_OP::ASSIGN;
        };

    while (!workStack.empty()) 
    {
        auto variable = workStack.top();
        workStack.pop();
        if (!variable || visited.count(variable))
            continue;

        visited.insert(variable);

        for (auto block : blocks) 
        {
            for (auto instrWrapper : block->getInstructions()) 
            {
                auto instr = instrWrapper->getInstruction();
                if (!instr || instr->getResult() != variable)
                    continue;

                auto op = instr->getOperation();
                auto arg1 = instr->getArg1();
                auto arg2 = instr->getArg2();

                if (isBinaryOp(op) && arg1 && arg2) 
                {
                    if (arg1->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
                        result.insert(arg1);
                    else if (arg1->getType() == SAPFOR::CFG_ARG_TYPE::REG)
                        workStack.push(arg1);

                    if (arg2->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
                        result.insert(arg2);
                    else if (arg2->getType() == SAPFOR::CFG_ARG_TYPE::REG)
                        workStack.push(arg2);
                }
                else if (isUnaryOp(op) && arg1) 
                {
                    if (arg1->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
                        result.insert(arg1);
                    else if (arg1->getType() == SAPFOR::CFG_ARG_TYPE::REG)
                        workStack.push(arg1);
                }
            }
        }
    }

    return result;
}

static std::vector<SAPFOR::Instruction*> getPhiArguments(SAPFOR::BasicBlock* block, SAPFOR::Instruction* phiInstr) 
{
    std::vector<SAPFOR::Instruction*> result;

    auto& instructions = block->getInstructions();
    bool collecting = false;
    for (int i = instructions.size() - 1; i >= 0; --i) {
        auto instr = instructions[i]->getInstruction();

        if (collecting) 
        {
            if (instr->getOperation() == SAPFOR::CFG_OP::PARAM) 
            {
                auto arg = instr->getArg1();
                if (arg)
                    result.push_back(instr);                
            }
            else 
                break;            
        }

        if (!instr)
            continue;

        if (instr == phiInstr) 
        {
            collecting = true;
            continue;
        }
    }

    std::reverse(result.begin(), result.end());
    return result;
}

SAPFOR::BasicBlock* findInstructionBlock(SAPFOR::Instruction* targetInstr, const std::vector<SAPFOR::BasicBlock*>& blocks)
{
    for (auto& block : blocks) 
    {
        for (auto& instrWrapper : block->getInstructions()) 
        {
            auto instr = instrWrapper->getInstruction();
            if (instr == targetInstr)
                return block;
        }
    }
    return NULL;
}

static SAPFOR::BasicBlock* findInstructionBlockByNumber(int number, const std::vector<SAPFOR::BasicBlock*>& blocks)
{
    for (auto& block : blocks) 
    {
        for (auto& instrWrapper : block->getInstructions()) 
        {
            auto instr = instrWrapper->getInstruction();
            if (instr->getNumber() == number)
                return block;
        }
    }
    return NULL;
}

static void findInductiveVars(const std::vector<SAPFOR::BasicBlock*>& blocks, 
                              const std::vector<SAPFOR::BasicBlock*>& Loopblocks, SAPFOR::BasicBlock* loopHeader, 
                              const set<SAPFOR::BasicBlock*>& loopExits) 
{
    set<string> inductiveVars;
    set<SAPFOR::BasicBlock*> relevantBlocks = { loopHeader };

    for (auto block : relevantBlocks) 
    {
        for (auto instrWrapper : block->getInstructions()) 
        {
            auto instr = instrWrapper->getInstruction();
            if (!instr) 
                continue;

            auto op = instr->getOperation();
            auto res = instr->getResult();
            auto arg1 = instr->getArg1();
            auto arg2 = instr->getArg2();

            if (op == SAPFOR::CFG_OP::JUMP_IF) 
            {
                if (arg1 && arg1->getType() == SAPFOR::CFG_ARG_TYPE::VAR) 
                    inductiveVars.insert(arg1->getValue());
                
                if (arg1 && arg1->getType() == SAPFOR::CFG_ARG_TYPE::REG) 
                {
                    auto foundVariables = findRegisterSourceVariables(blocks, arg1);
                    for (auto var : foundVariables)
                        inductiveVars.insert(var->getValue());
                }
            }
        }
    }

    set<string> finalInductiveVars;

    for (auto instrWrapper : loopHeader->getInstructions()) 
    {
        auto instr = instrWrapper->getInstruction();
        if (!instr || instr->getOperation() != SAPFOR::CFG_OP::F_CALL || !instr->getArg1() || instr->getArg1()->getValue() != "FI_FUNCTION") 
            continue;

        auto phiRes = instr->getResult();
        if (!phiRes || !inductiveVars.count(phiRes->getValue()))
            continue;

        auto currentBlock = findInstructionBlock(instr, blocks);
        if (!currentBlock)
            continue;

        auto phiArgs = getPhiArguments(currentBlock, instr);
        bool hasInLoopDefinition = false;

        for (const auto& argInstr : phiArgs) 
        {
            if (!argInstr) 
                continue;

            int definitionInstrNumber = stoi(argInstr->getArg1()->getValue());
            if (definitionInstrNumber == -1)
                continue;

            auto phiBlock = findInstructionBlockByNumber(definitionInstrNumber, blocks);
            if (!phiBlock)
                continue;

            if (std::find(Loopblocks.begin(), Loopblocks.end(), phiBlock) != Loopblocks.end())
                hasInLoopDefinition = true;
        }

        if (hasInLoopDefinition)
            finalInductiveVars.insert(phiRes->getValue());
    }

    for (auto i : finalInductiveVars)
        cout << "Confirmed inductive variable: " << i << endl;

    if (finalInductiveVars.empty())
        cout << "No confirmed inductive variables found." << endl;
}

static SAPFOR::Instruction* findInstructionAfterLoop(const std::vector<SAPFOR::BasicBlock*>& loopBody) 
{
    set<SAPFOR::BasicBlock*> loopSet(loopBody.begin(), loopBody.end());

    for (auto block : loopBody) 
    {
        for (auto succ : block->getNext()) 
        {
            if (!loopSet.count(succ)) 
            {
                // <20><><EFBFBD><EFBFBD><EFBFBD> <20><><EFBFBD><EFBFBD><EFBFBD> <20><> <20><><EFBFBD><EFBFBD><EFBFBD> <20> <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD> <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD> <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
                auto instructions = succ->getInstructions();
                for (auto wrapper : instructions)
                    if (auto instr = wrapper->getInstruction())
                        return instr;
            }
        }
    }

    return NULL; // <20><> <20><><EFBFBD><EFBFBD><EFBFBD>
}

void findImplicitLoops(const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& fullIR_SSA, const char* fileName)
{
    for (auto& i : fullIR_SSA)
    {
        //for (auto j : i.second)
        //    printblock(j);

        if (fileName != i.first->fileName)
            continue;

        map<int, int> visited;
        for (auto i : i.second)
            visited[i->getNumber()] = UNVISITED;

        //vector<int> visited(i.second.size(), UNVISITED);
        vector<pair<SAPFOR::BasicBlock*, SAPFOR::BasicBlock*>> startAndEnd;
        dfs(i.second[0], visited, startAndEnd, NULL);

        vector<LoopGraph*> loops;

        for (auto& [tail, header] : startAndEnd) {
            set<SAPFOR::BasicBlock*> loopExits;

            for (auto succ : tail->getNext()) {
                if (succ != header) {
                    loopExits.insert(succ);
                }
            }

            vector<SAPFOR::BasicBlock*> loopBody;
            getLoopBody(header, loopExits, loopBody);

            cout << "LOOP DETECTED:" << endl;
            cout << "  Header: " << header->getNumber() << endl;
            cout << "  Tail: " << tail->getNumber() << endl;
            cout << "  Body blocks: ";
            for (auto block : loopBody) {
                cout << block->getNumber() << " ";
            }
            cout << endl;

            findInductiveVars(i.second, loopBody, header, loopExits);

            SAPFOR::Instruction* instructionAfterLoop = findInstructionAfterLoop(loopBody);

            if (instructionAfterLoop == NULL) {
                cout << "Warning: instruction after loop not found!" << endl;
                continue;
            }

            auto firstInstruction = header->getInstructions()[0]->getInstruction();
            auto lastInstruction = tail->getInstructions().back()->getInstruction();

            cout << "first - " << firstInstruction->getNumber() << " last - " << lastInstruction->getNumber() << " after - " << instructionAfterLoop->getNumber() << endl;
            auto x = firstInstruction->getOperator();

            auto tmpLoop = new LoopGraph();
            tmpLoop->isFor = true;
            tmpLoop->lineNum = firstInstruction->getOperator()->lineNumber();
            tmpLoop->lineNumAfterLoop = instructionAfterLoop->getOperator()->lineNumber();

            if (firstInstruction->getOperator()->variant() == FOR_NODE) {
                SgForStmt* stmt = isSgForStmt(firstInstruction->getOperator());
                cout << "for loop" << endl;// << stmt->sunparse() << endl;
            }
            else if (firstInstruction->getOperator()->variant() == WHILE_NODE) {
                SgWhileStmt* stmt = isSgWhileStmt(firstInstruction->getOperator());
                cout << (stmt->conditional() == NULL ? "infinit" : "") << "while loop" << endl;//<< stmt->sunparse() << endl;
            }
            else if (firstInstruction->getOperator()->variant() == DO_WHILE_NODE) {
                SgWhileStmt* stmt = isSgDoWhileStmt(firstInstruction->getOperator());
                cout << "do while loop" << endl;// << stmt->sunparse() << endl;
            }
            else if (firstInstruction->getOperator()->variant() == LOOP_NODE) {
                cout << "not known loop" << endl;// << firstInstruction->getOperator()->sunparse() << endl;
            }
            else {
                cout << "goto loop" << endl;// firstInstruction->getOperator()->sunparse() << endl;
            }

            cout << "loop start line " << tmpLoop->lineNum << endl;
            cout << "after loop line " << tmpLoop->lineNumAfterLoop << endl << endl;

            loops.push_back(tmpLoop);
        }
    }
}