// 假设 Mem2Reg.h 看起来像这样 (你需要根据实际情况调整)
#ifndef SYSY_MEM2REG_H
#define SYSY_MEM2REG_H

#include <vector>
#include <unordered_map>
#include <unordered_set>
#include <stack>
#include <queue> // For computeIteratedDomFrontiers

// Include your IR and analysis headers
#include "IR.h"
#include "IRBuilder.h"
#include "SysYIRAnalyser.h"
#include "SysYIROptUtils.h"

namespace sysy {

class Mem2Reg {
private:
  Module* pModule;
  IRBuilder* pBuilder;
  ControlFlowAnalysis* controlFlowAnalysis;
  ActiveVarAnalysis* activeVarAnalysis;
  DataFlowAnalysisUtils dataFlowAnalysisUtils; // If this is part of Mem2Reg or an external helper

public:
  Mem2Reg(Module* module, IRBuilder* builder, ControlFlowAnalysis* cfa, ActiveVarAnalysis* ava)
    : pModule(module), pBuilder(builder), controlFlowAnalysis(cfa), activeVarAnalysis(ava) {}
    // Constructor initializes members
  void run();

  // --- 新增的私有成员变量和方法，用于SSA转换上下文 ---
  // 这是核心，用于存储 SSA 转换过程中的状态
  std::vector<AllocaInst*> currentFunctionAllocas; // 当前函数中所有可提升的 alloca
  // alloca -> set of BasicBlocks where it's defined (stored into)
  std::unordered_map<AllocaInst*, std::unordered_set<BasicBlock*>> allocaDefsBlock; 
  // alloca -> set of BasicBlocks where it's used (loaded from)
  std::unordered_map<AllocaInst*, std::unordered_set<BasicBlock*>> allocaUsesBlock; 

  // BasicBlock -> Map of (PhiInst, Original AllocaInst)
  // 用于在 rename 阶段通过 phi 指令找到它代表的原始 alloca
  std::unordered_map<BasicBlock*, std::unordered_map<PhiInst*, AllocaInst*>> phiMap;
  std::vector<PhiInst*> allPhiInstructions; // 收集所有创建的 Phi 指令以便后续简化和清理

  // --- 核心 SSA 转换辅助函数 ---
  // 计算给定定义块集合的迭代支配边界
  std::unordered_set<BasicBlock*> computeIteratedDomFrontiers(const std::unordered_set<BasicBlock*>& blocks);

  // 分析一个 alloca 的所有 uses，填充 allocaDefsBlock 和 allocaUsesBlock
  void allocaAnalysis(AllocaInst* alloca);

  // 判断一个 alloca 是否可以被提升为寄存器 (无地址逃逸，标量类型)
  bool is_promoted(AllocaInst* alloca);

  // 在迭代支配边界处插入 Phi 指令
  void insertPhiNodes(Function* func);

  // 递归地重命名基本块中的变量并填充 Phi 指令
  // 这里的 `count` 和 `stacks` 是临时的，用于 DFS 过程中传递状态
  void renameBlock(BasicBlock* block, 
                      std::unordered_map<AllocaInst*, Value*>& currentIncomings,
                      std::unordered_set<BasicBlock*>& visitedBlocks); // 修改为传递 map 和 set

  // 简化冗余的 Phi 指令 (当所有输入都相同时)
  void simplifyphi(PhiInst* phi);

  // 获取前驱块在后继块前驱列表中的索引，用于 Phi 指令入边
  int getPredIndex(BasicBlock* pred, BasicBlock* succ);

  // --- Mem2Reg 的主要工作流函数 ---
  // 对单个函数执行内存到寄存器的提升
  bool promoteMemoryToRegisters(Function* func);

};

} // namespace sysy

#endif // SYSY_MEM2REG_H