#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);
}


vector<SAPFOR::BasicBlock*> findBlocksInLoopsByFullIR(
    SgStatement* funcStmt,
    const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR)
{
    vector<SAPFOR::BasicBlock*> result;
    if (!funcStmt)
        return result;

    const vector<SAPFOR::BasicBlock*> funcBlocks = findFuncBlocksByFuncStatement(funcStmt, FullIR);
    const auto loopsMapping = findAndAnalyzeLoops(funcStmt, funcBlocks);

    set<SAPFOR::BasicBlock*> uniq;
    for (const auto& kv : loopsMapping)
        for (auto* bb : kv.second)
            if (bb)
                uniq.insert(bb);

    result.assign(uniq.begin(), uniq.end());
    sort(result.begin(), result.end(),
         [](const SAPFOR::BasicBlock* a, const SAPFOR::BasicBlock* b)
         {
             if (!a || !b)
                 return a < b;
             return a->getNumber() < b->getNumber();
         });

    return result;
}

static map<SgStatement*, vector<SgStatement*>> analyzeBasicBlockIntraDependencies(const SAPFOR::BasicBlock* bb)
{
    map<SgStatement*, vector<SgStatement*>> result;
    if (!bb)
        return result;

    auto isCompoundStmt = [](SgStatement* st) -> bool
    {
        if (!st)
            return true;
        const int v = st->variant();
        return loop_tags.count(v) || control_tags.count(v) || control_end_tags.count(v);
    };

    auto isTrackable = [](const SAPFOR::Argument* a) -> bool
    {
        if (!a)
            return false;
        const auto t = a->getType();
        return t == SAPFOR::CFG_ARG_TYPE::VAR || t == SAPFOR::CFG_ARG_TYPE::REG;
    };

    auto argKey = [&](const SAPFOR::Argument* a) -> string
    {
        if (!a)
            return string();
        return to_string((int)a->getType()) + "#" + to_string((int)a->getMemType()) + "#" + a->getValue();
    };

    auto memKeyFromInstr = [&](const SAPFOR::Instruction* instr) -> string
    {
        if (!instr)
            return string();
        SgExpression* ex = instr->getExpression();
        if (!ex || !ex->unparse())
            return string();
        auto normalizeExprText = [](const string& raw) -> string
        {
            string t;
            t.reserve(raw.size());
            for (unsigned char c : raw)
                if (!std::isspace(c))
                    t.push_back((char)c);

            auto stripOneLayer = [](const string& x) -> string
            {
                if (x.size() < 2 || x.front() != '(' || x.back() != ')')
                    return x;
                int bal = 0;
                for (size_t i = 0; i + 1 < x.size(); ++i)
                {
                    if (x[i] == '(') bal++;
                    else if (x[i] == ')') bal--;
                    if (bal == 0 && i + 1 < x.size() - 1)
                        return x;
                }
                return x.substr(1, x.size() - 2);
            };
            while (true)
            {
                const string stripped = stripOneLayer(t);
                if (stripped == t)
                    break;
                t = stripped;
            }
            return t;
        };
        return "MEMEX#" + normalizeExprText(ex->unparse());
    };

    auto isBarrier = [&](const SAPFOR::Instruction* instr) -> bool
    {
        if (!instr)
            return true;
        const auto op = instr->getOperation();
        switch (op)
        {
        case SAPFOR::CFG_OP::F_CALL:
        case SAPFOR::CFG_OP::IO_PARAM:
        case SAPFOR::CFG_OP::DVM_DIR:
        case SAPFOR::CFG_OP::SPF_DIR:
        case SAPFOR::CFG_OP::POINTER_ASS:
        case SAPFOR::CFG_OP::EXIT:
            return true;
        default:
            return false;
        }
    };

    auto isDef = [&](const SAPFOR::Instruction* instr) -> bool
    {
        if (!instr)
            return false;
        SAPFOR::Argument* r = instr->getResult();
        if (!isTrackable(r))
            return false;

        const auto op = instr->getOperation();
        if (op == SAPFOR::CFG_OP::STORE || op == SAPFOR::CFG_OP::REC_REF_STORE)
            return false;

        return true;
    };

    // Reaching definitions inside the BasicBlock in straight-line order:
    // lastDef[var] = last operator in this block that defined it.
    map<string, pair<SgStatement*, const SAPFOR::Argument*>> lastDef;
    map<string, pair<SgStatement*, const SAPFOR::Argument*>> lastMemDef;
    map<SgStatement*, set<SgStatement*>> depsSets;

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

        const SAPFOR::Instruction* instr = ir->getInstruction();
        SgStatement* opStmt = instr->getOperator();
        if (!opStmt)
            continue;
        if (isCompoundStmt(opStmt))
            continue;

        if (isBarrier(instr))
        {
            for (auto it = lastDef.begin(); it != lastDef.end();)
            {
                const SAPFOR::Argument* a = it->second.second;
                if (!a || a->isMemGlobal() || a->isParameter())
                    it = lastDef.erase(it);
                else
                    ++it;
            }
            for (auto it = lastMemDef.begin(); it != lastMemDef.end();)
            {
                const SAPFOR::Argument* a = it->second.second;
                if (!a || a->isMemGlobal() || a->isParameter())
                    it = lastMemDef.erase(it);
                else
                    ++it;
            }
        }

        if (!result.count(opStmt))
            result[opStmt] = {};

        auto addDep = [&](SAPFOR::Argument* use)
        {
            if (!isTrackable(use))
                return;
            const string k = argKey(use);
            auto it = lastDef.find(k);
            if (it == lastDef.end())
                return;
            if (it->second.first && it->second.first != opStmt)
                depsSets[opStmt].insert(it->second.first);
        };

        auto addMemDep = [&](const string& key)
        {
            if (key.empty())
                return;
            auto it = lastMemDef.find(key);
            if (it == lastMemDef.end())
                return;
            if (it->second.first && it->second.first != opStmt)
                depsSets[opStmt].insert(it->second.first);
        };

        addDep(instr->getArg1());
        addDep(instr->getArg2());
        if (instr->getOperation() == SAPFOR::CFG_OP::STORE || instr->getOperation() == SAPFOR::CFG_OP::REC_REF_STORE)
            addDep(instr->getResult());

        if (instr->getOperation() == SAPFOR::CFG_OP::LOAD || instr->getOperation() == SAPFOR::CFG_OP::REC_REF_LOAD)
        {
            const string memKey = memKeyFromInstr(instr);
            addMemDep(memKey);
        }

        if (isDef(instr))
        {
            const string dk = argKey(instr->getResult());
            lastDef[dk] = { opStmt, instr->getResult() };
        }

        if (instr->getOperation() == SAPFOR::CFG_OP::STORE || instr->getOperation() == SAPFOR::CFG_OP::REC_REF_STORE)
        {
            const string k = memKeyFromInstr(instr);
            SAPFOR::Argument* base = instr->getArg1();
            if (!k.empty() && base)
                lastMemDef[k] = { opStmt, base };
            addMemDep(k);
        }
    }

    for (auto& kv : result)
    {
        SgStatement* op = kv.first;
        auto it = depsSets.find(op);
        if (it == depsSets.end())
            continue;

        kv.second.assign(it->second.begin(), it->second.end());
        sort(kv.second.begin(), kv.second.end(),
             [](SgStatement* a, SgStatement* b)
             {
                 const int la = a ? a->lineNumber() : -1;
                 const int lb = b ? b->lineNumber() : -1;
                 if (la != lb)
                     return la < lb;
                 return a < b;
             });
    }

    return result;
}

