#pragma once

#include "Pass.h" // 包含 Pass 框架
#include "IR.h"   // 包含 IR 定义
#include <map>
#include <set>
#include <vector>
#include <algorithm>

namespace sysy {

// 支配树分析结果类 (保持不变)
class DominatorTree : public AnalysisResultBase {
public:
    DominatorTree(Function* F);
    const std::set<BasicBlock*>* getDominators(BasicBlock* BB) const;
    BasicBlock* getImmediateDominator(BasicBlock* BB) const;
    const std::set<BasicBlock*>* getDominanceFrontier(BasicBlock* BB) const;
    const std::map<BasicBlock*, std::set<BasicBlock*>>& getDominatorsMap() const { return Dominators; }
    const std::map<BasicBlock*, BasicBlock*>& getIDomsMap() const { return IDoms; }
    const std::map<BasicBlock*, std::set<BasicBlock*>>& getDominanceFrontiersMap() const { return DominanceFrontiers; }
    void computeDominators(Function* F);
    void computeIDoms(Function* F);
    void computeDominanceFrontiers(Function* F);
private:
    Function* AssociatedFunction;
    std::map<BasicBlock*, std::set<BasicBlock*>> Dominators;
    std::map<BasicBlock*, BasicBlock*> IDoms;
    std::map<BasicBlock*, std::set<BasicBlock*>> DominanceFrontiers;
};


// 支配树分析遍
class DominatorTreeAnalysisPass : public AnalysisPass {
public:
    // 唯一的 Pass ID
    static void *ID;

    DominatorTreeAnalysisPass() : AnalysisPass("DominatorTreeAnalysis", Pass::Granularity::Function) {}

    // 实现 getPassID
    void* getPassID() const override { return &ID; }

    bool runOnFunction(Function* F, AnalysisManager &AM) override;

    std::unique_ptr<AnalysisResultBase> getResult() override;

private:
    std::unique_ptr<DominatorTree> CurrentDominatorTree;
};

} // namespace sysy