#include "ir/analysis/DominatorTree.h"

#include <algorithm>
#include <functional>
#include <unordered_set>

namespace ir {

namespace {

std::vector<BasicBlock*> GetBlockSuccessors(BasicBlock* block) {
  std::vector<BasicBlock*> succs;
  if (!block) {
    return succs;
  }

  const auto& instructions = block->GetInstructions();
  if (instructions.empty()) {
    return succs;
  }

  Instruction* term = instructions.back().get();
  if (!term->IsTerminator()) {
    return succs;
  }

  if (term->GetOpcode() == Opcode::Br) {
    auto* br = static_cast<BranchInst*>(term);
    succs.push_back(br->GetTarget());
  } else if (term->GetOpcode() == Opcode::CondBr) {
    auto* br = static_cast<BranchInst*>(term);
    succs.push_back(br->GetTrueTarget());
    succs.push_back(br->GetFalseTarget());
  }
  return succs;
}

}  // namespace

void DominatorTree::Recalculate(Function& function) {
  BuildCFG(function);
  if (blocks_.empty()) {
    return;
  }

  idom_.clear();
  ComputeIDoms();
  ComputeDominanceFrontiers();
}

BasicBlock* DominatorTree::GetRoot() const {
  if (blocks_.empty()) {
    return nullptr;
  }
  return blocks_.front();
}

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

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

  auto it = idom_.find(b);
  while (it != idom_.end() && it->second != b) {
    if (it->second == a) {
      return true;
    }
    b = it->second;
    it = idom_.find(b);
  }
  return false;
}

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

const std::vector<BasicBlock*>& DominatorTree::GetDominanceFrontier(BasicBlock* block) const {
  auto it = dominance_frontier_.find(block);
  if (it == dominance_frontier_.end()) {
    static const std::vector<BasicBlock*> empty;
    return empty;
  }
  return it->second;
}

const std::vector<BasicBlock*>& DominatorTree::GetPredecessors(BasicBlock* block) const {
  auto it = preds_.find(block);
  if (it == preds_.end()) {
    static const std::vector<BasicBlock*> empty;
    return empty;
  }
  return it->second;
}

const std::vector<BasicBlock*>& DominatorTree::GetSuccessors(BasicBlock* block) const {
  auto it = succs_.find(block);
  if (it == succs_.end()) {
    static const std::vector<BasicBlock*> empty;
    return empty;
  }
  return it->second;
}

void DominatorTree::BuildCFG(Function& function) {
  blocks_.clear();
  preds_.clear();
  succs_.clear();
  idom_.clear();
  children_.clear();
  dominance_frontier_.clear();
  dfs_number_.clear();

  BasicBlock* entry = function.GetEntry();
  if (!entry) {
    return;
  }

  std::unordered_set<BasicBlock*> visited;
  int next_number = 0;

  std::function<void(BasicBlock*)> dfs = [&](BasicBlock* block) {
    if (!block || visited.count(block)) {
      return;
    }
    visited.insert(block);
    dfs_number_[block] = next_number++;
    blocks_.push_back(block);

    auto successors = GetBlockSuccessors(block);
    succs_[block] = successors;
    for (BasicBlock* succ : successors) {
      preds_[succ].push_back(block);
      dfs(succ);
    }
  };

  dfs(entry);
}

void DominatorTree::ComputeIDoms() {
  if (blocks_.empty()) {
    return;
  }

  BasicBlock* entry = blocks_.front();
  idom_[entry] = entry;

  bool changed = true;
  while (changed) {
    changed = false;
    for (BasicBlock* block : blocks_) {
      if (block == entry) {
        continue;
      }

      const auto& predecessors = preds_[block];
      BasicBlock* new_idom = nullptr;
      for (BasicBlock* pred : predecessors) {
        auto pred_it = idom_.find(pred);
        if (pred_it == idom_.end()) {
          continue;
        }
        if (!new_idom) {
          new_idom = pred;
        } else {
          new_idom = Intersect(pred, new_idom);
        }
      }

      if (!new_idom) {
        continue;
      }
      if (idom_.find(block) == idom_.end() || idom_[block] != new_idom) {
        idom_[block] = new_idom;
        changed = true;
      }
    }
  }

  children_.clear();
  for (const auto& pair : idom_) {
    BasicBlock* block = pair.first;
    BasicBlock* parent = pair.second;
    if (block != parent) {
      children_[parent].push_back(block);
    }
  }
}

void DominatorTree::ComputeDominanceFrontiers() {
  dominance_frontier_.clear();

  for (BasicBlock* block : blocks_) {
    const auto& predecessors = preds_[block];
    if (predecessors.size() < 2) {
      continue;
    }

    for (BasicBlock* pred : predecessors) {
      BasicBlock* runner = pred;
      while (runner != idom_[block]) {
        auto& frontier = dominance_frontier_[runner];
        if (std::find(frontier.begin(), frontier.end(), block) == frontier.end()) {
          frontier.push_back(block);
        }
        runner = idom_[runner];
      }
    }
  }
}

BasicBlock* DominatorTree::Intersect(BasicBlock* first, BasicBlock* second) const {
  std::unordered_set<BasicBlock*> first_ancestors;
  for (BasicBlock* block = first; block;) {
    first_ancestors.insert(block);
    auto it = idom_.find(block);
    if (it == idom_.end() || it->second == block) {
      break;
    }
    block = it->second;
  }

  for (BasicBlock* block = second; block;) {
    if (first_ancestors.count(block)) {
      return block;
    }
    auto it = idom_.find(block);
    if (it == idom_.end() || it->second == block) {
      break;
    }
    block = it->second;
  }

  return GetRoot();
}

}  // namespace ir