SAPFOR/src/Transformations/MoveOperators/move_operators.cpp

#include <map>
#include <set>
#include <vector>
#include <queue>
#include <iostream>
#include <algorithm>
#include <limits>
#include <cstring>
#include <cstdlib>

#include "../../Utils/errors.h"
#include "../../Utils/SgUtils.h"
#include "../../GraphCall/graph_calls.h"
#include "../../GraphCall/graph_calls_func.h"
#include "../../CFGraph/CFGraph.h"
#include "../../CFGraph/IR.h"
#include "../../GraphLoop/graph_loops.h"
#include "move_operators.h"

using namespace std;

string getNameByArg(SAPFOR::Argument* arg);

static vector<SAPFOR::IR_Block*> findInstructionsFromStatement(SgStatement* st, const vector<SAPFOR::BasicBlock*>& blocks)
{
    vector<SAPFOR::IR_Block*> result;
    if (!st)
        return result;

    const int stmtId = st->id();
    for (auto* bb : blocks)
    {
        if (!bb)
            continue;

        for (auto* ir : bb->getInstructions())
        {
            if (!ir || !ir->getInstruction())
                continue;

            SgStatement* op = ir->getInstruction()->getOperator();
            if (op && op->id() == stmtId)
                result.push_back(ir);
        }
    }

    sort(result.begin(), result.end(),
         [](const SAPFOR::IR_Block* a, const SAPFOR::IR_Block* b) { return a->getNumber() < b->getNumber(); });
    return result;
}

set<int> loop_tags = {FOR_NODE};
set<int> control_tags = {IF_NODE, ELSEIF_NODE, DO_WHILE_NODE, WHILE_NODE, LOGIF_NODE};
set<int> control_end_tags = {CONTROL_END};

static bool isBarrierIR(const SAPFOR::Instruction* instr)
{
    if (!instr)
        return true;

    using SAPFOR::CFG_OP;
    const auto op = instr->getOperation();

    switch (op)
    {
    case CFG_OP::F_CALL:
    case CFG_OP::STORE:
    case CFG_OP::REC_REF_STORE:
    case CFG_OP::IO_PARAM:
    case CFG_OP::DVM_DIR:
    case CFG_OP::SPF_DIR:
    case CFG_OP::POINTER_ASS:
    case CFG_OP::EXIT:
        return true;
    default:
        return false;
    }
}

