#include "Dom.h"
#include <limits> // for std::numeric_limits
#include <queue>

namespace sysy {

// 初始化 支配树静态 ID
void *DominatorTreeAnalysisPass::ID = (void *)&DominatorTreeAnalysisPass::ID;
// ==============================================================
// DominatorTree 结果类的实现
// ==============================================================

DominatorTree::DominatorTree(Function *F) : AssociatedFunction(F) {
  // 构造时可以不计算，在分析遍运行里计算并填充
}

const std::set<BasicBlock *> *DominatorTree::getDominators(BasicBlock *BB) const {
  auto it = Dominators.find(BB);
  if (it != Dominators.end()) {
    return &(it->second);
  }
  return nullptr;
}

BasicBlock *DominatorTree::getImmediateDominator(BasicBlock *BB) const {
  auto it = IDoms.find(BB);
  if (it != IDoms.end()) {
    return it->second;
  }
  return nullptr;
}

const std::set<BasicBlock *> *DominatorTree::getDominanceFrontier(BasicBlock *BB) const {
  auto it = DominanceFrontiers.find(BB);
  if (it != DominanceFrontiers.end()) {
    return &(it->second);
  }
  return nullptr;
}

void DominatorTree::computeDominators(Function *F) {
  // 经典的迭代算法计算支配者集合
  // TODO: 可以替换为更高效的算法，如 Lengauer-Tarjan 算法
  BasicBlock *entryBlock = F->getEntryBlock();

  for (const auto &bb_ptr : F->getBasicBlocks()) {
    BasicBlock *bb = bb_ptr.get();
    if (bb == entryBlock) {
      Dominators[bb].insert(bb);
    } else {
      for (const auto &all_bb_ptr : F->getBasicBlocks()) {
        Dominators[bb].insert(all_bb_ptr.get());
      }
    }
  }

  bool changed = true;
  while (changed) {
    changed = false;
    for (const auto &bb_ptr : F->getBasicBlocks()) {
      BasicBlock *bb = bb_ptr.get();
      if (bb == entryBlock)
        continue;

      std::set<BasicBlock *> newDom;
      bool firstPred = true;
      for (BasicBlock *pred : bb->getPredecessors()) {
        if (Dominators.count(pred)) {
          if (firstPred) {
            newDom = Dominators[pred];
            firstPred = false;
          } else {
            std::set<BasicBlock *> intersection;
            std::set_intersection(newDom.begin(), newDom.end(), Dominators[pred].begin(), Dominators[pred].end(),
                                  std::inserter(intersection, intersection.begin()));
            newDom = intersection;
          }
        }
      }
      newDom.insert(bb);

      if (newDom != Dominators[bb]) {
        Dominators[bb] = newDom;
        changed = true;
      }
    }
  }
}

void DominatorTree::computeIDoms(Function *F) {
  // 采用与之前类似的简化实现。TODO:Lengauer-Tarjan等算法。
  BasicBlock *entryBlock = F->getEntryBlock();
  IDoms[entryBlock] = nullptr;

  for (const auto &bb_ptr : F->getBasicBlocks()) {
    BasicBlock *bb = bb_ptr.get();
    if (bb == entryBlock)
      continue;

    BasicBlock *currentIDom = nullptr;
    const std::set<BasicBlock *> *domsOfBB = getDominators(bb);
    if (!domsOfBB)
      continue;

    for (BasicBlock *D : *domsOfBB) {
      if (D == bb)
        continue;

      bool isCandidateIDom = true;
      for (BasicBlock *candidate : *domsOfBB) {
        if (candidate == bb || candidate == D)
          continue;
        const std::set<BasicBlock *> *domsOfCandidate = getDominators(candidate);
        if (domsOfCandidate && domsOfCandidate->count(D) == 0 && domsOfBB->count(candidate)) {
          isCandidateIDom = false;
          break;
        }
      }
      if (isCandidateIDom) {
        currentIDom = D;
        break;
      }
    }
    IDoms[bb] = currentIDom;
  }
}

void DominatorTree::computeDominanceFrontiers(Function *F) {
  // 经典的支配边界计算算法
  for (const auto &bb_ptr_X : F->getBasicBlocks()) {
    BasicBlock *X = bb_ptr_X.get();
    DominanceFrontiers[X].clear();

    for (BasicBlock *Y : X->getSuccessors()) {
      const std::set<BasicBlock *> *domsOfY = getDominators(Y);
      if (domsOfY && domsOfY->find(X) == domsOfY->end()) {
        DominanceFrontiers[X].insert(Y);
      }
    }

    const std::set<BasicBlock *> *domsOfX = getDominators(X);
    if (!domsOfX)
      continue;
    for (const auto &bb_ptr_Z : F->getBasicBlocks()) {
      BasicBlock *Z = bb_ptr_Z.get();
      if (Z == X)
        continue;
      const std::set<BasicBlock *> *domsOfZ = getDominators(Z);
      if (domsOfZ && domsOfZ->count(X) && Z != X) {

        for (BasicBlock *Y : Z->getSuccessors()) {
          const std::set<BasicBlock *> *domsOfY = getDominators(Y);
          if (domsOfY && domsOfY->find(X) == domsOfY->end()) {
            DominanceFrontiers[X].insert(Y);
          }
        }
      }
    }
  }
}

// ==============================================================
// DominatorTreeAnalysisPass 的实现
// ==============================================================


bool DominatorTreeAnalysisPass::runOnFunction(Function* F, AnalysisManager &AM) {
  CurrentDominatorTree = std::make_unique<DominatorTree>(F);
  CurrentDominatorTree->computeDominators(F);
  CurrentDominatorTree->computeIDoms(F);
  CurrentDominatorTree->computeDominanceFrontiers(F);
  return false;
}

std::unique_ptr<AnalysisResultBase> DominatorTreeAnalysisPass::getResult() {
  // 返回计算好的 DominatorTree 实例，所有权转移给 AnalysisManager
  return std::move(CurrentDominatorTree);
}

} // namespace sysy