#include "replace_dist_arrays_in_io.h"

#include <string>
#include <map>
#include <set>

#include "expr_transform.h"
#include "../ParallelizationRegions/resolve_par_reg_conflicts.h"

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

#define DEBUG_TRACE 0

static bool hasArrayRef(SgExpression* ex)
{
    if (ex)
        return isArrayRef(ex) || hasArrayRef(ex->lhs()) || hasArrayRef(ex->rhs());

    return false;
}

static inline bool isArrayDeclaration(SgStatement* st)
{
    return isSgDeclarationStatement(st) || isSgVarListDeclStmt(st) || isSgNestedVarListDeclStmt(st);
}

static SgExpression* findExprWithVariant(SgExpression* exp, int variant)
{
    if (exp)
    {
        if (exp->variant() == variant)
            return exp;

        auto *l = findExprWithVariant(exp->lhs(), variant);
        if (l)
            return l;

        auto *r = findExprWithVariant(exp->rhs(), variant);
        if (r)
            return r;
    }

    return NULL;
}

static bool switchToDeclarationFile(DIST::Array* array_p, 
                                    pair<string, int>& decl_place, 
                                    const string &current_file_name,
                                    FuncInfo *current_func)
{
    if (!array_p)
        return false;

    // try to find declaration from current function
    for (const auto &p : array_p->GetDeclInfo())
    {
        if (p.first == current_file_name && 
            current_func->linesNum.first <= p.second && 
            p.second <= current_func->linesNum.second)
        {
            decl_place = p;
            return true;
        }
    }

    for (const auto &p : array_p->GetDeclInfo())
    {
        if (SgFile::switchToFile(p.first) != -1)
        {
            decl_place = p;
            return true;
        }
    }

    return false;
}

static bool checkAssumedShape(SgStatement* decl, const string& array_name)
{
    SgExpression* list = decl->expr(0);
    SgExpression* dim_list = NULL;

    while (list)
    {
        auto *arr_ref = list->lhs();

        if (arr_ref && 
            arr_ref->symbol() && 
            arr_ref->symbol()->identifier() && 
            arr_ref->symbol()->identifier() == array_name)
        {
            dim_list = arr_ref->lhs();
            break;
        }

        list = list->rhs();
    }

    if (!dim_list)
    {
        auto *dim_expr = findExprWithVariant(decl->expr(2), DIMENSION_OP);
        if (dim_expr)
            dim_list = dim_expr->lhs();
    }

    while (dim_list)
    {
        auto *lhs = dim_list->lhs();

        if (lhs && lhs->variant() == DDOT && (!lhs->lhs() || !lhs->rhs()))
            return true;

        dim_list = dim_list->rhs();
    }

    return false;
}

static bool checkAssumedSize(const string &array_name, DIST::Array* array_p, const string &current_file_name, FuncInfo *current_func)
{
    if (!array_p)
        return true; // raise true to prevent array from copying

    bool found = false;

    pair<string, int> decl_place;

    if (!switchToDeclarationFile(array_p, decl_place, current_file_name, current_func))
        return true;

    auto *st = SgStatement::getStatementByFileAndLine(decl_place.first, decl_place.second);
    checkNull(st, convertFileName(__FILE__).c_str(), __LINE__);
    
    SgExpression* list = st->expr(0);    
    while (list)
    {
        if (list->lhs() && 
            list->lhs()->symbol() && 
            list->lhs()->symbol()->identifier() && 
            list->lhs()->symbol()->identifier() == array_name
        )
        {
            if(findExprWithVariant(list->lhs(), STAR_RANGE))
                found = true;

            break;
        }

        list = list->rhs();
    }

    if (!found)
    {
        auto *dim_expr = findExprWithVariant(st->expr(2), DIMENSION_OP);
        if (dim_expr && findExprWithVariant(dim_expr, STAR_RANGE))
            found = true;
    }

    if (SgFile::switchToFile(current_file_name) == -1)
        printInternalError(convertFileName(__FILE__).c_str(), __LINE__); // can't switch back to file with array usage

    return found;
}

