#include "ir/PassManager.h"

#include "ir/IR.h"
#include "PassUtils.h"

#include <vector>

namespace ir {
namespace {

bool TryGetConstBranchTarget(CondBrInst* br, BasicBlock*& target, BasicBlock*& removed) {
  if (!br) {
    return false;
  }
  auto* then_block = br->GetThenBlock();
  auto* else_block = br->GetElseBlock();
  if (then_block == else_block) {
    target = then_block;
    removed = nullptr;
    return true;
  }
  if (auto* cond = dyncast<ConstantI1>(br->GetCondition())) {
    target = cond->GetValue() ? then_block : else_block;
    removed = cond->GetValue() ? else_block : then_block;
    return true;
  }
  return false;
}

bool SimplifyBlockTerminator(BasicBlock* block) {
  if (!block || block->GetInstructions().empty()) {
    return false;
  }
  auto* term = block->GetInstructions().back().get();
  auto* condbr = dyncast<CondBrInst>(term);
  if (!condbr) {
    return false;
  }

  BasicBlock* target = nullptr;
  BasicBlock* removed = nullptr;
  if (!TryGetConstBranchTarget(condbr, target, removed)) {
    return false;
  }

  if (removed) {
    passutils::RemoveIncomingFromSuccessor(removed, block);
    removed->RemovePredecessor(block);
    block->RemoveSuccessor(removed);
  }
  passutils::ReplaceTerminatorWithBr(block, target);
  return true;
}

bool SimplifyPhiNodes(Function& function) {
  bool changed = false;
  for (const auto& block_ptr : function.GetBlocks()) {
    bool local_changed = true;
    while (local_changed) {
      local_changed = false;
      for (const auto& inst_ptr : block_ptr->GetInstructions()) {
        auto* phi = dyncast<PhiInst>(inst_ptr.get());
        if (!phi) {
          break;
        }
        if (!passutils::SimplifyPhiInst(phi)) {
          continue;
        }
        local_changed = true;
        changed = true;
        break;
      }
    }
  }
  return changed;
}

bool RunCFGSimplifyOnFunction(Function& function) {
  if (function.IsExternal() || function.GetEntryBlock() == nullptr) {
    return false;
  }

  bool changed = false;

  for (const auto& block_ptr : function.GetBlocks()) {
    changed |= SimplifyBlockTerminator(block_ptr.get());
  }

  changed |= passutils::RemoveUnreachableBlocks(function);
  changed |= SimplifyPhiNodes(function);

  return changed;
}

}  // namespace

bool RunCFGSimplify(Module& module) {
  bool changed = false;
  for (const auto& function : module.GetFunctions()) {
    if (function) {
      changed |= RunCFGSimplifyOnFunction(*function);
    }
  }
  return changed;
}

}  // namespace ir