package midend.optimize;

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

import midend.llvm.IrBuilder;
import midend.llvm.constant.IrConstantInt;
import midend.llvm.instr.AllocateInstr;
import midend.llvm.instr.IrInstr;
import midend.llvm.instr.StoreInstr;
import midend.llvm.instr.LoadInstr;
import midend.llvm.instr.PhiInstr;
import midend.llvm.use.IrUser;
import midend.llvm.value.IrBasicBlock;
import midend.llvm.value.IrValue;

public class PhiInsert {
    private ArrayList<IrInstr> defineInstrs;
    private ArrayList<IrInstr> useInstrs;
    private AllocateInstr allocInstr;
    private IrBasicBlock entryBlock;
    private ArrayList<IrBasicBlock> defBlocks;
    private ArrayList<IrBasicBlock> useBlocks;
    private Stack<IrValue> workList;

    public PhiInsert(AllocateInstr allocInstr, IrBasicBlock entryBlock) {
        this.defineInstrs = new ArrayList<>();
        this.useInstrs = new ArrayList<>();
        this.allocInstr = allocInstr;
        this.entryBlock = entryBlock;
        this.defBlocks = new ArrayList<>();
        this.useBlocks = new ArrayList<>();
        this.workList = new Stack<>();
    }

    public IrValue getNewValue() {
        if (workList.isEmpty()) {
            return new IrConstantInt(0);
        }
        return workList.peek();
    }

    public void run() {
        makeDefAndUse();
        insertPhiNodes();
        changeLoadAndStoreToReg(entryBlock);
    }

    public void makeDefAndUse() {
        ArrayList<IrUser> users = new ArrayList<>(allocInstr.getUsers());
        for (IrUser user : users) {
            IrInstr instr = (IrInstr) user;
            if (instr instanceof StoreInstr) {
                defineInstrs.add(instr);
                IrBasicBlock defBlock = instr.getBBlock();
                if (!defBlocks.contains(defBlock)) {
                    defBlocks.add(defBlock);
                }
            } else if (instr instanceof LoadInstr){
                useInstrs.add(instr);
                IrBasicBlock useBlock = instr.getBBlock();
                if (!useBlocks.contains(useBlock)) {
                    useBlocks.add(useBlock);
                }
            }
        }
    }

    public void insertPhiNodes() {
        HashSet<IrBasicBlock> hasPhiBlocks = new HashSet<>();
        Stack<IrBasicBlock> defStack = new Stack<>();
        for (IrBasicBlock defBlock : defBlocks) {
            defStack.push(defBlock);
        }
        while (!defStack.isEmpty()) {
            IrBasicBlock currBlock = defStack.pop();
            for (IrBasicBlock frontierBlock : currBlock.getDomiFrontier()) {
                if (hasPhiBlocks.contains(frontierBlock)) {
                    continue;
                }
                //插入phi节点
                insertPhiNode(frontierBlock);
                hasPhiBlocks.add(frontierBlock);
                if (!defBlocks.contains(frontierBlock)) {
                    defStack.push(frontierBlock);
                }
            }
        }
    }

    public void insertPhiNode(IrBasicBlock block) {
        PhiInstr phiInstr = new PhiInstr(allocInstr.getPointeeType(), block, IrBuilder.getLocalName(block.getFunc()));
        block.addInstr(phiInstr, 0);
        defineInstrs.add(phiInstr);
        useInstrs.add(phiInstr);
    }

    public void changeLoadAndStoreToReg(IrBasicBlock block) {
        Stack<IrValue> workListCopy = new Stack<>();
        workListCopy.addAll(workList);
        ArrayList<IrInstr> deleteInstrs = new ArrayList<>();
        for(IrInstr instr : block.getInstrs()) {
            if (instr instanceof LoadInstr && useInstrs.contains(instr)) {
                instr.replaceUserToAnother(getNewValue());
                deleteInstrs.add(instr);
            } else if (instr instanceof StoreInstr && defineInstrs.contains(instr)) {
                StoreInstr storeInstr = (StoreInstr) instr;
                IrValue value = storeInstr.getValue();
                workList.push(value);
                deleteInstrs.add(instr);
            } else if (instr instanceof PhiInstr && defineInstrs.contains(instr)) {
                workList.push(instr);
            } else if (instr == allocInstr) {
                deleteInstrs.add(instr);
            }
        }
        for (IrInstr instr : deleteInstrs) {
            block.deleteInstr(instr);
        }
        for (IrBasicBlock succ : block.getSuccs()) {
            IrInstr firstInstr = succ.getFirstInstr();
            if (firstInstr instanceof PhiInstr && useInstrs.contains(firstInstr)) {
                PhiInstr phiInstr = (PhiInstr) firstInstr;
                phiInstr.setValueForPred(block, getNewValue());
            }
        }
        for (IrBasicBlock domiBB : block.getDirectDomies()) {
            changeLoadAndStoreToReg(domiBB);
        }
        workList = workListCopy;
    }
}