static void findArrays(SgExpression* exp, set<SgSymbol*>& arrays)
{
    if (exp)
    {
        if (isArrayRef(exp))
            arrays.insert(exp->symbol());

        findArrays(exp->lhs(), arrays);
        findArrays(exp->rhs(), arrays);
    }
}

static bool populateDistributedIoArrays(map<SgSymbol*, set<SgStatement*>>& arrays, 
                                        SgStatement* stat, 
                                        const string& current_file_name,
                                        FuncInfo *current_func)
{
    auto var = stat->variant();

    if (var != READ_STAT && var != PRINT_STAT && var != WRITE_STAT)
        return false;

    // check if such IO allowed in dvm: 
    // list should consist only of single array and format string should be *

    bool need_replace = false;

    SgExpression* ioList = stat->expr(0);

    if (!ioList)
        return false;

    if (ioList->variant() != EXPR_LIST)
        return false;

    if (ioList->rhs() == NULL)
    {
        SgExpression* arg = ioList->lhs();
        if (!arg)
            return false;

        if (hasArrayRef(arg))
        {
            if (isArrayRef(arg) && arg->lhs())
                need_replace = true;
        }
        else
        {
            return false;
        }
    }
    else
    {
        need_replace = true;
    }

    if (!need_replace)
    {
        switch (var)
        {
        case PRINT_STAT:
        {
            SgExpression* fmt = stat->expr(1);
            if (!fmt || fmt->variant() != SPEC_PAIR || fmt->lhs()->variant() != KEYWORD_VAL)
                printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

            if (fmt->rhs()->variant() != KEYWORD_VAL || fmt->rhs()->sunparse() != "*")
                need_replace = true;

            break;
        }
        case READ_STAT:
        case WRITE_STAT:
        {
            SgExpression* spec = stat->expr(1);

            if (spec->variant() == SPEC_PAIR)
            {
                if (spec->rhs()->variant() != KEYWORD_VAL || spec->rhs()->sunparse() != "*")
                {
                    need_replace = true;
                    break;
                }
            }
            else
            {
                while (spec)
                {
                    auto *kv = spec->lhs();
                    if (!kv || kv->variant() != SPEC_PAIR || !kv->rhs())
                        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

                    if (kv->rhs()->variant() != KEYWORD_VAL || kv->rhs()->sunparse() != "*")
                    {
                        need_replace = true;
                        break;
                    }

                    spec = spec->rhs();
                }
            }
            break;
        }
        default:
            break;
        }
    }

    if (!need_replace)
        return false;

    bool ret = false;

    set<SgSymbol*> found_arrays;

    for (int i = 0; i < 3; i++)
        findArrays(stat->expr(i), found_arrays);

    for (auto* by_symb : found_arrays)
    {
        string array_name = string(by_symb->identifier());
        DIST::Array* array_p = getArrayFromDeclarated(declaratedInStmt(by_symb), array_name);
        if (array_p && 
            array_p->GetDistributeFlagVal() == Distribution::distFlag::IO_PRIV && 
            !checkAssumedSize(array_name, array_p, current_file_name, current_func))
        {
            auto inserted = arrays[by_symb].insert(stat).second;

            if (inserted)
                __spf_print(DEBUG_TRACE, "[%d]: add array %s %p\n", stat->lineNumber(), array_p->GetName().c_str(), by_symb);

            ret = true;
        }
    }

    __spf_print(DEBUG_TRACE, "[replace]\n");
    return ret;
}

