#include "liverange.hpp"

#include "Function.h"
#include "Instruction.h"

#include <algorithm>
#include <iterator>

void
LiveRangeAnalyzer::clear() {
    IN.clear();
    OUT.clear();
    liverange.clear();
    instr_id.clear();
}

LiveRangeAnalyzer::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))]);
        // std::cout << "# " + it->print() << std::endl;
    }
    // std::cout << "\tget out: " << print_liveSet(out) << std::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
    int ir_cnt = 0;
    for (auto &bb : func->get_basic_blocks()) {
        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::run(Function *func) {
    clear();
    make_id(func);
#ifdef __LRA_PRINT__
    print(func, false);
#endif
    bool cont = true;
    while (cont) {
        cont = false;
        // reverse traverse BasicBlocks
        for (auto rit_bb = func->get_basic_blocks().rbegin();
             rit_bb != func->get_basic_blocks().rend();
             ++rit_bb) {
            auto bb = &(*rit_bb);

            LiveSet bef_in, out;
            bool last_ir = true;

            // reverse traverse instructions
            for (auto rit_ir = rit_bb->get_instructions().rbegin();
                 rit_ir != rit_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;
                        }
                    }
                }
            }
        }
    }
#ifdef __LRA_PRINT__
    print(func);
#endif
}

LiveRangeAnalyzer::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);
        if (ins) {
            assert(not ins->is_void() && "instr as op should not be void");
            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) { // for debug
    for (auto &bb : func->get_basic_blocks()) {
        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);
                    std::cout
                        << cpstmt_id[pr] << ". " << lv->get_name() << " = "
                        << (rv->get_name() == "" ? rv->print() : rv->get_name())
                        << std::endl;
                    if (not printSet)
                        continue;
                    auto &in = IN.at(idx);
                    auto &out = OUT.at(idx);
                    std::cout << "\tin-set: " + print_liveSet(in) << "\n";
                    std::cout << "\tout-set: " + print_liveSet(out) << "\n";
                }
            }
            // normal ir
            std::cout << instr_id[&instr] << ". " << instr.print() << " # "
                      << &instr << std::endl;
            if (not printSet)
                continue;
            auto idx = instr_id.at(&instr);
            auto &in = IN.at(idx);
            auto &out = OUT.at(idx);
            std::cout << "\tin-set: " + print_liveSet(in) << "\n";
            std::cout << "\tout-set: " + print_liveSet(out) << "\n";
        }
        /* if (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);
         *         std::cout << cpstmt_id[pr] << ". " << lv->get_name() << " = "
         *                   << (rv->get_name() == "" ? rv->print()
         *                                            : rv->get_name())
         *                   << std::endl;
         *         if (not printSet)
         *             continue;
         *         auto &in = IN.at(idx);
         *         auto &out = OUT.at(idx);
         *         std::cout << "\tin-set: " + print_liveSet(in) << "\n";
         *         std::cout << "\tout-set: " + print_liveSet(out) << "\n";
         *     }
         * } */
    }
}