#include "BrMerge.hpp"

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

void
BrMerge::run() {
    BranchInst *br;
    ReturnInst *ret;
    for (auto &func : m_->get_functions()) {
        bool cont = true;
        while (cont) {
            cont = false;
            for (auto &bb : func.get_basic_blocks()) {
                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);
                            succ = static_cast<BasicBlock *>(
                                dynamic_cast<ConstantInt *>(br->get_operand(0))
                                        ->get_value()
                                    ? br->get_operand(1)
                                    : br->get_operand(2));
                        } else
                            continue;
                        for (auto pre : bb.get_pre_basic_blocks()) {
                            // 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;
                                    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);
                        cont = true;
                    } else { // ret: do not change
                    }
                } else {
                }
            }
        }
    }
}