#include "LoopUnroll.hpp"

#include "BasicBlock.h"
#include "Constant.h"
#include "Function.h"
#include "Instruction.h"
#include "syntax_analyzer.h"

#include <vector>

using std::find;

using namespace Graph;
using namespace Analysis;

void
LoopUnroll::run() {
    for (auto &_f : m_->get_functions()) {
        if (_f.is_declaration())
            continue;
        auto func = &_f;
        // cout << func->get_name() << endl;
        auto belist = detect_back(func);
        auto sloops = check_sloops(belist);
        cout << "get simple loops for function " << func->get_name() << ":\n";
        for (auto sl : sloops) {
            cout << "\t";
            for (auto p : sl)
                cout << p->get_name() << " ";
            cout << "\n";
        }
    }
}

LoopAnalysis::LoopAnalysis(const Graph::SimpleLoop &sl) {
    auto b = sl.front();
    auto e = sl.back();

    // In `p`, get stop number(const)
    Value *i;
    auto rit = b->get_instructions().rbegin();
    assert(dynamic_cast<BranchInst *>(&*rit) &&
           "The end instruction of a block should be branch");
    i = (rit++)->get_operand(0);
    assert(i == &*rit && dynamic_cast<CmpInst *>(&*rit) &&
           static_cast<CmpInst *>(&*rit)->get_cmp_op() == CmpInst::NE &&
           "should be neq 0");
    i = (rit++)->get_operand(0);
    assert(i == &*rit && dynamic_cast<ZextInst *>(&*rit) && "neqz");
    i = (rit++)->get_operand(0);
    assert(
        i == &*rit &&
        (dynamic_cast<CmpInst *>(&*rit) or dynamic_cast<FCmpInst *>(&*rit)) &&
        "cmp or fcmp");
    if (dynamic_cast<Constant *>(rit->get_operand(0)) or
        dynamic_cast<Constant *>(rit->get_operand(1))) {
        if (dynamic_cast<CmpInst *>(&*rit)) {
            Type = FLOAT;
            auto constfloat = dynamic_cast<ConstantFP *>(rit->get_operand(1));
            assert(constfloat &&
                   "the case const at operand(0) not implemented");
            threshold.fv = constfloat->get_value();
        } else {
            Type = INT;
            auto constint = dynamic_cast<ConstantInt *>(rit->get_operand(1));
            assert(constint && "the case const at operand(0) not implemented");
            threshold.v = constint->get_value();
        }
    } else {
        Type = UNDEF;
        return;
    }

    // get control value and initial value
    auto control = rit->get_operand(0);
    auto it = b->get_instructions().begin();
    for (; it != b->get_instructions().end(); ++it) {
        auto phi = dynamic_cast<PhiInst *>(&*it);
        assert(phi && "unexpected structure for while block");
        if (&*it != control)
            continue;
        assert(phi->get_operand(3) == e && "the orther case not implemented");
        if (dynamic_cast<Constant *>(phi->get_operand(0))) {
            switch (Type) {
                case INT:
                    initial.v = static_cast<ConstantInt *>(phi->get_operand(0))
                                    ->get_value();
                    break;
                case FLOAT:
                    initial.fv = static_cast<ConstantFP *>(phi->get_operand(0))
                                     ->get_value();
                    break;
                case UNDEF:
                    assert(false);
                    break;
            }
        } else {
            Type = UNDEF;
            return;
        }
    }

    // get delta

    // check correctness

    // count loop 
}

vector<SimpleLoop>
LoopUnroll::check_sloops(const BackEdgeList &belist) const {
    vector<SimpleLoop> sloops;
    for (auto [e, b] : belist) {
        SimpleLoop sl;
        if (b->get_succ_basic_blocks().size() != 2 or
            b->get_pre_basic_blocks().size() != 2)
            continue;
        bool flag = true;
        // the start node should have 2*2 degree, and others have 1*1 degree
        // one exception: br to neg_idx_except
        for (auto p = e; p != b; p = *p->get_pre_basic_blocks().begin()) {
            if (p->get_pre_basic_blocks().size() != 1) {
                flag = false;
                break;
            }
            auto succ_bbs = p->get_succ_basic_blocks();
            if (succ_bbs.size() != 1) {
                assert(succ_bbs.size() == 2);
                flag = false;
                for (auto bb : succ_bbs) {
                    auto instr = &*bb->get_instructions().begin();
                    if (instr->is_call() and
                        instr->get_operand(0) == neg_func) {
                        flag = true;
                        break;
                    }
                }
            }
            if (not flag)
                break;
            sl.insert(sl.begin(), p);
        }

        if (not flag)
            continue;
        sl.insert(sl.begin(), b);
        if (flag)
            sloops.push_back(sl);
    }
    return sloops;
}

BackEdgeList
LoopUnroll::detect_back(Function *func) {
    BackEdgeSearcher search(func->get_entry_block());
    return search.edges;
}

void
BackEdgeSearcher::dfsrun(BasicBlock *bb) {
    vis[bb] = true;
    path.push_back(bb);
    for (auto succ : bb->get_succ_basic_blocks()) {
        if (vis[succ]) {
            string type;
            Edge edge(bb, succ);
            if (find(path.rbegin(), path.rend(), succ) == path.rend()) {
                type = "cross-edge";
            } else {
                type = "back-edge";
                edges.push_back(edge);
            }
            /* cout << "find " << type << ": " << LoopUnroll::str(edge)
             *      << "\n"; */
        } else
            dfsrun(succ);
    }

    path.pop_back();
}