static void replaceArrayRec(SgSymbol* arr, SgSymbol* replace_by, SgExpression* exp, bool& has_read, bool& has_write, bool from_read, bool from_write)
{
    if (!exp)
        return;

    if (isArrayRef(exp) && 
        strcmp(exp->symbol()->identifier(), arr->identifier()) == 0)
    {
        has_read |= from_read;
        has_write |= from_write;
        exp->setSymbol(replace_by);
    }

    switch (exp->variant())
    {
    case FUNC_CALL:
    {
        replaceArrayRec(arr, replace_by, exp->rhs(), has_read, has_write, true, false);
        replaceArrayRec(arr, replace_by, exp->lhs(), has_read, has_write, true, true);
        break;
    }
    case EXPR_LIST:
    {
        replaceArrayRec(arr, replace_by, exp->lhs(), has_read, has_write, from_read, from_write);
        replaceArrayRec(arr, replace_by, exp->rhs(), has_read, has_write, from_read, from_write);
        break;
    }
    default:
    {
        replaceArrayRec(arr, replace_by, exp->lhs(), has_read, has_write, true, false);
        replaceArrayRec(arr, replace_by, exp->rhs(), has_read, has_write, true, false);
        break;
    }
    }
}

static void replaceArrayRec(SgSymbol* arr, SgSymbol* replace_by, SgStatement* st, bool& has_read, bool& has_write)
{
    if (!st)
        return;

    switch (st->variant())
    {
    case ASSIGN_STAT:
    case READ_STAT:
    {
        replaceArrayRec(arr, replace_by, st->expr(0), has_read, has_write, false, true);
        replaceArrayRec(arr, replace_by, st->expr(1), has_read, has_write, true, false);
        break;
    }
    case PROC_STAT:
    case FUNC_STAT:
    {
        replaceArrayRec(arr, replace_by, st->expr(0), has_read, has_write, true, false);
        replaceArrayRec(arr, replace_by, st->expr(1), has_read, has_write, true, true);
        break;
    }
    default:
    {
        for (int i = 0; i < 3; i++)
            replaceArrayRec(arr, replace_by, st->expr(i), has_read, has_write, true, false);

        break;
    }
    }
}

static void copyArrayBetweenStatements(SgSymbol* replace_symb, 
                                       SgSymbol* replace_by,
                                       SgStatement* start, 
                                       SgStatement* last, 
                                       bool start_is_scope,
                                       FuncInfo *func_info)
{
    while (start->lexNext() && !isSgExecutableStatement(start->lexNext()))
        start = start->lexNext();

    bool has_read = false, has_write = false;

    for (auto* st = start->lexNext(); st && st != last; st = st->lexNext())
        replaceArrayRec(replace_symb, replace_by, st, has_read, has_write);


    if (has_read)
    {
        // A_copy = A
        SgAssignStmt* assign = new SgAssignStmt(*new SgArrayRefExp(*replace_by), *new SgArrayRefExp(*replace_symb));
        assign->setlineNumber(getNextNegativeLineNumber()); // before region
        auto* parent = start_is_scope ? start : start->controlParent();
        if (parent && parent->lastNodeOfStmt() == start)
            parent = parent->controlParent();
        checkNull(parent, convertFileName(__FILE__).c_str(), __LINE__);

        start->insertStmtAfter(*assign, *parent);
    }

    if (has_write)
    {
        for (int i = 0; i < func_info->funcParams.identificators.size(); i++)
        {
            if (func_info->funcParams.identificators[i] == replace_symb->identifier())
            {
                has_write &= func_info->funcParams.isArgOut(i);
                break;
            }
        }
    }

    if (has_write)
    {
        // A = A_reg
        SgAssignStmt* assign = new SgAssignStmt(*new SgArrayRefExp(*replace_symb), *new SgArrayRefExp(*replace_by));
        //TODO: bug with insertion
        //assign->setlineNumber(getNextNegativeLineNumber()); // after region
        last->insertStmtBefore(*assign, *(last->controlParent()));
    }
}

