package midend.optimize;

import java.util.HashSet;

import midend.llvm.instr.IrInstr;
import midend.llvm.instr.IrInstrType;
import midend.llvm.instr.JumpInstr;
import midend.llvm.value.IrBasicBlock;
import midend.llvm.value.IrFuncValue;
import midend.llvm.instr.BranchInstr;

public class CfgMake extends Optimizer {
    public void optimize() {
        for (IrFuncValue func : getIrModule().getFuncs()) {
            for (IrBasicBlock bb : func.getBBlocks()) {
                bb.clearCfg();
            }
        }
        makeCfg();
    }

    public void makeCfg() {
        for (IrFuncValue func : getIrModule().getFuncs()) {
            for (IrBasicBlock bb : func.getBBlocks()) {
                for (IrInstr instr : bb.getInstrs()) {
                    if (instr.getInstrType() == IrInstrType.BR) {
                        BranchInstr branchInstr = (BranchInstr) instr;
                        IrBasicBlock trueBlock = branchInstr.getTrueBB();
                        IrBasicBlock falseBlock = branchInstr.getFalseBB();
                        bb.addSucc(trueBlock);
                        bb.addSucc(falseBlock);
                        trueBlock.addPred(bb);
                        falseBlock.addPred(bb);
                    } else if (instr.getInstrType() == IrInstrType.JUMP) {
                        JumpInstr jumpInstr = (JumpInstr) instr;
                        IrBasicBlock jumpBlock = jumpInstr.getTargetBlock();
                        bb.addSucc(jumpBlock);
                        jumpBlock.addPred(bb);
                    }
                }
            }
        }
    }

    public void makeDomination() {
        for (IrFuncValue func : getIrModule().getFuncs()) {
            for (IrBasicBlock bb : func.getBBlocks()) {
                HashSet<IrBasicBlock> canArrive = new HashSet<>();
                handleArrive(func.getBBlock(0), bb, canArrive);
                for (IrBasicBlock block : func.getBBlocks()) {
                    if (!canArrive.contains(block)) {
                        block.addDomied(bb);
                    }
                }
            }
        }
    }

    public void handleArrive(IrBasicBlock entryBlock,
            IrBasicBlock deleteBlock, HashSet<IrBasicBlock> canArrive) {
        if (entryBlock == deleteBlock) {
            return;
        } else {
            if (!canArrive.contains(entryBlock)) {
                canArrive.add(entryBlock);
            }
            for (IrBasicBlock succ : entryBlock.getSuccs()) {
                if (succ != deleteBlock && !canArrive.contains(succ)) {
                    handleArrive(succ, deleteBlock, canArrive);
                }
            }
        }
    }

    public void makeDirectDomi() {
        for (IrFuncValue func : getIrModule().getFuncs()) {
            for (IrBasicBlock bb : func.getBBlocks()) {
                for (IrBasicBlock domi : bb.getDomied()) {
                    HashSet<IrBasicBlock> bbDomi = bb.getDomied();
                    HashSet<IrBasicBlock> domiDomi = domi.getDomied();
                    for (IrBasicBlock domiDomiBB : domiDomi) {
                        if (bbDomi.contains(domiDomiBB)) {
                            bbDomi.remove(domiDomiBB);
                        }
                    }
                    if (bbDomi.size() == 1 && bbDomi.contains(bb)) {
                        bb.setDirectDomi(domi);
                        domi.addDirectDomies(bb);
                        break;
                    }
                }
            }
        }
    }

    public void makeDomiFrontier() {
        for (IrFuncValue func : getIrModule().getFuncs()) {
            for (IrBasicBlock bb : func.getBBlocks()) {
                for (IrBasicBlock succ : bb.getSuccs()) {
                    IrBasicBlock currentBB = bb;
                    while (!succ.getDomied().contains(currentBB) || currentBB == succ) {
                        currentBB.addDomiFrontier(succ);
                        currentBB = currentBB.getDirectDomi();
                        if (currentBB == null) {
                            break;
                        }
                    }
                }
            }
        }
    }
}