// 支配树分析：
// - 构建/查询 Dominator Tree 及相关关系
// - 使用 Cooper-Harvey-Kennedy 算法，近线性时间复杂度

#include "ir/IR.h"

#include <algorithm>
#include <functional>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace ir {

void DominatorTree::Compute(Function& func) {
  func.RebuildCFG();

  // Build block list and reverse postorder (RPO)
  std::vector<BasicBlock*> blocks;
  std::unordered_map<BasicBlock*, int> rpo;
  {
    std::vector<BasicBlock*> rpo_vec;
    std::unordered_set<BasicBlock*> visited;
    std::function<void(BasicBlock*)> dfs = [&](BasicBlock* bb) {
      if (!bb || visited.count(bb)) return;
      visited.insert(bb);
      for (auto* succ : bb->GetSuccessors()) {
        dfs(succ);
      }
      rpo_vec.push_back(bb);
    };
    dfs(func.GetEntry());
    // Reverse to get RPO (postorder reversed)
    std::reverse(rpo_vec.begin(), rpo_vec.end());
    blocks = rpo_vec;
    for (int i = 0; i < (int)blocks.size(); ++i) {
      rpo[blocks[i]] = i;
    }
  }
  if (blocks.empty()) return;
  int n = (int)blocks.size();

  auto* entry = func.GetEntry();
  if (!entry) return;

  // ─── 1. CHK algorithm for immediate dominators ─────────────────────────
  idom_.clear();
  idom_[entry] = entry;  // entry is its own dominator

  // Intersect: find common ancestor of b1 and b2 walking up the dom tree
  // Uses RPO number: a dominator always has lower RPO number
  auto intersect = [&](BasicBlock* b1, BasicBlock* b2) -> BasicBlock* {
    auto i1 = rpo.find(b1), i2 = rpo.find(b2);
    if (i1 == rpo.end() || i2 == rpo.end()) return entry;
    int r1 = i1->second, r2 = i2->second;
    while (b1 != b2) {
      while (r1 > r2) {
        auto it = idom_.find(b1);
        if (it == idom_.end() || it->second == b1) return b1;
        b1 = it->second;
        r1 = rpo[b1];
      }
      while (r2 > r1) {
        auto it = idom_.find(b2);
        if (it == idom_.end() || it->second == b2) return b2;
        b2 = it->second;
        r2 = rpo[b2];
      }
    }
    return b1;
  };

  bool changed = true;
  while (changed) {
    changed = false;
    // Process in RPO (skip entry which is first in RPO)
    for (int i = 1; i < n; ++i) {
      auto* bb = blocks[i];
      // Find first predecessor with defined IDOM
      BasicBlock* new_idom = nullptr;
      for (auto* pred : bb->GetPredecessors()) {
        if (idom_.count(pred) && pred != bb) {
          new_idom = pred;
          break;
        }
      }
      if (!new_idom) continue;

      // Intersect with remaining predecessors
      for (auto* pred : bb->GetPredecessors()) {
        if (pred == new_idom || pred == bb) continue;
        if (idom_.count(pred)) {
          new_idom = intersect(pred, new_idom);
        }
      }

      auto old = idom_.find(bb);
      if (old == idom_.end() || old->second != new_idom) {
        idom_[bb] = new_idom;
        changed = true;
      }
    }
  }

  // Entry is its own IDOM, set to nullptr for external queries
  idom_[entry] = nullptr;

  // Unreached blocks get entry as IDOM
  for (auto* bb : blocks) {
    if (!idom_.count(bb)) idom_[bb] = entry;
  }

  // ─── 2. Build children map and dom levels ──────────────────────────────
  children_.clear();
  dom_level_.clear();
  for (auto& [child, parent] : idom_) {
    if (parent) children_[parent].push_back(child);
  }

  // BFS to compute dom levels
  std::queue<BasicBlock*> q;
  dom_level_[entry] = 0;
  q.push(entry);
  while (!q.empty()) {
    auto* cur = q.front();
    q.pop();
    size_t cur_level = dom_level_[cur];
    auto it = children_.find(cur);
    if (it != children_.end()) {
      for (auto* child : it->second) {
        dom_level_[child] = cur_level + 1;
        q.push(child);
      }
    }
  }

  // ─── 3. Compute dominance frontier ─────────────────────────────────────
  df_.clear();
  for (int i = 0; i < n; ++i) {
    auto* b = blocks[i];
    if (b->GetPredecessors().size() < 2) continue;
    for (auto* p : b->GetPredecessors()) {
      auto* runner = p;
      auto* b_idom = GetIDom(b);
      while (runner != b_idom) {
        if (!runner) break;
        df_[runner].push_back(b);
        runner = GetIDom(runner);
      }
    }
  }
  // Deduplicate DF entries
  for (auto& [bb, vec] : df_) {
    std::sort(vec.begin(), vec.end());
    vec.erase(std::unique(vec.begin(), vec.end()), vec.end());
  }

  // ─── 4. Compute DFS order of dominator tree ────────────────────────────
  df_order_.clear();
  visited_.clear();
  std::function<void(BasicBlock*)> dfs_tree = [&](BasicBlock* bb) {
    if (!bb || visited_.count(bb)) return;
    visited_.insert(bb);
    df_order_.push_back(bb);
    auto it = children_.find(bb);
    if (it != children_.end()) {
      for (auto* child : it->second) {
        dfs_tree(child);
      }
    }
  };
  dfs_tree(entry);
}

BasicBlock* DominatorTree::GetIDom(BasicBlock* bb) const {
  auto it = idom_.find(bb);
  return (it != idom_.end()) ? it->second : nullptr;
}

const std::vector<BasicBlock*>& DominatorTree::GetChildren(
    BasicBlock* bb) const {
  static const std::vector<BasicBlock*> empty;
  auto it = children_.find(bb);
  return (it != children_.end()) ? it->second : empty;
}

const std::vector<BasicBlock*>& DominatorTree::GetDominanceFrontier(
    BasicBlock* bb) const {
  static const std::vector<BasicBlock*> empty;
  auto it = df_.find(bb);
  return (it != df_.end()) ? it->second : empty;
}

bool DominatorTree::Dominates(BasicBlock* a, BasicBlock* b) const {
  if (a == b) return true;
  BasicBlock* runner = b;
  while (runner) {
    runner = GetIDom(runner);
    if (runner == a) return true;
  }
  return false;
}

}  // namespace ir