#pragma once

#include "ir/IR.h"

#include <unordered_map>
#include <vector>

namespace ir {

class DominatorTree {
 public:
  DominatorTree() = default;
  ~DominatorTree() = default;

  void Recalculate(Function& function);

  BasicBlock* GetRoot() const;
  BasicBlock* GetIDom(BasicBlock* block) const;
  bool Dominates(BasicBlock* a, BasicBlock* b) const;
  const std::vector<BasicBlock*>& GetChildren(BasicBlock* block) const;
  const std::vector<BasicBlock*>& GetDominanceFrontier(BasicBlock* block) const;
  const std::vector<BasicBlock*>& GetPredecessors(BasicBlock* block) const;
  const std::vector<BasicBlock*>& GetSuccessors(BasicBlock* block) const;

 private:
  void BuildCFG(Function& function);
  void ComputeIDoms();
  void ComputeDominanceFrontiers();
  BasicBlock* Intersect(BasicBlock* first, BasicBlock* second) const;

  std::vector<BasicBlock*> blocks_;
  std::unordered_map<BasicBlock*, std::vector<BasicBlock*>> preds_;
  std::unordered_map<BasicBlock*, std::vector<BasicBlock*>> succs_;
  std::unordered_map<BasicBlock*, BasicBlock*> idom_;
  std::unordered_map<BasicBlock*, std::vector<BasicBlock*>> children_;
  std::unordered_map<BasicBlock*, std::vector<BasicBlock*>> dominance_frontier_;
  std::unordered_map<BasicBlock*, int> dfs_number_;
};

}  // namespace ir