#include "Mem2Reg.h" // 包含 Mem2Reg 遍的头文件
#include "Dom.h"     // 包含支配树分析的头文件
#include "Liveness.h"
#include "IR.h"      // 包含 IR 相关的定义
#include "SysYIROptUtils.h"
#include <cassert>   // 用于断言
#include <iostream>  // 用于调试输出

namespace sysy {

void *Mem2Reg::ID = (void *)&Mem2Reg::ID;

void Mem2RegContext::run(Function *func, AnalysisManager *AM) {
  if (func->getBasicBlocks().empty()) {
    return;
  }

  // 清空所有状态，确保每次运行都是新的状态
  promotableAllocas.clear();
  allocaToPhiMap.clear();
  allocaToValueStackMap.clear();
  allocaToStoresMap.clear();
  allocaToDefBlocksMap.clear();

  // 获取支配树分析结果
  dt = AM->getAnalysisResult<DominatorTree, DominatorTreeAnalysisPass>(func);
  assert(dt && "DominatorTreeAnalysisResult not available for Mem2Reg!");

  // --------------------------------------------------------------------
  // 阶段1: 识别可提升的 AllocaInst 并收集其 Store 指令
  // --------------------------------------------------------------------
  // 遍历函数入口块?中的所有指令，寻找 AllocaInst
  // 必须是要入口块的吗
  for (auto &inst : func->getEntryBlock()->getInstructions_Range()) {
    Value *allocainst = inst.get();
    if (auto alloca = dynamic_cast<AllocaInst *>(allocainst)) {
      if (isPromotableAlloca(alloca)) {
        promotableAllocas.push_back(alloca);
        collectStores(alloca); // 收集所有对该 alloca 的 store
      }
    }
  }

  // --------------------------------------------------------------------
  // 阶段2: 插入 Phi 指令
  // --------------------------------------------------------------------
  for (auto alloca : promotableAllocas) {
    // 为每个可提升的 alloca 插入 Phi 指令
    insertPhis(alloca, allocaToDefBlocksMap[alloca]);
  }

  // --------------------------------------------------------------------
  // 阶段3: 变量重命名
  // --------------------------------------------------------------------
  // 为每个可提升的 alloca 初始化其值栈
  for (auto alloca : promotableAllocas) {
    // 初始值通常是 undef 或 null，取决于 IR 类型系统
    UndefinedValue *undefValue =  UndefinedValue::get(alloca->getType()->as<PointerType>()->getBaseType());
    allocaToValueStackMap[alloca].push(undefValue); // 压入一个初始的“未定义”值
  }

  // 从入口基本块开始，对支配树进行 DFS 遍历，进行变量重命名
  renameVariables(nullptr, func->getEntryBlock()); // 第一个参数 alloca 在这里不使用，因为是递归入口点

  // --------------------------------------------------------------------
  // 阶段4: 清理
  // --------------------------------------------------------------------
  cleanup();
}

// 判断一个 AllocaInst 是否可以被提升到寄存器
bool Mem2RegContext::isPromotableAlloca(AllocaInst *alloca) {
  // 1. 必须是标量类型（非数组、非结构体）sysy不支持结构体
  if (alloca->getType()->as<PointerType>()->getBaseType()->isArray()) {
    return false;
  }

  // 2. 其所有用途都必须是 LoadInst 或 StoreInst
  // (或 GetElementPtrInst，但 GEP 的结果也必须只被 Load/Store 使用)
  for (auto use : alloca->getUses()) {
    auto user = use->getUser();
    if (!user)
      return false; // 用户无效

    if (dynamic_cast<LoadInst *>(user)) {
      // OK
    } else if (dynamic_cast<StoreInst *>(user)) {
      // OK
    } else if (auto gep = dynamic_cast<GetElementPtrInst *>(user)) {
      // 如果是 GetElementPtrInst (GEP)
      // 需要判断这个 GEP 是否代表了数组元素的访问，而非简单的指针操作
      // LLVM 的 mem2reg 通常不提升用于数组元素访问的 alloca。
      // 启发式判断：
      // 如果 GEP 有多个索引（例如 `getelementptr i32, i32* %ptr, i32 0, i32 %idx`），
      // 或者第一个索引（对于指针类型）不是常量 0，则很可能是数组访问。
      // 对于 `alloca i32* %a.param` (对应 `int a[]` 参数)，其 `allocatedType()` 是 `i32*`。
      // 访问 `a[i]` 会生成类似 `getelementptr i32, i32* %a.param, i32 %i` 的 GEP。
      // 这种 GEP 有两个操作数：基指针和索引。

      // 检查 GEP 的操作数数量和索引值
      // GEP 的操作数通常是：<base_pointer>, <index_1>, <index_2>, ...
      // 对于一个 `i32*` 类型的 `alloca`，如果它被 GEP 使用，那么 GEP 的第一个索引通常是 `0`
      // （表示解引用指针本身），后续索引才是数组元素的索引。
      // 如果 GEP 的操作数数量大于 2 (即 `base_ptr` 和 `index_0` 之外还有其他索引)，
      // 或者 `index_0` 不是常量 0，则它可能是一个复杂的数组访问。
      // 假设 `gep->getNumOperands()` 和 `gep->getOperand(idx)->getValue()`
      // 假设 `ConstantInt` 类用于表示常量整数值
      if (gep->getNumOperands() > 2) { // 如果有超过一个索引（除了基指针的第一个隐式索引）
        // std::cerr << "Mem2Reg: Not promotable (GEP with multiple indices): " << alloca->name() << std::endl;
        return false; // 复杂 GEP，通常表示数组或结构体字段访问
      }
      if (gep->getNumOperands() == 2) {                            // 只有基指针和一个索引
        Value *firstIndexVal = gep->getOperand(1); // 获取第一个索引值
        if (auto constInt = dynamic_cast<ConstantInteger *>(firstIndexVal)) {
          if (constInt->getInt() != 0) {
            // std::cerr << "Mem2Reg: Not promotable (GEP with non-zero first index): " << alloca->name() << std::endl;
            return false; // 索引不是0，表示访问数组的非第一个元素
          }
        } else {
          // std::cerr << "Mem2Reg: Not promotable (GEP with non-constant first index): " << alloca->name() <<
          // std::endl;
          return false; // 索引不是常量，表示动态数组访问
        }
      }

      // 此外，GEP 的结果也必须只被 LoadInst 或 StoreInst 使用
      for (auto gep_use : gep->getUses()) {
        auto gep_user = gep_use->getUser();
        if (!gep_user) {
          // std::cerr << "Mem2Reg: Not promotable (GEP result null user): " << alloca->name() << std::endl;
          return false;
        }
        if (!dynamic_cast<LoadInst *>(gep_user) && !dynamic_cast<StoreInst *>(gep_user)) {
          // std::cerr << "Mem2Reg: Not promotable (GEP result used by non-load/store): " << alloca->name() <<
          // std::endl;
          return false; // GEP 结果被其他指令使用，地址逃逸或复杂用途
        }
      }
    } else {
      // 其他类型的用户，如 CallInst (如果地址逃逸)，则不能提升
      return false;
    }
  }
  // 3. 不能是 volatile 内存访问 (假设 AllocaInst 有 isVolatile() 方法)
  // if (alloca->isVolatile()) return false; // 如果有这样的属性

  return true;
}

// 收集所有对给定 AllocaInst 进行存储的 StoreInst
void Mem2RegContext::collectStores(AllocaInst *alloca) {
  // 遍历 alloca 的所有用途
  for (auto use : alloca->getUses()) {
    auto user = use->getUser();
    if (!user)
      continue;

    if (auto storeInst = dynamic_cast<StoreInst *>(user)) {
      allocaToStoresMap[alloca].insert(storeInst);
      allocaToDefBlocksMap[alloca].insert(storeInst->getParent());
    } else if (auto gep = dynamic_cast<GetElementPtrInst *>(user)) {
      // 如果是 GEP，递归收集其下游的 store
      for (auto gep_use : gep->getUses()) {
        if (auto gep_store = dynamic_cast<StoreInst *>(gep_use->getUser())) {
          allocaToStoresMap[alloca].insert(gep_store);
          allocaToDefBlocksMap[alloca].insert(gep_store->getParent());
        }
      }
    }
  }
}

// 为给定的 AllocaInst 插入必要的 Phi 指令
void Mem2RegContext::insertPhis(AllocaInst *alloca, const std::unordered_set<BasicBlock *> &defBlocks) {
  std::queue<BasicBlock *> workQueue;
  std::unordered_set<BasicBlock *> phiHasBeenInserted; // 记录已插入 Phi 的基本块

  // 将所有定义块加入工作队列
  for (auto bb : defBlocks) {
    workQueue.push(bb);
  }

  while (!workQueue.empty()) {
    BasicBlock *currentDefBlock = workQueue.front();
    workQueue.pop();

    // 遍历当前定义块的支配边界 (Dominance Frontier)
    const std::set<BasicBlock *> *frontierBlocks = dt->getDominanceFrontier(currentDefBlock);
    for (auto frontierBlock : *frontierBlocks) {
      // 如果该支配边界块还没有为当前 alloca 插入 Phi 指令
      if (phiHasBeenInserted.find(frontierBlock) == phiHasBeenInserted.end()) {
        // 在支配边界块的开头插入一个新的 Phi 指令
        // Phi 指令的类型与 alloca 的类型指向的类型相同

        builder->setPosition(frontierBlock, frontierBlock->begin()); // 设置插入位置为基本块开头
        PhiInst *phiInst = builder->createPhiInst(alloca->getAllocatedType(), {}, {}, "");
        
        allocaToPhiMap[alloca][frontierBlock] = phiInst; // 记录 Phi 指令

        phiHasBeenInserted.insert(frontierBlock); // 标记已插入 Phi

        // 如果这个支配边界块本身也是一个定义块（即使没有 store，但插入了 Phi），
        // 那么它的支配边界也可能需要插入 Phi
        // 例如一个xx型的cfg，如果在第一个交叉处插入phi节点，那么第二个交叉处可能也需要插入phi
        workQueue.push(frontierBlock);
      }
    }
  }
}

// 对支配树进行深度优先遍历，重命名变量并替换 load/store 指令
void Mem2RegContext::renameVariables(AllocaInst *currentAlloca, BasicBlock *currentBB) {
  // 维护一个局部栈，用于存储当前基本块中为 Phi 和 Store 创建的 SSA 值，以便在退出时弹出
  std::stack<Value *> localStackPushed;

  // --------------------------------------------------------------------
  // 处理当前基本块的指令
  // --------------------------------------------------------------------
  for (auto instIter = currentBB->getInstructions().begin(); instIter != currentBB->getInstructions().end();) {
    Instruction *inst = instIter->get();
    bool instDeleted = false;

    // 处理 Phi 指令 (如果是当前 alloca 的 Phi)
    if (auto phiInst = dynamic_cast<PhiInst *>(inst)) {
      // 检查这个 Phi 是否是为某个可提升的 alloca 插入的
      for (auto alloca : promotableAllocas) {
        if (allocaToPhiMap[alloca].count(currentBB) && allocaToPhiMap[alloca][currentBB] == phiInst) {
          // 为 Phi 指令的输出创建一个新的 SSA 值，并压入值栈
          allocaToValueStackMap[alloca].push(phiInst);
          localStackPushed.push(phiInst); // 记录以便弹出
          break;                          // 找到对应的 alloca，处理下一个指令
        }
      }
    }
    // 处理 LoadInst
    else if (auto loadInst = dynamic_cast<LoadInst *>(inst)) {
      // 检查这个 LoadInst 是否是为某个可提升的 alloca
      for (auto alloca : promotableAllocas) {
        if (loadInst->getPointer() == alloca) { 
          // loadInst->getPointer() 返回 AllocaInst*
          // 将 LoadInst 的所有用途替换为当前 alloca 值栈顶部的 SSA 值
          assert(!allocaToValueStackMap[alloca].empty() && "Value stack empty for alloca during load replacement!");
          if(DEBUG){
            std::cout << "Mem2Reg: Replacing load " << loadInst->getPointer()->getName() << " with SSA value." << std::endl;
          }
          loadInst->replaceAllUsesWith(allocaToValueStackMap[alloca].top());
          instIter = SysYIROptUtils::usedelete(instIter);
          instDeleted = true;
          // std::cerr << "Mem2Reg: Replaced load " << loadInst->name() << " with SSA value." << std::endl;
          break;
        }
      }
    }
    // 处理 StoreInst
    else if (auto storeInst = dynamic_cast<StoreInst *>(inst)) {
      // 检查这个 StoreInst 是否是为某个可提升的 alloca
      for (auto alloca : promotableAllocas) {
        if (storeInst->getPointer() == alloca) { 
          // 假设 storeInst->getPointer() 返回 AllocaInst*
          // 将 StoreInst 存储的值作为新的 SSA 值，压入值栈
          if(DEBUG){
            std::cout << "Mem2Reg: Replacing store to " << storeInst->getPointer()->getName() << " with SSA value." << std::endl;
          }
          allocaToValueStackMap[alloca].push(storeInst->getValue());
          localStackPushed.push(storeInst->getValue());          // 记录以便弹出
          instIter = SysYIROptUtils::usedelete(instIter);
          instDeleted = true;
          // std::cerr << "Mem2Reg: Replaced store to " << storeInst->ptr()->name() << " with SSA value." << std::endl;
          break;
        }
      }
    }

    if (!instDeleted) {
      ++instIter; // 如果指令没有被删除，移动到下一个
    }
  }

  // --------------------------------------------------------------------
  // 处理后继基本块的 Phi 指令参数
  // --------------------------------------------------------------------
  for (auto successorBB : currentBB->getSuccessors()) {
    if (!successorBB)
      continue;
    for (auto alloca : promotableAllocas) {
      // 如果后继基本块包含为当前 alloca 插入的 Phi 指令
      if (allocaToPhiMap[alloca].count(successorBB)) {
        auto phiInst = allocaToPhiMap[alloca][successorBB];
        // 为 Phi 指令添加来自当前基本块的参数
        // 参数值是当前 alloca 值栈顶部的 SSA 值
        assert(!allocaToValueStackMap[alloca].empty() && "Value stack empty for alloca when setting phi operand!");
        phiInst->addIncoming(allocaToValueStackMap[alloca].top(), currentBB);
      }
    }
  }

  // --------------------------------------------------------------------
  // 递归访问支配树的子节点
  // --------------------------------------------------------------------
  const std::set<BasicBlock *> *dominatedBlocks = dt->getDominatorTreeChildren(currentBB);
  if(dominatedBlocks){
    for (auto dominatedBB : *dominatedBlocks) {
      if (dominatedBB) {
        if(DEBUG){
          std::cout << "Mem2Reg: Recursively renaming variables in dominated block: " << dominatedBB->getName() << std::endl;
        }
        renameVariables(currentAlloca, dominatedBB);
      }
    }
  }
  

  // --------------------------------------------------------------------
  // 退出基本块时，弹出在此块中压入值栈的 SSA 值
  // --------------------------------------------------------------------
  while (!localStackPushed.empty()) {
    Value *val = localStackPushed.top();
    localStackPushed.pop();
    // 找到是哪个 alloca 对应的栈
    for (auto alloca : promotableAllocas) {
      if (!allocaToValueStackMap[alloca].empty() && allocaToValueStackMap[alloca].top() == val) {
        allocaToValueStackMap[alloca].pop();
        break;
      }
    }
  }
}

// 删除所有原始的 AllocaInst、LoadInst 和 StoreInst
void Mem2RegContext::cleanup() {
  for (auto alloca : promotableAllocas) {
    if (alloca && alloca->getParent()) {
      // 删除 alloca 指令本身
      SysYIROptUtils::usedelete(alloca);
      
      // std::cerr << "Mem2Reg: Deleted alloca " << alloca->name() << std::endl;
    }
  }
  // LoadInst 和 StoreInst 已经在 renameVariables 阶段被删除了
}

// Mem2Reg 遍的 runOnFunction 方法实现
bool Mem2Reg::runOnFunction(Function *F, AnalysisManager &AM) {
  // 记录初始的指令数量，用于判断优化是否发生了改变
  size_t initial_inst_count = 0;
  for (auto &bb : F->getBasicBlocks()) {
    initial_inst_count += bb->getInstructions().size();
  }

  Mem2RegContext ctx(builder);
  ctx.run(F, &AM); // 运行 Mem2Reg 优化

  // 运行优化后，再次计算指令数量
  size_t final_inst_count = 0;
  for (auto &bb : F->getBasicBlocks()) {
    final_inst_count += bb->getInstructions().size();
  }

  // 如果指令数量发生变化（通常是减少，因为 load/store 被删除，phi 被添加），说明 IR 被修改了
  // TODO：不保险，后续修改为更精确的判断
  // 直接在添加和删除指令时维护changed值
  bool changed = (initial_inst_count != final_inst_count);

  // 如果 IR 被修改，则使相关的分析结果失效
  if (changed) {
    // Mem2Reg 会显著改变 IR 结构，特别是数据流和控制流（通过 Phi）。
    // 这会使几乎所有数据流分析和部分控制流分析失效。
    // AM.invalidateAnalysis(&DominatorTreeAnalysisPass::ID, F); // 支配树可能间接改变（如果基本块被删除）
    // AM.invalidateAnalysis(&LivenessAnalysisPass::ID, F); // 活跃性分析肯定失效
    // AM.invalidateAnalysis(&LoopInfoAnalysisPass::ID, F); // 循环信息可能失效
    // AM.invalidateAnalysis(&SideEffectInfoAnalysisPass::ID); // 副作用分析可能失效（如果 Alloca/Load/Store
    // 被替换为寄存器）
    // ... 其他数据流分析，如到达定义、可用表达式等，也应失效
  }
  return changed;
}

// 声明Mem2Reg遍的分析依赖和失效信息
void Mem2Reg::getAnalysisUsage(std::set<void *> &analysisDependencies, std::set<void *> &analysisInvalidations) const {
  // Mem2Reg 强烈依赖于支配树分析来插入 Phi 指令
  analysisDependencies.insert(&DominatorTreeAnalysisPass::ID); // 假设 DominatorTreeAnalysisPass 的 ID

  // Mem2Reg 会删除 Alloca/Load/Store 指令，插入 Phi 指令，这会大幅改变 IR 结构。
  // 因此，它会使许多分析结果失效。
  analysisInvalidations.insert(&DominatorTreeAnalysisPass::ID); // 支配树可能受影响
  analysisInvalidations.insert(&LivenessAnalysisPass::ID); // 活跃性分析肯定失效
  // analysisInvalidations.insert(&LoopInfoAnalysisPass::ID); // 循环信息可能失效
  // analysisInvalidations.insert(&SideEffectInfoAnalysisPass::ID); // 副作用分析可能失效
  // 其他所有依赖于数据流或 IR 结构的分析都可能失效。
}

} // namespace sysy