#pragma once

#include "ir/IR.h"
#include <vector>
#include <unordered_map>
#include <unordered_set>

namespace ir {

// Dominator Tree Analysis
class DominatorTree {
 public:
  explicit DominatorTree(Function* func);
  void Run();

  // Query interfaces
  BasicBlock* GetIdom(BasicBlock* bb) const;
  const std::vector<BasicBlock*>& GetDominatedBlocks(BasicBlock* bb) const;
  const std::vector<BasicBlock*>& GetDominanceFrontier(BasicBlock* bb) const;
  bool Dominates(BasicBlock* a, BasicBlock* b) const;

 private:
  Function* func_;
  std::vector<BasicBlock*> rpo_;
  std::unordered_map<BasicBlock*, BasicBlock*> idom_;
  std::unordered_map<BasicBlock*, std::vector<BasicBlock*>> dom_tree_;
  std::unordered_map<BasicBlock*, std::vector<BasicBlock*>> df_;

  void ComputeRPO();
  void ComputeIdom();
  void ComputeDomTree();
  void ComputeDF();
};

// Individual Pass Declarations
bool RunMem2Reg(Function* func, Context& ctx);
bool RunConstProp(Function* func, Context& ctx);
bool RunConstFold(Function* func, Context& ctx);
bool RunDCE(Function* func);
bool RunCFGSimplify(Function* func);
bool RunCSE(Function* func);

// Run the optimization pipeline on a Function or Module
void RunOptimizationPasses(Module& module);
void RunFunctionOptimizationPasses(Function* func, Context& ctx);

} // namespace ir