#include "liverange.hpp"

#include "BasicBlock.h"
#include "Function.h"

#include <deque>

using std::cout;
using std::endl;
using namespace LRA;

void
LiveRangeAnalyzer::clear() {
    ir_cnt = 0;
    IN.clear();
    OUT.clear();
    instr_id.clear();
    cpstmt_id.clear();
    intervalmap.clear();
    liveIntervals.clear();
    BB_DFS_Order.clear();
}

LiveSet
LiveRangeAnalyzer::joinFor(BasicBlock *bb) {
    LiveSet out;
    for (auto succ : bb->get_succ_basic_blocks()) {
        auto &irs = succ->get_instructions();
        auto it = irs.begin();
        while (it != irs.end() and it->is_phi())
            ++it;
        assert(it != irs.end() && "need to find first_ir from copy-stmt");
        union_ip(out, IN[instr_id.at(&(*it))]);
        // cout << "# " + it->print() << endl;
    }
    // cout << "\tget out: " << print_liveSet(out) << endl;
    return out;
}

void
LiveRangeAnalyzer::make_id(Function *func) {
    // instruction numbering
    // this is also the structure of the IR logically:
    // ignore phi, add copy-statement
#ifdef __USE_DFS_ORDER__
    for (auto bb : BB_DFS_Order) {
#else
    for (auto &bb_ : func->get_basic_blocks()) {
        auto bb = &bb_;
#endif
        for (auto &instr : bb->get_instructions()) {
            if (instr.is_phi())
                continue;
            if (instr.is_br() or instr.is_ret()) {
                // insert copy-stmt in front of the jump ir
                auto it = phi_map.find(bb);
                if (it != phi_map.end()) {
                    for (auto pr : it->second) {
                        cpstmt_id[pr] = ++ir_cnt;
                    }
                }
            }
            instr_id[&instr] = ++ir_cnt;
        }
    }
}

void
LiveRangeAnalyzer::get_dfs_order(Function *func) {
    map<BasicBlock *, bool> vis;
    std::deque<BasicBlock *> Q;

    Q.push_back(func->get_entry_block());

    while (not Q.empty()) {
        auto bb = Q.front();
        Q.pop_front();
        if (vis[bb])
            continue;
        vis[bb] = true;
        BB_DFS_Order.push_back(bb);
        for (auto succ : bb->get_succ_basic_blocks())
            Q.push_front(succ);
    }
    cout << "DFS order for function " << func->get_name() << ":\n";
    for (auto bb : BB_DFS_Order)
        cout << bb->get_name() << " ";
    cout << endl;
}

void
LiveRangeAnalyzer::run(Function *func) {
    clear();
    get_dfs_order(func);
    make_id(func);
    bool cont = true;
    while (cont) {
        cont = false;
        // reverse traverse BasicBlocks
#ifdef __USE_DFS_ORDER__
        for (auto rit_bb = BB_DFS_Order.rbegin(); rit_bb != BB_DFS_Order.rend();
             ++rit_bb) {
            auto bb = *rit_bb;
#else
        for (auto rit_bb = func->get_basic_blocks().rbegin();
             rit_bb != func->get_basic_blocks().rend();
             ++rit_bb) {
            auto bb = &(*rit_bb);
#endif

            LiveSet bef_in, out;
            bool last_ir = true;

            // reverse traverse instructions
            for (auto rit_ir = bb->get_instructions().rbegin();
                 rit_ir != bb->get_instructions().rend();
                 ++rit_ir) {
                auto instr = &(*rit_ir);
                if (instr->is_phi()) {
                    assert(not last_ir && "If phi is the last ir, then data "
                                          "flow fails due to ignorance of phi");
                    continue;
                }
                //
                // get out-set for this instruction
                if (last_ir) {
                    last_ir = false;
                    out = joinFor(bb);
                } else {
                    out = bef_in;
                }
                OUT[instr_id.at(instr)] = out;
                //
                // get in-set
                bef_in = transferFunction(instr);
                cont |= bef_in != IN[instr_id.at(instr)];
                IN[instr_id.at(instr)] = bef_in;

                if (instr->is_ret() or instr->is_br()) {
                    // deal with copy-stmt
                    auto it = phi_map.find(bb);
                    if (it != phi_map.end()) {
                        for (auto rit = it->second.rbegin();
                             rit != it->second.rend();
                             ++rit) {
                            auto [lv, rv] = *rit;
                            if (last_ir) {
                                last_ir = false;
                                out = joinFor(bb);
                            } else
                                out = bef_in;

                            OUT[cpstmt_id.at(*rit)] = out;
                            // transfer manually
                            out.erase(lv);
                            if (dynamic_cast<Instruction *>(rv))
                                out.insert(rv);
                            cont |= out != IN[cpstmt_id.at(*rit)];
                            bef_in = IN[cpstmt_id.at(*rit)] = out;
                        }
                    }
                }
            }
        }
    }
    // argument should be in the IN-set of Entry
    assert(IN.find(0) == IN.end() and OUT.find(0) == OUT.end() &&
           "no instr_id will be mapped to 0");
    IN[0] = OUT[0] = {};
    for (auto arg : func->get_args())
        IN[0].insert(arg);
    make_interval(func);

#ifdef __LRA_PRINT__
    print(func, true, true);
#endif
}

void
LiveRangeAnalyzer::make_interval(Function *) {
    for (int time = 0; time <= ir_cnt; ++time) {
        for (auto op : IN.at(time)) {
            auto &interval = intervalmap[op];
            if (interval.i == UNINITIAL) // uninitialized
                interval.i = interval.j = time;
            else
                interval.j = time;
        }
        for (auto op : OUT.at(time)) {
            auto &interval = intervalmap[op];
            if (interval.i == UNINITIAL) // uninitialized
                interval.i = interval.j = time + 1;
            else
                interval.j = time + 1;
        }
    }
    for (auto &[op, interval] : intervalmap)
        liveIntervals.insert({interval, op});
}

LiveSet
LiveRangeAnalyzer::transferFunction(Instruction *instr) {
    LiveSet in, out = OUT[instr_id.at(instr)];
    LiveSet use;

    // calculate use for this instruction
    for (auto op : instr->get_operands()) {
        // op type:
        // - global
        // - function
        // - Constant
        // - instruction
        // - argument*
        // - BasicBlock
        auto ins = dynamic_cast<Instruction *>(op);
        auto arg = dynamic_cast<Argument *>(op);
        if (ins) {
            assert(not ins->is_void() && "instr as op should not be void");
            use.insert(op);
        } else if (arg) {
            use.insert(op);
        }
    }
    // in = use + (out - def)
    auto iter = out.find(instr);
    if (iter != out.end())
        out.erase(iter);
    std::set_union(out.begin(),
                   out.end(),
                   use.begin(),
                   use.end(),
                   std::inserter(in, in.begin()));

    return in;
}

void
LiveRangeAnalyzer::print(Function *func,
                         bool printSet,
                         bool printInt) const { // for debug
    cout << "Function " << func->get_name() << endl;
    cout << "0. Entry" << endl;
    if (printSet) {
        cout << "\tin-set: " + print_liveSet(IN.at(0)) << "\n";
        cout << "\tout-set: " + print_liveSet(OUT.at(0)) << "\n";
    }
#ifdef __USE_DFS_ORDER__
    for (auto bb : BB_DFS_Order) {
#else
    for (auto &bb_ : func->get_basic_blocks()) {
        auto bb = &bb_;
#endif
        for (auto &instr : bb->get_instructions()) {
            if (instr.is_phi()) // ignore phi
                continue;
            // insert copy-stmt
            if ((instr.is_br() or instr.is_ret()) and
                phi_map.find(bb) != phi_map.end()) {
                for (auto pr : phi_map.find(bb)->second) {
                    auto [lv, rv] = pr;
                    auto idx = cpstmt_id.at(pr);
                    cout << cpstmt_id.at(pr) << ". " << lv->get_name() << " = "
                         << (rv->get_name() == "" ? rv->print()
                                                  : rv->get_name())
                         << endl;
                    if (not printSet)
                        continue;
                    auto &in = IN.at(idx);
                    auto &out = OUT.at(idx);
                    cout << "\tin-set: " + print_liveSet(in) << "\n";
                    cout << "\tout-set: " + print_liveSet(out) << "\n";
                }
            }
            // normal ir
            cout << instr_id.at(&instr) << ". " << instr.print() << endl;
            if (not printSet)
                continue;
            auto idx = instr_id.at(&instr);
            auto &in = IN.at(idx);
            auto &out = OUT.at(idx);
            cout << "\tin-set: " + print_liveSet(in) << "\n";
            cout << "\tout-set: " + print_liveSet(out) << "\n";
        }
    }

    if (printInt) {
        for (auto [interval, op] : liveIntervals)
            cout << op->get_name() << ": " << print_interval(interval) << endl;
    }
}