/**
 * @file: Sysyoptimization.cpp
 * @brief  CFG优化
 * @Author : Ixeux email:you@domain.com
 * @Version : 1.0
 * @Creat Date : 2024-08-10
 *
 */
#include "SysYIROptPre.h"
#include <cassert>
#include <list>
#include <map>
#include <memory>
#include <string>
#include "IR.h"
#include "IRBuilder.h"

namespace sysy {

/**
 * use删除operand,以免扰乱后续分析
 * instr: 要删除的指令
 */
void SysYOptPre::usedelete(Instruction *instr) {
  for (auto &use : instr->getOperands()) {
    Value* val = use->getValue();
    // std::cout << delete << val->getName() << std::endl;
    val->removeUse(use);
  }
}


// 删除br后的无用指令
void SysYOptPre::SysYDelInstAfterBr() {
  auto &functions = pModule->getFunctions();
  for (auto &function : functions) {
    auto basicBlocks = function.second->getBasicBlocks();
    for (auto &basicBlock : basicBlocks) {
      bool Branch = false;
      auto &instructions = basicBlock->getInstructions();
      auto Branchiter = instructions.end();
      for (auto iter = instructions.begin(); iter != instructions.end(); ++iter) {
        if (Branch) 
          usedelete(iter->get());
        else if ((*iter)->isTerminator()){
          Branch = true;
          Branchiter = iter;
        }
      }
      if (Branchiter != instructions.end()) ++Branchiter;
      while (Branchiter != instructions.end()) 
        Branchiter = instructions.erase(Branchiter);
      
      if (Branch) {  // 更新前驱后继关系
        auto thelastinstinst = basicBlock->getInstructions().end();
        --thelastinstinst;
        auto &Successors = basicBlock->getSuccessors();
        for (auto iterSucc = Successors.begin(); iterSucc != Successors.end();) {
          (*iterSucc)->removePredecessor(basicBlock.get());
          basicBlock->removeSuccessor(*iterSucc);
        }
        if (thelastinstinst->get()->isUnconditional()) {
          BasicBlock* branchBlock = dynamic_cast<BasicBlock *>(thelastinstinst->get()->getOperand(0));
          basicBlock->addSuccessor(branchBlock);
          branchBlock->addPredecessor(basicBlock.get());
        } else if (thelastinstinst->get()->isConditional()) {
          BasicBlock* thenBlock = dynamic_cast<BasicBlock *>(thelastinstinst->get()->getOperand(1));
          BasicBlock* elseBlock = dynamic_cast<BasicBlock *>(thelastinstinst->get()->getOperand(2));
          basicBlock->addSuccessor(thenBlock);
          basicBlock->addSuccessor(elseBlock);
          thenBlock->addPredecessor(basicBlock.get());
          elseBlock->addPredecessor(basicBlock.get());
        }
      }
    }
  }
}


void SysYOptPre::SysYBlockMerge() {
  auto &functions = pModule->getFunctions(); //std::map<std::string, std::unique_ptr<Function>>
  for (auto &function : functions) {
    // auto basicBlocks = function.second->getBasicBlocks();
    auto &func = function.second;
    for (auto blockiter = func->getBasicBlocks().begin();
         blockiter != func->getBasicBlocks().end();) {
      if (blockiter->get()->getNumSuccessors() == 1) {
        // 如果当前块只有一个后继块
        // 且后继块只有一个前驱块
        // 则将当前块和后继块合并
        if (((blockiter->get())->getSuccessors()[0])->getNumPredecessors() == 1) {
          // std::cout << "merge block: " << blockiter->get()->getName() << std::endl;
          BasicBlock* block = blockiter->get();
          BasicBlock* nextBlock = blockiter->get()->getSuccessors()[0];
          auto nextarguments = nextBlock->getArguments();
          // 删除br指令
          if (block->getNumInstructions() != 0) {
            auto thelastinstinst = block->end();
            (--thelastinstinst);
            if (thelastinstinst->get()->isUnconditional()) {
              usedelete(thelastinstinst->get());
              block->getInstructions().erase(thelastinstinst);
            } else if (thelastinstinst->get()->isConditional()) {
              // 如果是条件分支，判断条件是否相同，主要优化相同布尔表达式
              if (thelastinstinst->get()->getOperand(1)->getName() == thelastinstinst->get()->getOperand(1)->getName()) {
                usedelete(thelastinstinst->get());
                block->getInstructions().erase(thelastinstinst);
              }
            }
          }
          // 将后继块的指令移动到当前块
          // 并将后继块的父指针改为当前块
          for (auto institer = nextBlock->begin(); institer != nextBlock->end();) {
            institer->get()->setParent(block);
            block->getInstructions().emplace_back(institer->release());
            institer = nextBlock->getInstructions().erase(institer);   
          }
          // 合并参数
          // TODO：是否需要去重?
          for (auto &argm : nextarguments) {
            argm->setParent(block);
            block->insertArgument(argm);
          }
          // 更新前驱后继关系，类似树节点操作
          block->removeSuccessor(nextBlock);
          nextBlock->removePredecessor(block);
          std::list<BasicBlock *> succshoulddel;
          for (auto &succ : nextBlock->getSuccessors()) {
            block->addSuccessor(succ);
            succ->replacePredecessor(nextBlock, block);
            succshoulddel.push_back(succ);
          }
          for (auto del : succshoulddel) {
            nextBlock->removeSuccessor(del);
          }

          func->removeBasicBlock(nextBlock);

        } else {
          blockiter++;
        }
      } else {
        blockiter++;
      }
    }
  }
}

// 删除无前驱块，兼容SSA后的处理
void SysYOptPre::SysYDelNoPreBLock() {
  
  auto &functions = pModule->getFunctions(); // std::map<std::string, std::unique_ptr<sysy::Function>>
  for (auto &function : functions) {
    auto &func = function.second;

    for (auto &block : func->getBasicBlocks()) {
      block->setreachableFalse();
    }
    // 对函数基本块做一个拓扑排序，排查不可达基本块
    auto entryBlock = func->getEntryBlock();
    entryBlock->setreachableTrue();
    std::queue<BasicBlock *> blockqueue;
    blockqueue.push(entryBlock);
    while (!blockqueue.empty()) {
      auto block = blockqueue.front();
      blockqueue.pop();
      for (auto &succ : block->getSuccessors()) {
        if (!succ->getreachable()) {
          succ->setreachableTrue();
          blockqueue.push(succ);
        }
      }
    }

    // 删除不可达基本块指令
    for (auto blockIter = func->getBasicBlocks().begin(); blockIter != func->getBasicBlocks().end();blockIter++) {
      
      if (!blockIter->get()->getreachable()) 
        for (auto &iterInst : blockIter->get()->getInstructions()) 
          usedelete(iterInst.get());
    
    }

    
    for (auto blockIter = func->getBasicBlocks().begin(); blockIter != func->getBasicBlocks().end();) {
      if (!blockIter->get()->getreachable()) {
        for (auto succblock : blockIter->get()->getSuccessors()) {
          int indexphi = 1;
          for (auto pred : succblock->getPredecessors()) {
            if (pred == blockIter->get()) {
              break;
            }
            indexphi++;
          }
          for (auto &phiinst : succblock->getInstructions()) {
            if (phiinst->getKind() != Instruction::kPhi) {
              break;
            }
            phiinst->removeOperand(indexphi);
          }
        }
        // 删除不可达基本块，注意迭代器不可达问题
        func->removeBasicBlock((blockIter++)->get());
      } else {
        blockIter++;
      }
    }
  }
}

void SysYOptPre::SysYDelEmptyBlock() {
  auto &functions = pModule->getFunctions();
  for (auto &function : functions) {
    // 收集不可达基本块
    // 这里的不可达基本块是指没有实际指令的基本块
    // 当一个基本块没有实际指令例如只有phi指令和一个uncondbr指令时，也会被视作不可达
    auto basicBlocks = function.second->getBasicBlocks();
    std::map<sysy::BasicBlock *, BasicBlock *> EmptyBlocks;
    // 空块儿和后继的基本块的映射
    for (auto &basicBlock : basicBlocks) {
      if (basicBlock->getNumInstructions() == 0) {
        if (basicBlock->getNumSuccessors() == 1) {
          EmptyBlocks[basicBlock.get()] = basicBlock->getSuccessors().front();
        }
      }
      else{
        // 如果只有phi指令和一个uncondbr。(phi)*(uncondbr)?
          // 判断除了最后一个指令之外是不是只有phi指令
        bool onlyPhi = true;
        for (auto &inst : basicBlock->getInstructions()) {
          if (!inst->isPhi() && !inst->isUnconditional()) {
            onlyPhi = false;
            break;
          }
        }
        if(onlyPhi)
          EmptyBlocks[basicBlock.get()] = basicBlock->getSuccessors().front();
      }
      
      
    }
    // 更新基本块信息，增加必要指令
    for (auto &basicBlock : basicBlocks) {
      // 把空块转换成只有跳转指令的不可达块
      if (distance(basicBlock->begin(), basicBlock->end()) == 0) {
        if (basicBlock->getNumSuccessors() == 0) {
          continue;
        }
        if (basicBlock->getNumSuccessors() > 1) {
          assert("");
        }
        pBuilder->setPosition(basicBlock.get(), basicBlock->end());
        pBuilder->createUncondBrInst(basicBlock->getSuccessors()[0], {});
        continue;
      }

      auto thelastinst = basicBlock->getInstructions().end();
      --thelastinst;

      // 根据br指令传递的后继块信息，跳过空块链
      if (thelastinst->get()->isUnconditional()) {
        BasicBlock* OldBrBlock = dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0));
        BasicBlock *thelastBlockOld = nullptr;
        // 如果空块链表为多个块
        while (EmptyBlocks.find(dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0))) !=
               EmptyBlocks.end()) {
          thelastBlockOld = dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0));
          thelastinst->get()->replaceOperand(0, EmptyBlocks[thelastBlockOld]);
        }

        basicBlock->removeSuccessor(OldBrBlock);
        OldBrBlock->removePredecessor(basicBlock.get());
        basicBlock->addSuccessor(dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0)));
        dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0))->addPredecessor(basicBlock.get());

        if (thelastBlockOld != nullptr) {
          int indexphi = 0;
          for (auto &pred : dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0))->getPredecessors()) {
            if (pred == thelastBlockOld) {
              break;
            }
            indexphi++;
          }

          // 更新phi指令的操作数
          // 移除thelastBlockOld对应的phi操作数
          for (auto &InstInNew : dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0))->getInstructions()) {
            if (InstInNew->isPhi()) {
              dynamic_cast<PhiInst *>(InstInNew.get())->removeOperand(indexphi + 1);
            } else {
              break;
            }
          }
        }

      } else if (thelastinst->get()->getKind() == Instruction::kCondBr) {
        auto OldThenBlock = dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1));
        auto OldElseBlock = dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(2));

        BasicBlock *thelastBlockOld = nullptr;
        while (EmptyBlocks.find(dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1))) !=
               EmptyBlocks.end()) {
          thelastBlockOld = dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1));
          thelastinst->get()->replaceOperand(
              1, EmptyBlocks[dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1))]);
        }
        basicBlock->removeSuccessor(OldThenBlock);
        OldThenBlock->removePredecessor(basicBlock.get());
        // 处理 then 和 else 分支合并的情况
        if (dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1)) ==
            dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(2))) {
          auto thebrBlock = dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1));
          usedelete(thelastinst->get());
          thelastinst = basicBlock->getInstructions().erase(thelastinst);
          pBuilder->setPosition(basicBlock.get(), basicBlock->end());
          pBuilder->createUncondBrInst(thebrBlock, {});
          continue;
        }
        basicBlock->addSuccessor(dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1)));
        dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1))->addPredecessor(basicBlock.get());
        // auto indexInNew = dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0))->getPredecessors().

        if (thelastBlockOld != nullptr) {
          int indexphi = 0;
          for (auto &pred : dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1))->getPredecessors()) {
            if (pred == thelastBlockOld) {
              break;
            }
            indexphi++;
          }

          for (auto &InstInNew : dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1))->getInstructions()) {
            if (InstInNew->isPhi()) {
              dynamic_cast<PhiInst *>(InstInNew.get())->removeOperand(indexphi + 1);
            } else {
              break;
            }
          }
        }

        thelastBlockOld = nullptr;
        while (EmptyBlocks.find(dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(2))) !=
               EmptyBlocks.end()) {
          thelastBlockOld = dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(2));
          thelastinst->get()->replaceOperand(
              2, EmptyBlocks[dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(2))]);
        }
        basicBlock->removeSuccessor(OldElseBlock);
        OldElseBlock->removePredecessor(basicBlock.get());
        // 处理 then 和 else 分支合并的情况
        if (dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1)) ==
            dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(2))) {
          auto thebrBlock = dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(1));
          usedelete(thelastinst->get());
          thelastinst = basicBlock->getInstructions().erase(thelastinst);
          pBuilder->setPosition(basicBlock.get(), basicBlock->end());
          pBuilder->createUncondBrInst(thebrBlock, {});
          continue;
        }
        basicBlock->addSuccessor(dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(2)));
        dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(2))->addPredecessor(basicBlock.get());

        if (thelastBlockOld != nullptr) {
          int indexphi = 0;
          for (auto &pred : dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(2))->getPredecessors()) {
            if (pred == thelastBlockOld) {
              break;
            }
            indexphi++;
          }
          for (auto &InstInNew : dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(2))->getInstructions()) {
            if (InstInNew->isPhi()) {
              dynamic_cast<PhiInst *>(InstInNew.get())->removeOperand(indexphi + 1);
            } else {
              break;
            }
          }
        }
      } else {
        if (basicBlock->getNumSuccessors() == 1) {
          pBuilder->setPosition(basicBlock.get(), basicBlock->end());
          pBuilder->createUncondBrInst(basicBlock->getSuccessors()[0], {});
          auto thelastinst = basicBlock->getInstructions().end();
          (--thelastinst);
          auto OldBrBlock = dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0));
          sysy::BasicBlock *thelastBlockOld = nullptr;
          while (EmptyBlocks.find(dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0))) !=
                 EmptyBlocks.end()) {
            thelastBlockOld = dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0));

            thelastinst->get()->replaceOperand(
                0, EmptyBlocks[dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0))]);
          }
          
          basicBlock->removeSuccessor(OldBrBlock);
          OldBrBlock->removePredecessor(basicBlock.get());
          basicBlock->addSuccessor(dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0)));
          dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0))->addPredecessor(basicBlock.get());
          if (thelastBlockOld != nullptr) {
            int indexphi = 0;
            for (auto &pred : dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0))->getPredecessors()) {
              if (pred == thelastBlockOld) {
                break;
              }
              indexphi++;
            }

            for (auto &InstInNew : dynamic_cast<BasicBlock *>(thelastinst->get()->getOperand(0))->getInstructions()) {
              if (InstInNew->isPhi()) {
                dynamic_cast<PhiInst *>(InstInNew.get())->removeOperand(indexphi + 1);
              } else {
                break;
              }
            }
          }
        }
      }
    }

    for (auto iter = function.second->getBasicBlocks().begin(); iter != function.second->getBasicBlocks().end();) {
      
      if (EmptyBlocks.find(iter->get()) != EmptyBlocks.end()) {
        // EntryBlock跳过
        if (iter->get() == function.second->getEntryBlock()) {
          ++iter;
          continue;
        }

        for (auto &iterInst : iter->get()->getInstructions()) 
          usedelete(iterInst.get());
        // 删除不可达基本块的phi指令的操作数
        for (auto &succ : iter->get()->getSuccessors()) {
          int index = 0;
          for (auto &pred : succ->getPredecessors()) {
            if (pred == iter->get()) {
              break;
            }
            index++;
          }

          for (auto &instinsucc : succ->getInstructions()) {
            if (instinsucc->isPhi()) {
              dynamic_cast<PhiInst *>(instinsucc.get())->removeOperand(index);
            } else {
              break;
            }
          }
        }

        function.second->removeBasicBlock((iter++)->get());
      } else {
        ++iter;
      }
    }
  }
}

// 如果函数没有返回指令，则添加一个默认返回指令(主要解决void函数没有返回指令的问题)
void SysYOptPre::SysYAddReturn() {
  auto &functions = pModule->getFunctions();
  for (auto &function : functions) {
    auto &func = function.second;
    auto basicBlocks = func->getBasicBlocks();
    for (auto &block : basicBlocks) {
      if (block->getNumSuccessors() == 0) {
        // 如果基本块没有后继块，则添加一个返回指令
        if (block->getNumInstructions() == 0) {
          pBuilder->setPosition(block.get(), block->end());
          pBuilder->createReturnInst({});
        }
        auto thelastinst = block->getInstructions().end();
        --thelastinst;
        if (thelastinst->get()->getKind() != Instruction::kReturn) {
          pBuilder->setPosition(block.get(), block->end());
          // TODO: 如果int float函数缺少返回值是否需要报错
          if (func->getReturnType()->isInt()) {
            pBuilder->createReturnInst(ConstantValue::get(0));
          } else if (func->getReturnType()->isFloat()) {
            pBuilder->createReturnInst(ConstantValue::get(0.0F));
          } else {
            pBuilder->createReturnInst({});
          }
        }
      }
    }
  }
}

}  // namespace sysy