#pragma once

#include "llir.h"
#include "pass.h"

namespace CompSysY
{
void gen_dominance(FunctionPtr_t func);
void gen_dominance_frontier(FunctionPtr_t func);
void update_dfs_numbers(BasicBlockPtr_t bb, bool rst);

struct CFGNode
{
  BasicBlock *bb                   = nullptr;
  int special_node                 = 0;
  std::vector<CFGNode *> pred_list = {};
  std::vector<CFGNode *> succ_list = {};
  std::unordered_set<CFGNode *> DF = {};
  CFGNode *idom;                           // immediate dominator of this node
  std::unordered_set<CFGNode *> dom;       // dominators of this node
  std::unordered_set<CFGNode *> dominated; // nodes immediately dominated by this node
  static sptr(CFGNode) New(BasicBlock *bb = nullptr)
  {
    auto newnode = std::make_shared<CFGNode>();
    newnode->bb  = bb;
    return newnode;
  }
  static sptr(CFGNode) New(int special_node)
  {
    auto newnode          = std::make_shared<CFGNode>();
    newnode->bb           = nullptr;
    newnode->special_node = special_node;
    return newnode;
  }
  std::string to_string() const
  {
    if (is_entry()) return "ENTRY";
    if (is_exit()) return "EXIT";
    return bb->name;
  }
  bool is_entry() const
  {
    return special_node == 1;
  }
  bool is_exit() const
  {
    return special_node == 2;
  }
  void print_cfg(std::ostream &ostr) const
  {
    ostr << "---" << to_string() << "---\n";
    ostr << "  pred: [";
    for (auto pred : pred_list)
      {
        ostr << pred->to_string() << ", ";
      }
    ostr << "]\n";
    ostr << "  succ: [";
    for (auto succ : succ_list)
      {
        ostr << succ->to_string() << ", ";
      }
    ostr << "]\n";
    ostr.flush();
  }
  void print_dom(std::ostream &ostr) const
  {
    print_cfg(ostr);
    ostr << "  idom: " << (!idom ? "NULL" : idom->to_string()) << "\n";
    ostr << "  domer: [";
    for (auto d : dom)
      {
        ostr << d->to_string() << ", ";
      }
    ostr << "]\n";
    ostr << "  domee: [";
    for (auto domee : dominated)
      {
        ostr << domee->to_string() << ", ";
      }
    ostr << "]\n";
    ostr << "  DF: [";
    for (auto df : DF)
      {
        ostr << df->to_string() << ", ";
      }
    ostr << "]\n";
    ostr.flush();
  }
};

class CFG
{
public:
  std::list<sptr(CFGNode)> node_pool;
  CFGNode *entry_node;
  CFGNode *exit_node;
  std::unordered_map<BasicBlock *, CFGNode *> bb_node;
  void build_from_func(Function *func, bool reversed);
  void print_cfg();
  void print_dom();
  void calculate();

private:
  std::unordered_set<CFGNode *> visit;
  void _print_cfg(CFGNode *);
  void _print_dom(CFGNode *);
};

} // namespace CompSysY