static void replaceArrayInFragment(SgSymbol* replace_symb, 
                                   const set<SgStatement*> usages, 
                                   SgSymbol* replace_by, 
                                   SgStatement* start, 
                                   SgStatement* last, 
                                   FuncInfo *func_info)
{
    while (start->lexNext() && (!isSgExecutableStatement(start->lexNext()) || start->lexNext()->variant() == ALLOCATE_STMT))
        start = start->lexNext();

    set<SgStatement*> not_opened, not_closed, copied;

    for (auto* it = start; it; it = it->controlParent())
        not_opened.insert(it);

    for (auto* it = last; it; it = it->controlParent())
        not_closed.insert(it);

    for (auto* io_stmt : usages)
    {
        bool already_copied = false;
        SgStatement* copy_scope = NULL;

        for (auto* par = io_stmt; par; par = par->controlParent())
        {
            if (copied.find(par) != copied.end())
            {
                already_copied = true;
                break;
            }
            else if (not_opened.find(par) != not_opened.end() || not_closed.find(par) != not_closed.end())
            {
                copy_scope = par;
                break;
            }
        }


        if (already_copied)
            continue;

        auto* scope_start = copy_scope, * scope_end = copy_scope->lastNodeOfStmt();
        __spf_print(DEBUG_TRACE, "[scope to copy] %d\n", copy_scope->lineNumber());

        if (not_opened.find(copy_scope) != not_opened.end() && start != copy_scope)
        {
            auto* from = start;
            for (auto* st = from; st; st = st->controlParent())
            {
                __spf_print(DEBUG_TRACE, "[find start of parent %d] %d\n", copy_scope->lineNumber(), st->lineNumber());

                if (st->controlParent() == copy_scope)
                {
                    scope_start = st->lastNodeOfStmt() ? st->lastNodeOfStmt() : st;
                    break;
                }
            }
        }

        if (not_closed.find(copy_scope) != not_closed.end() && last != copy_scope)
        {
            for (auto* st = last; st; st = st->controlParent())
            {
                __spf_print(DEBUG_TRACE, "[find end of parent %d] %d\n", copy_scope->lineNumber(), st->lineNumber());

                if (st->controlParent() == copy_scope)
                {
                    scope_end = st;
                    break;
                }
            }
        }

        __spf_print(DEBUG_TRACE, "[copy %s] [%d, %d]\n", replace_symb->identifier(), scope_start->lineNumber(), scope_end->lineNumber());
        copyArrayBetweenStatements(replace_symb, replace_by, scope_start, scope_end, copy_scope == scope_start, func_info);
        copied.insert(copy_scope);
    }
}

static bool ioRegionBorder(SgStatement* stat, SgStatement* last_io_bound)
{
    auto var = stat->variant();

    static const set<int> border_stats =
    {
        PROC_STAT,
        FUNC_STAT,
        PROG_HEDR,
        FUNC_HEDR,
        PROC_HEDR,
        FOR_NODE,
        LOOP_NODE,
        RETURN_NODE,
        RETURN_STAT,
        STOP_STAT,
        STOP_NODE,
        EXIT_STMT,
        EXIT_NODE,
        GOTO_NODE
    };

    if (border_stats.find(var) != border_stats.end())
        return true;

    if (stat->hasLabel())
        return true;

    if (last_io_bound && last_io_bound->lastNodeOfStmt() && last_io_bound->lastNodeOfStmt() == stat)
        return true;

    if (var == CONTROL_END && border_stats.find(stat->controlParent()->variant()) != border_stats.end())
        return true;

    return false;
}

FuncInfo* getCurrentFuncInfo(const vector<FuncInfo*>& fileFuncInfo, int line)
{
    for (auto* func : fileFuncInfo)
    {
        if (func->linesNum.first <= line && line <= func->linesNum.second)
            return func;
    }

    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
    return NULL;
}

