package midend.optimize;

import java.util.HashSet;
import java.util.ArrayList;

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();
            }
        }
        deleteUselessInstrs();
        deleteUnreachedBlocks();
        makeCfg();
        for (IrFuncValue func : getIrModule().getFuncs()) {
            func.deleteDeadBlock();
        }
        makeDomination();
        makeDirectDomi();
        makeDomiFrontier();
    }

    public void deleteUselessInstrs() {
        for (IrFuncValue func : getIrModule().getFuncs()) {
            for (IrBasicBlock bb : func.getBBlocks()) {
                ArrayList<IrInstr> toDelete = new ArrayList<>();
                boolean deleteUseless = false;
                for (IrInstr instr : bb.getInstrs()) {
                    if (!deleteUseless) {
                        if (instr.getInstrType() == IrInstrType.RET ||
                                instr.getInstrType() == IrInstrType.BR ||
                                instr.getInstrType() == IrInstrType.JUMP) {
                            deleteUseless = true;
                        }
                    } else {
                        toDelete.add(instr);
                    }
                }
                for (IrInstr instr : toDelete) {
                    bb.deleteInstr(instr);
                }
            }
        }
    }

    public void deleteUnreachedBlocks() {
        for (IrFuncValue func : getIrModule().getFuncs()) {
            HashSet<IrBasicBlock> canArrive = new HashSet<>();
            IrBasicBlock entryBlock = func.getBBlock(0);
            dfsArrive(entryBlock, canArrive);
            ArrayList<IrBasicBlock> toDelete = new ArrayList<>();
            for (IrBasicBlock bb : func.getBBlocks()) {
                if (!canArrive.contains(bb)) {
                    toDelete.add(bb);
                }
            }
            for (IrBasicBlock bb : toDelete) {
                func.deleteBBlock(bb);
            }
        }
    }

    public void dfsArrive(IrBasicBlock bb, HashSet<IrBasicBlock> canArrive) {
        if (!canArrive.contains(bb)) {
            canArrive.add(bb);
        }
        IrInstr lastInstr = bb.getLastInstr();
        if (lastInstr != null) {
            if (lastInstr.getInstrType() == IrInstrType.BR) {
                BranchInstr branchInstr = (BranchInstr) lastInstr;
                IrBasicBlock trueBlock = branchInstr.getTrueBB();
                IrBasicBlock falseBlock = branchInstr.getFalseBB();
                if (!canArrive.contains(trueBlock)) {
                    dfsArrive(trueBlock, canArrive);
                }
                if (!canArrive.contains(falseBlock)) {
                    dfsArrive(falseBlock, canArrive);
                }
            } else if (lastInstr.getInstrType() == IrInstrType.JUMP) {
                JumpInstr jumpInstr = (JumpInstr) lastInstr;
                IrBasicBlock jumpBlock = jumpInstr.getTargetBlock();
                if (!canArrive.contains(jumpBlock)) {
                    dfsArrive(jumpBlock, canArrive);
                }
            }
        }
    }

    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 = new HashSet<>(bb.getDomied());
                    HashSet<IrBasicBlock> domiDomi = new HashSet<>(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;
                        }
                    }
                }
            }
        }
    }
}