#include "ir/Analysis.h"

#include <algorithm>
#include <functional>

namespace ir {
namespace {

std::vector<BasicBlock*> BuildReversePostOrder(Function& function) {
  std::vector<BasicBlock*> post_order;
  auto* entry = function.GetEntryBlock();
  if (!entry) {
    return post_order;
  }

  std::unordered_set<BasicBlock*> visited;
  std::function<void(BasicBlock*)> dfs = [&](BasicBlock* block) {
    if (!block || !visited.insert(block).second) {
      return;
    }
    for (auto* succ : block->GetSuccessors()) {
      dfs(succ);
    }
    post_order.push_back(block);
  };
  dfs(entry);
  std::reverse(post_order.begin(), post_order.end());
  return post_order;
}

}  // namespace

DominatorTree::DominatorTree(Function& function) : function_(&function) {
  Recalculate();
}

void DominatorTree::Recalculate() {
  reverse_post_order_ = BuildReversePostOrder(*function_);
  block_index_.clear();
  dominates_.clear();
  immediate_dominator_.clear();
  dom_children_.clear();

  const auto num_blocks = reverse_post_order_.size();
  for (std::size_t i = 0; i < num_blocks; ++i) {
    block_index_.emplace(reverse_post_order_[i], i);
  }
  if (num_blocks == 0) {
    return;
  }

  dominates_.assign(num_blocks, std::vector<std::uint8_t>(num_blocks, 1));
  dominates_[0].assign(num_blocks, 0);
  dominates_[0][0] = 1;

  bool changed = true;
  while (changed) {
    changed = false;
    for (std::size_t i = 1; i < num_blocks; ++i) {
      auto* block = reverse_post_order_[i];
      std::vector<std::uint8_t> next(num_blocks, 1);
      bool has_reachable_pred = false;
      for (auto* pred : block->GetPredecessors()) {
        auto pred_it = block_index_.find(pred);
        if (pred_it == block_index_.end()) {
          continue;
        }
        has_reachable_pred = true;
        const auto& pred_dom = dominates_[pred_it->second];
        for (std::size_t bit = 0; bit < num_blocks; ++bit) {
          next[bit] &= pred_dom[bit];
        }
      }
      if (!has_reachable_pred) {
        next.assign(num_blocks, 0);
      }
      next[i] = 1;
      if (next != dominates_[i]) {
        dominates_[i] = std::move(next);
        changed = true;
      }
    }
  }

  for (std::size_t i = 1; i < num_blocks; ++i) {
    auto* block = reverse_post_order_[i];
    BasicBlock* idom = nullptr;
    for (std::size_t candidate = 0; candidate < num_blocks; ++candidate) {
      if (candidate == i || !dominates_[i][candidate]) {
        continue;
      }
      auto* candidate_block = reverse_post_order_[candidate];
      bool immediate = true;
      for (std::size_t other = 0; other < num_blocks; ++other) {
        if (other == i || other == candidate || !dominates_[i][other]) {
          continue;
        }
        if (Dominates(reverse_post_order_[other], candidate_block)) {
          immediate = false;
          break;
        }
      }
      if (immediate) {
        idom = candidate_block;
        break;
      }
    }
    immediate_dominator_.emplace(block, idom);
    if (idom) {
      dom_children_[idom].push_back(block);
    }
  }
}

bool DominatorTree::IsReachable(BasicBlock* block) const {
  return block != nullptr && block_index_.find(block) != block_index_.end();
}

bool DominatorTree::Dominates(BasicBlock* dom, BasicBlock* node) const {
  if (!dom || !node) {
    return false;
  }
  const auto dom_it = block_index_.find(dom);
  const auto node_it = block_index_.find(node);
  if (dom_it == block_index_.end() || node_it == block_index_.end()) {
    return false;
  }
  return dominates_[node_it->second][dom_it->second] != 0;
}

bool DominatorTree::Dominates(Instruction* dom, Instruction* user) const {
  if (!dom || !user) {
    return false;
  }
  if (dom == user) {
    return true;
  }

  auto* dom_block = dom->GetParent();
  auto* user_block = user->GetParent();
  if (dom_block != user_block) {
    return Dominates(dom_block, user_block);
  }

  for (const auto& inst_ptr : dom_block->GetInstructions()) {
    if (inst_ptr.get() == dom) {
      return true;
    }
    if (inst_ptr.get() == user) {
      return false;
    }
  }
  return false;
}

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

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

}  // namespace ir