#pragma once

#include "ir/IR.h"

#include <cstdint>
#include <memory>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace ir {

class DominatorTree {
 public:
  explicit DominatorTree(Function& function);

  void Recalculate();

  Function& GetFunction() const { return *function_; }
  bool IsReachable(BasicBlock* block) const;
  bool Dominates(BasicBlock* dom, BasicBlock* node) const;
  bool Dominates(Instruction* dom, Instruction* user) const;
  BasicBlock* GetIDom(BasicBlock* block) const;
  const std::vector<BasicBlock*>& GetChildren(BasicBlock* block) const;
  const std::vector<BasicBlock*>& GetReversePostOrder() const {
    return reverse_post_order_;
  }

 private:
  Function* function_ = nullptr;
  std::vector<BasicBlock*> reverse_post_order_;
  std::unordered_map<BasicBlock*, std::size_t> block_index_;
  std::vector<std::vector<std::uint8_t>> dominates_;
  std::unordered_map<BasicBlock*, BasicBlock*> immediate_dominator_;
  std::unordered_map<BasicBlock*, std::vector<BasicBlock*>> dom_children_;
};

struct Loop {
  BasicBlock* header = nullptr;
  std::unordered_set<BasicBlock*> blocks;
  std::vector<BasicBlock*> block_list;
  std::vector<BasicBlock*> latches;
  std::vector<BasicBlock*> exiting_blocks;
  std::vector<BasicBlock*> exit_blocks;
  BasicBlock* preheader = nullptr;
  Loop* parent = nullptr;
  std::vector<Loop*> subloops;

  bool Contains(BasicBlock* block) const;
  bool Contains(const Loop* other) const;
  bool IsInnermost() const;
};

class LoopInfo {
 public:
  LoopInfo(Function& function, const DominatorTree& dom_tree);

  void Recalculate();

  const std::vector<std::unique_ptr<Loop>>& GetLoops() const { return loops_; }
  std::vector<Loop*> GetTopLevelLoops() const;
  std::vector<Loop*> GetLoopsInPostOrder() const;
  Loop* GetLoopFor(BasicBlock* block) const;

 private:
  Function* function_ = nullptr;
  const DominatorTree* dom_tree_ = nullptr;
  std::vector<std::unique_ptr<Loop>> loops_;
  std::vector<Loop*> top_level_loops_;
  std::unordered_map<BasicBlock*, Loop*> block_to_loop_;
};

}  // namespace ir