static bool isMovableStmtIR(const vector<SAPFOR::IR_Block*>& irs,
                            set<SAPFOR::Argument*>& uses,
                            set<SAPFOR::Argument*>& defs,
                            bool& spansMultipleBB)
{
    uses.clear();
    defs.clear();
    spansMultipleBB = false;

    if (irs.empty())
        return false;

    SAPFOR::BasicBlock* bb = irs.front()->getBasicBlock();
    for (auto* ir : irs)
        if (ir && ir->getBasicBlock() != bb)
            spansMultipleBB = true;

    for (auto* ir : irs)
    {
        if (!ir || !ir->getInstruction())
            return false;

        const auto* instr = ir->getInstruction();
        if (isBarrierIR(instr))
            return false;

        auto addUse = [&](SAPFOR::Argument* a)
        {
            if (!a)
                return;
            if (a->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
                uses.insert(a);
            else if (a->getType() == SAPFOR::CFG_ARG_TYPE::ARRAY || a->getType() == SAPFOR::CFG_ARG_TYPE::RECORD)
                uses.insert(a);
        };

        auto addDef = [&](SAPFOR::Argument* a)
        {
            if (!a)
                return;
            if (a->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
                defs.insert(a);
            else if (a->getType() == SAPFOR::CFG_ARG_TYPE::ARRAY || a->getType() == SAPFOR::CFG_ARG_TYPE::RECORD)
                defs.insert(a);
        };

        addUse(instr->getArg1());
        addUse(instr->getArg2());
        addDef(instr->getResult());
    }

    for (auto* a : uses)
        if (a && (a->getType() == SAPFOR::CFG_ARG_TYPE::ARRAY || a->getType() == SAPFOR::CFG_ARG_TYPE::RECORD))
            return false;
    for (auto* a : defs)
        if (a && (a->getType() == SAPFOR::CFG_ARG_TYPE::ARRAY || a->getType() == SAPFOR::CFG_ARG_TYPE::RECORD))
            return false;

    if (spansMultipleBB)
        return false;

    if (defs.empty())
        return false;

    return true;
}

static bool isStatementEmbedded(SgStatement* stmt, SgStatement* parent) {
    if (!stmt || !parent || stmt == parent)
        return false;
    
    if (parent->variant() == LOGIF_NODE) {
        if (stmt->lineNumber() == parent->lineNumber())
            return true;
        
        SgStatement* current = parent;
        SgStatement* lastNode = parent->lastNodeOfStmt();

        while (current && current != lastNode) {
            if (current == stmt)
                return true;

            if (current->isIncludedInStmt(*stmt))
                return true;

            current = current->lexNext();
        }
    }
    
    if (parent->isIncludedInStmt(*stmt))
        return true;
    
    return false;
}

static bool isLoopBoundary(SgStatement* stmt) {
    if (!stmt)
        return false;
    
    if (stmt->variant() == FOR_NODE || stmt->variant() == CONTROL_END)
        return true;
    
    return false;
}

static bool isPartOfNestedLoop(SgStatement* stmt, SgForStmt* loop) {
    if (!stmt || !loop)
        return false;
    
    SgStatement* loopStart = loop->lexNext();
    SgStatement* loopEnd = loop->lastNodeOfStmt();
    if (!loopStart || !loopEnd)
        return false;
    
    if (stmt->lineNumber() < loopStart->lineNumber() || stmt->lineNumber() > loopEnd->lineNumber())
        return false;
    
    SgStatement* current = loopStart;
    
    while (current && current != loopEnd) {
        
        if (current->variant() == FOR_NODE && current != loop) {
            SgForStmt* nestedLoop = (SgForStmt*)current;
            SgStatement* nestedStart = nestedLoop->lexNext();
            SgStatement* nestedEnd = nestedLoop->lastNodeOfStmt();
            
            if (nestedStart && nestedEnd && 
                stmt->lineNumber() >= nestedStart->lineNumber() && 
                stmt->lineNumber() <= nestedEnd->lineNumber()) {
                return true;
            }
        }
        
        current = current->lexNext();
    }
    
    return false;
}

static bool canSafelyExtract(SgStatement* stmt, SgForStmt* loop) {
    if (!stmt || !loop)
        return false;
    
    if (isLoopBoundary(stmt))
        return false;
    
    if (control_tags.find(stmt->variant()) != control_tags.end())
        return false;
    
    if (isPartOfNestedLoop(stmt, loop))
        return false;
    
    SgStatement* loopStart = loop->lexNext();
    SgStatement* loopEnd = loop->lastNodeOfStmt();
    if (!loopStart || !loopEnd) 
        return false;
    
    SgStatement* current = loopStart;
    
    while (current && current != loopEnd) {
        if (current->variant() == LOGIF_NODE && current->lineNumber() == stmt->lineNumber())
            return false;

        if (control_tags.find(current->variant()) != control_tags.end())
            if (isStatementEmbedded(stmt, current))
                return false;

        if (current == stmt) 
            break;
        current = current->lexNext();
    }
    
    return true;
}

namespace
{
    struct StmtInfo
    {
        SgStatement* stmt = nullptr;
        SAPFOR::BasicBlock* bb = nullptr;
        int firstInstr = std::numeric_limits<int>::max();
        int lastInstr = std::numeric_limits<int>::min();
        vector<SAPFOR::IR_Block*> irs;
        set<SAPFOR::Argument*> uses;
        set<SAPFOR::Argument*> defs;
    };

    using RDState = map<SAPFOR::Argument*, set<int>>;

    static void killGlobalsAndParams(RDState& st)
    {
        for (auto it = st.begin(); it != st.end();)
        {
            SAPFOR::Argument* a = it->first;
            if (!a)
            {
                it = st.erase(it);
                continue;
            }

            const bool kill = a->isMemGlobal() || a->isParameter();
            if (kill)
                it = st.erase(it);
            else
                ++it;
        }
    }

    static void transferRD(RDState& st, const SAPFOR::Instruction* instr)
    {
        if (!instr)
            return;

        if (isBarrierIR(instr))
        {
            killGlobalsAndParams(st);
            return;
        }

        using SAPFOR::CFG_OP;
        const auto op = instr->getOperation();
        if (op == CFG_OP::ASSIGN || op == CFG_OP::LOAD)
        {
            SAPFOR::Argument* res = instr->getResult();
            if (res && res->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
                st[res] = { instr->getNumber() };
        }
    }

    static RDState computeRD_BeforeInstr(const SAPFOR::BasicBlock* bb, int beforeInstrNum)
    {
        RDState st;
        if (!bb)
            return st;

        st = bb->getRD_In();
        for (auto* ir : bb->getInstructions())
        {
            if (!ir || !ir->getInstruction())
                continue;
            if (ir->getNumber() >= beforeInstrNum)
                break;
            transferRD(st, ir->getInstruction());
        }
        return st;
    }

    static bool topoSchedule(const vector<SgStatement*>& original,
                             const vector<vector<int>>& succ,
                             const vector<int>& indegInit,
                             vector<SgStatement*>& scheduled)
    {
        const int n = (int)original.size();
        scheduled.clear();
        scheduled.reserve(n);

        vector<int> indeg = indegInit;
        vector<int> maxPredPos(n, 0);

        struct Node
        {
            int idx;
            int keyMaxPredPos;
        };
        struct Cmp
        {
            bool operator()(const Node& a, const Node& b) const
            {
                if (a.keyMaxPredPos != b.keyMaxPredPos)
                    return a.keyMaxPredPos < b.keyMaxPredPos;
                return a.idx > b.idx;
            }
        };

        std::priority_queue<Node, vector<Node>, Cmp> ready;
        for (int i = 0; i < n; ++i)
            if (indeg[i] == 0)
                ready.push(Node{ i, 0 });

        int outPos = 0;
        while (!ready.empty())
        {
            Node cur = ready.top();
            ready.pop();

            const int u = cur.idx;
            scheduled.push_back(original[u]);
            ++outPos;

            for (int v : succ[u])
            {
                maxPredPos[v] = std::max(maxPredPos[v], outPos);
                if (--indeg[v] == 0)
                    ready.push(Node{ v, maxPredPos[v] });
            }
        }

        return (int)scheduled.size() == n;
    }

    static bool applyReorderContiguous(SgForStmt* loop,
                                       const vector<SgStatement*>& oldOrder,
                                       const vector<SgStatement*>& newOrder)
    {
        if (!loop || oldOrder.size() != newOrder.size() || oldOrder.size() < 2)
            return false;

        bool changed = false;
        for (size_t i = 0; i < oldOrder.size(); ++i)
            if (oldOrder[i] != newOrder[i])
            {
                changed = true;
                break;
            }
        if (!changed)
            return false;

        SgStatement* loopStart = loop->lexNext();
        SgStatement* loopEnd = loop->lastNodeOfStmt();
        if (!loopStart || !loopEnd)
            return false;

        SgStatement* first = oldOrder.front();
        SgStatement* anchor = loop;
        for (SgStatement* cur = loopStart; cur && cur != loopEnd; cur = cur->lexNext())
        {
            if (cur == first)
                break;
            anchor = cur;
        }

        map<SgStatement*, int> savedLine;
        map<SgStatement*, char*> savedComments;
        map<SgStatement*, SgStatement*> extracted;

        for (SgStatement* st : oldOrder)
        {
            if (!st || st == loop || st == loopEnd)
                return false;
            if (!canSafelyExtract(st, loop))
                return false;

            savedLine[st] = st->lineNumber();
            savedComments[st] = st->comments() ? strdup(st->comments()) : nullptr;

            SgStatement* ex = st->extractStmt();
            if (!ex)
                return false;
            extracted[st] = ex;
        }

        SgStatement* insertAfter = anchor;
        for (SgStatement* st : newOrder)
        {
            SgStatement* ex = extracted[st];
            if (!ex)
                continue;

            auto itC = savedComments.find(st);
            if (itC != savedComments.end() && itC->second)
                ex->setComments(itC->second);

            auto itL = savedLine.find(st);
            if (itL != savedLine.end())
                ex->setlineNumber(itL->second);

            insertAfter->insertStmtAfter(*ex, *loop);
            insertAfter = ex;
        }

        for (auto& kv : savedComments)
            if (kv.second)
                free(kv.second);

        return true;
    }
}

static bool reorderMovableRegionsInLoop(SgForStmt* loop, const vector<SAPFOR::BasicBlock*>& blocks)
{
    if (!loop)
        return false;

    SgStatement* loopStart = loop->lexNext();
    SgStatement* loopEnd = loop->lastNodeOfStmt();
    if (!loopStart || !loopEnd)
        return false;

    bool anyChange = false;

    vector<SgStatement*> region;
    auto flushRegion = [&]() {
        if (region.size() < 2)
        {
            region.clear();
            return;
        }

        vector<StmtInfo> infos;
        infos.reserve(region.size());
        map<SgStatement*, int> idxOf;

        map<int, int> defInstrToIdx;

        for (int i = 0; i < (int)region.size(); ++i)
        {
            SgStatement* st = region[i];
            idxOf[st] = i;

            StmtInfo info;
            info.stmt = st;
            info.irs = findInstructionsFromStatement(st, blocks);
            if (info.irs.empty())
            {
                infos.clear();
                region.clear();
                return;
            }

            info.bb = info.irs.front()->getBasicBlock();
            for (auto* ir : info.irs)
            {
                info.firstInstr = std::min(info.firstInstr, ir->getNumber());
                info.lastInstr = std::max(info.lastInstr, ir->getNumber());

                const auto* instr = ir->getInstruction();
                if (!instr)
                    continue;

                if (instr->getResult() && instr->getResult()->getType() == SAPFOR::CFG_ARG_TYPE::VAR &&
                    (instr->getOperation() == SAPFOR::CFG_OP::ASSIGN || instr->getOperation() == SAPFOR::CFG_OP::LOAD))
                {
                    defInstrToIdx[instr->getNumber()] = i;
                }
            }

            bool spansMultiple = false;
            if (!isMovableStmtIR(info.irs, info.uses, info.defs, spansMultiple) || spansMultiple)
            {
                infos.clear();
                region.clear();
                return;
            }

            infos.push_back(std::move(info));
        }

        const int n = (int)infos.size();
        vector<vector<int>> succ(n);
        vector<int> indeg(n, 0);

        auto addEdge = [&](int u, int v) {
            if (u == v)
                return;
            succ[u].push_back(v);
            indeg[v]++;
        };

        for (int i = 0; i < n; ++i)
        {
            const StmtInfo& cur = infos[i];
            RDState stRD = computeRD_BeforeInstr(cur.bb, cur.firstInstr);

            for (SAPFOR::Argument* use : cur.uses)
            {
                if (!use || use->getType() != SAPFOR::CFG_ARG_TYPE::VAR)
                    continue;

                auto it = stRD.find(use);
                if (it == stRD.end())
                    continue;

                for (int defInstr : it->second)
                {
                    if (defInstr < 0)
                        continue;

                    auto itDef = defInstrToIdx.find(defInstr);
                    if (itDef != defInstrToIdx.end())
                        addEdge(itDef->second, i);
                }
            }
        }

        map<SAPFOR::Argument*, vector<int>> defsByVar;
        for (int i = 0; i < n; ++i)
            for (auto* d : infos[i].defs)
                if (d && d->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
                    defsByVar[d].push_back(i);

        for (auto& kv : defsByVar)
        {
            auto& list = kv.second;
            sort(list.begin(), list.end());
            list.erase(unique(list.begin(), list.end()), list.end());
            for (int k = 0; k + 1 < (int)list.size(); ++k)
                addEdge(list[k], list[k + 1]);
        }

        map<SAPFOR::Argument*, vector<int>> defPositions;
        for (auto& kv : defsByVar)
            defPositions[kv.first] = kv.second;

        for (int i = 0; i < n; ++i)
        {
            for (SAPFOR::Argument* use : infos[i].uses)
            {
                if (!use || use->getType() != SAPFOR::CFG_ARG_TYPE::VAR)
                    continue;

                auto itDefs = defPositions.find(use);
                if (itDefs == defPositions.end())
                    continue;

                const auto& dpos = itDefs->second;
                auto it = std::upper_bound(dpos.begin(), dpos.end(), i);
                if (it != dpos.end())
                    addEdge(i, *it);
            }
        }

        for (int u = 0; u < n; ++u)
        {
            auto& out = succ[u];
            sort(out.begin(), out.end());
            out.erase(unique(out.begin(), out.end()), out.end());
        }
        fill(indeg.begin(), indeg.end(), 0);
        for (int u = 0; u < n; ++u)
            for (int v : succ[u])
                indeg[v]++;

        vector<SgStatement*> scheduled;
        vector<SgStatement*> original = region;
        if (!topoSchedule(original, succ, indeg, scheduled))
        {
            region.clear();
            return;
        }

        if (applyReorderContiguous(loop, original, scheduled))
            anyChange = true;

        region.clear();
    };

    SgStatement* current = loopStart;
    set<SgStatement*> visited;

    while (current && current != loopEnd)
    {
        if (!visited.insert(current).second)
            break;

        if (isLoopBoundary(current))
        {
            flushRegion();
            current = current->lexNext();
            continue;
        }

        if (current->variant() == FOR_NODE && current != loop)
        {
            flushRegion();
            SgStatement* nestedEnd = current->lastNodeOfStmt();
            current = nestedEnd ? nestedEnd->lexNext() : current->lexNext();
            continue;
        }

        if (!isSgExecutableStatement(current) || control_tags.count(current->variant()))
        {
            flushRegion();
            current = current->lexNext();
            continue;
        }

        const bool isTopLevel = (current->controlParent() == loop);
        if (isTopLevel && current->variant() == ASSIGN_STAT && canSafelyExtract(current, loop))
        {
            auto irs = findInstructionsFromStatement(current, blocks);
            set<SAPFOR::Argument*> uses, defs;
            bool spansMultiple = false;
            if (isMovableStmtIR(irs, uses, defs, spansMultiple) && !spansMultiple)
            {
                region.push_back(current);
                current = current->lexNext();
                continue;
            }
        }

        flushRegion();
        current = current->lexNext();
    }

    flushRegion();
    return anyChange;
}

vector<SAPFOR::BasicBlock*> findFuncBlocksByFuncStatement(SgStatement *st, const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR) {
    vector<SAPFOR::BasicBlock*> result;
    Statement* forSt = (Statement*)st;

    for (auto& func: FullIR) {
        if (func.first->funcPointer->getCurrProcessFile() == forSt->getCurrProcessFile()
            && func.first->funcPointer->lineNumber() == forSt->lineNumber())
        {
            result = func.second;
        }
    }
    return result;
}

map<SgForStmt*, vector<SAPFOR::BasicBlock*>> findAndAnalyzeLoops(SgStatement *st, const vector<SAPFOR::BasicBlock*>& blocks) {
    map<SgForStmt*, vector<SAPFOR::BasicBlock*>> result;
    SgStatement *lastNode = st->lastNodeOfStmt();
    
    while (st && st != lastNode) {
        if (loop_tags.find(st -> variant()) != loop_tags.end()) {
            SgForStmt *forSt = (SgForStmt*)st;
            SgStatement *loopBody = forSt -> body();
            SgStatement *lastLoopNode = st->lastNodeOfStmt();
            set<int> blocks_nums;
            
            while (loopBody && loopBody != lastLoopNode) {
                vector<SAPFOR::IR_Block*> irBlocks = findInstructionsFromStatement(loopBody, blocks);
                if (!irBlocks.empty()) {
                    SAPFOR::IR_Block* IR = irBlocks.front();
                    if (IR && IR->getBasicBlock()) {
                        if (blocks_nums.find(IR -> getBasicBlock() -> getNumber()) == blocks_nums.end()) {
                            result[forSt].push_back(IR -> getBasicBlock());
                            blocks_nums.insert(IR -> getBasicBlock() -> getNumber());
                        }
                    }
                }
                loopBody = loopBody -> lexNext();
            }
            sort(result[forSt].begin(), result[forSt].end());
        }
        st = st -> lexNext();
    }
    return result;
}

static void processLoopRecursively(SgForStmt* loop, const vector<SAPFOR::BasicBlock*>& blocks, 
                                   const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR) {
    if (!loop) 
        return;
    
    SgStatement* loopStart = loop->lexNext();
    SgStatement* loopEnd = loop->lastNodeOfStmt();
    if (loopStart && loopEnd) {
        SgStatement* current = loopStart;
        while (current && current != loopEnd) {
            if (current->variant() == FOR_NODE && current != loop) {
                SgForStmt* nestedLoop = (SgForStmt*)current;
                processLoopRecursively(nestedLoop, blocks, FullIR);
                SgStatement* nestedEnd = nestedLoop->lastNodeOfStmt();

                if (nestedEnd)
                    current = nestedEnd->lexNext();
                else
                    current = current->lexNext();
            } 
            else
                current = current->lexNext();
        }
    }

    reorderMovableRegionsInLoop(loop, blocks);
}

void moveOperators(SgFile *file, map<string, vector<LoopGraph*>>& loopGraph, 
                   const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR, 
                   int& countOfTransform) {
    countOfTransform += 1;

    const int funcNum = file -> numberOfFunctions();

    for (int i = 0; i < funcNum; ++i) {
        SgStatement *st = file -> functions(i);
        
        vector<SAPFOR::BasicBlock*> blocks = findFuncBlocksByFuncStatement(st, FullIR);        
        map<SgForStmt*, vector<SAPFOR::BasicBlock*>> loopsMapping = findAndAnalyzeLoops(st, blocks);
        
        for (auto& loopForAnalyze: loopsMapping)
            processLoopRecursively(loopForAnalyze.first, loopForAnalyze.second, FullIR);
    }
}