static bool reorderOperatorsInBasicBlockUsingDeps(SAPFOR::BasicBlock* bb)
{
    if (!bb)
        return false;

    auto isCompoundStmt = [](SgStatement* st) -> bool
    {
        if (!st)
            return true;
        const int v = st->variant();
        return loop_tags.count(v) || control_tags.count(v) || control_end_tags.count(v);
    };

    vector<SgStatement*> ops;
    ops.reserve(bb->getInstructions().size());
    set<SgStatement*> seen;
    for (auto* ir : bb->getInstructions())
    {
        if (!ir || !ir->getInstruction())
            continue;
        SgStatement* st = ir->getInstruction()->getOperator();
        if (!st || isCompoundStmt(st))
            continue;
        if (seen.insert(st).second)
            ops.push_back(st);
    }

    if (ops.size() < 2)
        return false;

    SgStatement* parent = ops.front()->controlParent();
    if (!parent)
        return false;
    for (auto* st : ops)
        if (!st || st->controlParent() != parent)
            return false;

    set<SgStatement*> opSet(ops.begin(), ops.end());
    {
        SgStatement* lastNode = parent->lastNodeOfStmt();
        SgStatement* cur = ops.front();
        size_t idx = 0;

        {
            SgStatement* scan = parent->lexNext();
            SgStatement* end = lastNode;
            bool found = false;
            set<SgStatement*> visited;
            while (scan && scan != end)
            {
                if (!visited.insert(scan).second)
                    return false;
                if (scan == ops.front())
                {
                    found = true;
                    break;
                }
                scan = scan->lexNext();
            }
            if (!found)
                return false;
        }

        while (cur && cur != lastNode && idx < ops.size())
        {
            if (cur == ops[idx])
                ++idx;
            else if (isSgExecutableStatement(cur) && !opSet.count(cur))
                return false;
            if (idx == ops.size())
                break;
            cur = cur->lexNext();
        }
        if (idx != ops.size())
            return false;
    }

    // Compute dependencies (inside BB) and build a new order by moving each statement
    // as close as possible after its last dependency (if any).
    const auto depsMap = analyzeBasicBlockIntraDependencies(bb);
    vector<SgStatement*> order = ops;

    auto buildPos = [&](const vector<SgStatement*>& v)
    {
        map<SgStatement*, int> pos;
        for (int i = 0; i < (int)v.size(); ++i)
            pos[v[i]] = i;
        return pos;
    };

    const int maxIterations = (int)order.size() * (int)order.size() + 10;
    bool anyMove = false;
    for (int iter = 0; iter < maxIterations; ++iter)
    {
        bool movedThisIter = false;
        const auto pos = buildPos(order);

        for (int i = 0; i < (int)order.size(); ++i)
        {
            SgStatement* curSt = order[i];
            auto it = depsMap.find(curSt);
            if (it == depsMap.end() || it->second.empty())
                continue;

            int lastDepIdx = -1;
            for (SgStatement* dep : it->second)
            {
                auto itP = pos.find(dep);
                if (itP != pos.end())
                    lastDepIdx = max(lastDepIdx, itP->second);
            }
            if (lastDepIdx < 0)
                continue;

            int target = lastDepIdx + 1;
            if (target == i)
                continue;

            SgStatement* moved = order[i];
            order.erase(order.begin() + i);
            if (target > i)
                target -= 1;
            if (target < 0)
                target = 0;
            if (target > (int)order.size())
                target = (int)order.size();
            order.insert(order.begin() + target, moved);

            movedThisIter = true;
            anyMove = true;
            break;
        }

        if (!movedThisIter)
            break;
    }

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

    SgStatement* anchor = parent;
    {
        SgStatement* scan = parent->lexNext();
        SgStatement* end = parent->lastNodeOfStmt();
        while (scan && scan != end)
        {
            if (scan == ops.front())
                break;
            anchor = scan;
            scan = scan->lexNext();
        }
    }

    // apply AST reordering
    map<SgStatement*, int> savedLine;
    map<SgStatement*, char*> savedComments;
    map<SgStatement*, SgStatement*> extracted;

    for (SgStatement* st : ops)
    {
        if (!st)
            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 : order)
    {
        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, *parent);
        insertAfter = ex;
    }

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

    return true;
}

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

    if (!file)
        return;

    const int funcNum = file->numberOfFunctions();

    for (int i = 0; i < funcNum; ++i) {
        SgStatement* st = file->functions(i);

        const auto loopBlocks = findBlocksInLoopsByFullIR(st, FullIR);
        for (auto* bb : loopBlocks)
        {
            if (!bb)
                continue;
            if (reorderOperatorsInBasicBlockUsingDeps(bb))
                countOfTransform += 1;
        }

        // 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);
    }
}