#include "BrMerge.hpp"

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

#include <iostream>
using std::cout;
using std::endl;

void
BrMerge::run() {
    m_->set_print_name();
    BranchInst *br;
    ReturnInst *ret;
    for (auto &func : m_->get_functions()) {
        if (func.is_declaration())
            continue;
        bool cont = true;
        while (cont) {
            cont = false;
            for (auto &bb : func.get_basic_blocks()) {
                cout << bb.get_name() << endl;
                if (&bb == func.get_entry_block())
                    continue;
                auto &instructions = bb.get_instructions();
                br = dynamic_cast<BranchInst *>(&*instructions.rbegin());
                ret = dynamic_cast<ReturnInst *>(&*instructions.rbegin());
                assert((br or ret) && "final Instruction");
                if (instructions.size() == 1) {
                    if (br) { // change: br or cond-br with constant jump
                        BasicBlock *succ;
                        if (not br->is_cond_br()) {
                            assert(bb.get_succ_basic_blocks().size() == 1);
                            succ = *bb.get_succ_basic_blocks().begin();
                        } else if (dynamic_cast<Constant *>(
                                       br->get_operand(0))) {
                            assert(bb.get_succ_basic_blocks().size() == 2);
                            auto const_bool =
                                dynamic_cast<ConstantInt *>(br->get_operand(0))
                                    ->get_value();
                            succ = static_cast<BasicBlock *>(
                                const_bool ? br->get_operand(1)
                                           : br->get_operand(2));
                            static_cast<BasicBlock *>(
                                (const_bool ? br->get_operand(2)
                                            : br->get_operand(1)))
                                ->remove_pre_basic_block(&bb);
                        } else
                            continue;
                        if (bb.get_pre_basic_blocks().size() != 1)
                            continue;
                        auto pre = *bb.get_pre_basic_blocks().begin();
                        // change br's op
                        auto ins = &*pre->get_instructions().rbegin();
                        bool set = false;
                        for (int i = 0; i < ins->get_num_operand(); ++i)
                            if (ins->get_operand(i) == &bb) {
                                ins->set_operand(i, succ);
                                set = true;
                                bb.remove_use(ins);
                                break;
                            }
                        assert(set);
                        // change pre's succ
                        pre->remove_succ_basic_block(&bb);
                        pre->add_succ_basic_block(&bb);
                        // change succ's pre
                        succ->add_pre_basic_block(pre);
                        // change succ's pre
                        succ->remove_pre_basic_block(&bb);
                        // remove useless block
                        func.get_basic_blocks().remove(&bb);
                        // replace use
                        bb.replace_all_use_with(pre);
                        cont = true;
                        break;
                    } else { // ret: do not change
                    }
                } else {
                }
            }
        }
    }
}