void replaceDistributedArraysInIO(vector<ParallelRegion*>& regions,
                                  const map<string, vector<FuncInfo*>>& allFuncInfo,
                                  map<string, vector<Messages>>& SPF_messages,
                                  map<string, map<int, set<string>>>& newDeclsToInclude)
{
    map<SgSymbol*, SgSymbol*> created_copies;
    map<string, map<int, set<string>>> copied_syms;

    for (auto& region : regions)
    {
        __spf_print(DEBUG_TRACE, "[%s]: enter region\n", region->GetName().c_str());

        for (auto& lines_by_file : region->GetAllLinesToModify())
        {
            const auto& current_file_name = lines_by_file.first;

            if (SgFile::switchToFile(current_file_name) == -1)
                printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

            FuncInfo *current_func_info = NULL;

            auto func_info_it = allFuncInfo.find(current_file_name);
            if (func_info_it == allFuncInfo.end())
                printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

            auto *lbound_symb = new SgSymbol(FUNCTION_NAME, "lbound");
            auto *ubound_symb = new SgSymbol(FUNCTION_NAME, "ubound");

            for (auto& lines : lines_by_file.second)
            {
                __spf_print(DEBUG_TRACE, "[fragment] %s: %d:%d %d\n", current_file_name.c_str(), lines.lines.first,
                    lines.lines.second, lines.isImplicit());

                SgStatement* curr_stmt, * end;

                if (lines.isImplicit())
                {
                    curr_stmt = current_file->SgStatementAtLine(lines.lines.first);
                    end = current_file->SgStatementAtLine(lines.lines.second);

                    if (end)
                        end = end->lexNext();
                }
                else
                {
                    curr_stmt = lines.stats.first->GetOriginal();
                    end = lines.stats.second->GetOriginal()->lexNext();
                    current_func_info = getCurrentFuncInfo(func_info_it->second, curr_stmt->lineNumber());
                }

                map<SgSymbol*, set<SgStatement*>> need_replace;
                SgStatement* last_io_bound = NULL;

                while (curr_stmt != end)
                {
                    if (!curr_stmt)
                        break;

                    auto var = curr_stmt->variant();

                    if (var == PROC_HEDR || var == PROG_HEDR || var == FUNC_HEDR)
                    {
                        current_func_info = getCurrentFuncInfo(func_info_it->second, curr_stmt->lineNumber());

                        curr_stmt = curr_stmt->lexNext();
                        while (curr_stmt && !isSgExecutableStatement(curr_stmt))
                        {
                            last_io_bound = curr_stmt;
                            curr_stmt = curr_stmt->lexNext();
                        }

                        if (!curr_stmt)
                            break;
                    }

                    if (ioRegionBorder(curr_stmt, last_io_bound))
                    {
                        for (const auto& by_array_to_copy : need_replace)
                        {
                            auto* array_to_copy = by_array_to_copy.first;
                            auto it = created_copies.find(array_to_copy);

                            if (it == created_copies.end())
                            {
                                string array_name = string(array_to_copy->identifier());
                                DIST::Array* array_p = getArrayFromDeclarated(declaratedInStmt(array_to_copy), array_name);

                                // at this point all considered arrays must have DIST::Array* references
                                if (!array_p)
                                    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

                                // ... and be declared in switchable files (not in headers)
                                pair<string, int> decl_place;
                                if (!switchToDeclarationFile(array_p, decl_place, current_file_name, current_func_info))
                                    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

                                string suffix = "_io_l";

                                switch (array_p->GetLocation().first)
                                {
                                    case DIST::l_MODULE:
                                        suffix = "_io_m";
                                        break;
                                    case DIST::l_COMMON:
                                        suffix = "_io_c";
                                        break;
                                }
                                
                                auto copied_symbol = copyArray(decl_place, 
                                                               array_p, 
                                                               lines_by_file.second, 
                                                               suffix + to_string(region->GetId()), 
                                                               decl_place.first, 
                                                               newDeclsToInclude, 
                                                               copied_syms);
                            
                                // original declaration statement
                                auto *decl_stmt = SgStatement::getStatementByFileAndLine(decl_place.first, decl_place.second);

                                if (!decl_stmt || !isArrayDeclaration(decl_stmt))
                                    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

                                auto dynamic_array = checkAssumedShape(decl_stmt, array_name);  

                                while (decl_stmt && !isSgExecutableStatement(decl_stmt))
                                {
                                    if (isArrayDeclaration(decl_stmt) &&
                                        decl_stmt->expr(0) &&
                                        decl_stmt->expr(0)->lhs() &&
                                        decl_stmt->expr(0)->lhs()->symbol() == copied_symbol.second)
                                    {
                                        break;
                                    }

                                    decl_stmt = decl_stmt->lexNext();
                                }
                                
                                // created declaration statement
                                if (!decl_stmt || !isArrayDeclaration(decl_stmt))
                                    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

                                // insert !$SPF ANALYSIS(PROCESS_PRIVATE(array)) directive before declaration statement
                                
                                string dir_str;
                                if (decl_stmt->comments())
                                {
                                    auto str_comment = string(decl_stmt->comments());
                                    if (str_comment.size() && str_comment.back() != '\n')
                                        dir_str += "\n";
                                }

                                dir_str += "!$SPF ANALYSIS(PROCESS_PRIVATE(" + string(copied_symbol.second->identifier()) + "))\n";
                                decl_stmt->addComment(dir_str.c_str());
                            
                                created_copies.insert({ array_to_copy, copied_symbol.second });

                                // make array copy allocatable in case of assumed-shape array
                                if(dynamic_array)
                                {
                                    // insert allocatable keyword in declaration
                                    auto *kword_list = decl_stmt->expr(2);
                                    if (!findExprWithVariant(kword_list, ALLOCATABLE_OP))
                                    {
                                        if (!kword_list)
                                        {
                                            kword_list = new SgExprListExp();
                                            decl_stmt->setExpression(2, *kword_list);
                                        }
                                        
                                        while (kword_list->rhs())
                                            kword_list = kword_list->rhs();
                                        
                                        if (kword_list->lhs())
                                        {
                                            kword_list->setRhs(new SgExprListExp());
                                            kword_list = kword_list->rhs();
                                        }

                                        kword_list->setLhs(new SgExpression(ALLOCATABLE_OP));
                                    }

                                    // can't switch back to file with array usage
                                    if (SgFile::switchToFile(current_file_name) == -1)
                                        printInternalError(convertFileName(__FILE__).c_str(), __LINE__); 
                                        
                                    // insert allocate(a_l(lbound(a, 1):ubound(a,1),...)) statement
                                    SgStatement* allocate_stmt_place = NULL;

                                    auto* func_stmt = curr_stmt->getScopeForDeclare();
                                    for (auto iterator = func_stmt->lexNext();
                                        iterator && !isSgExecutableStatement(iterator);
                                        iterator = iterator->lexNext())
                                    {
                                            allocate_stmt_place = iterator;
                                    }

                                    //NULL - no decl stats in function!
                                    if (!allocate_stmt_place)
                                        allocate_stmt_place = func_stmt;

                                    auto *allocate_stmt_parent = allocate_stmt_place->controlParent();
                                    if (allocate_stmt_parent->variant() == GLOBAL)
                                        allocate_stmt_parent = allocate_stmt_place;
                                    
                                    auto *created_array_ref = new SgArrayRefExp(*copied_symbol.second);

                                    auto* dim_list = new SgExprListExp();
                                    created_array_ref->setLhs(dim_list);
                                    
                                    int dim_len = array_p->GetSizes().size();
                                    for (int i = 1; i <= dim_len; i++)
                                    {
                                        auto *lcall = new SgFunctionCallExp(*lbound_symb);
                                        auto *rcall = new SgFunctionCallExp(*ubound_symb);
                                        
                                        for (auto *call : {lcall, rcall})
                                        {
                                            call->setLhs(new SgExprListExp());

                                            call->lhs()->setLhs(new SgArrayRefExp(*array_to_copy));
                                            call->lhs()->setRhs(new SgValueExp(i));
                                        }
                                        
                                        auto *dot_expr = new SgExpression(DDOT);
                                        dot_expr->setLhs(lcall);
                                        dot_expr->setRhs(rcall);

                                        dim_list->setLhs(dot_expr);

                                        if (i < dim_len)
                                        {
                                            auto *next = new SgExprListExp();
                                            dim_list->setRhs(next);
                                            dim_list = next;
                                        }
                                    }

                                    auto *allocate_stmt = new SgStatement(ALLOCATE_STMT);
                                    allocate_stmt->setExpression(0, created_array_ref);

                                    allocate_stmt_place->insertStmtAfter(*allocate_stmt, *allocate_stmt_parent);

                                    // insert deallocate statemens before all returns
                                    auto *find_return_stmt = func_stmt;
                                    while (find_return_stmt != func_stmt->lastNodeOfStmt())
                                    {
                                        auto *next = find_return_stmt->lexNext();
                                        if (next && (isSgReturnStmt(next) || next == func_stmt->lastNodeOfStmt()))
                                        {
                                            if (next->hasLabel())
                                            {
                                                moveLabelBefore(next);
                                                find_return_stmt = next->lexPrev();
                                            }
                                            
                                            auto *dealloc_stmt = new SgStatement(DEALLOCATE_STMT);

                                            dealloc_stmt->setExpression(0, new SgExprListExp());
                                            dealloc_stmt->expr(0)->setLhs(new SgArrayRefExp(*copied_symbol.second));

                                            find_return_stmt->insertStmtAfter(*dealloc_stmt, *next->controlParent());

                                            if (next == curr_stmt)
                                                curr_stmt = dealloc_stmt;
                                        }

                                        find_return_stmt = next;
                                    }
                                }

                                // can't switch back to file with array usage
                                if (SgFile::switchToFile(current_file_name) == -1)
                                    printInternalError(convertFileName(__FILE__).c_str(), __LINE__); 
                            }
                        }

                        if (last_io_bound)
                        {
                            __spf_print(DEBUG_TRACE, "[io region] [%d, %d]\n", last_io_bound->lineNumber(), curr_stmt->lineNumber());

                            for (const auto& p : need_replace)
                            {
                                auto it = created_copies.find(p.first);
                                if (it != created_copies.end())
                                    replaceArrayInFragment(p.first, p.second, it->second, last_io_bound, curr_stmt, current_func_info);
                                else
                                    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
                            }
                        }

                        need_replace.clear();
                        last_io_bound = curr_stmt;
                    }

                    __spf_print(DEBUG_TRACE, "[line] %d (%d)\n", curr_stmt->lineNumber(), curr_stmt->variant());

                    for (int i = 0; i < 3; i++)
                    {
                        if (curr_stmt->expr(i))
                        {
                            __spf_print(DEBUG_TRACE, "[%d: %s (%d)]\n", i, curr_stmt->expr(i)->unparse(), curr_stmt->expr(i)->variant());
                        }
                    }

                    auto need_fix_io = populateDistributedIoArrays(need_replace, curr_stmt, current_file_name, current_func_info);
                    
                    // incorrect IO statement with label
                    // move label to dummy statement and insert copy statements between dummy statement and IO
                    if (need_fix_io && curr_stmt->hasLabel())
                    {
                        moveLabelBefore(curr_stmt);
                        if (last_io_bound == curr_stmt) // always true
                            last_io_bound = curr_stmt->lexPrev();
                    }

                    curr_stmt = curr_stmt->lexNext();
                }
            }
        }
    }
}
