// Mem2Reg（SSA 构造）：
// - 将局部变量的 alloca/load/store 提升为 SSA 形式
// - 使用标准的支配边界 + PHI 节点插入算法

#include "include/ir/IR.h"

#include <algorithm>
#include <iostream>
#include <map>
#include <memory>
#include <set>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace ir {

namespace {

constexpr bool kDebugMem2Reg = false;

// 计算基本块的后继
std::vector<BasicBlock*> GetSuccessors(BasicBlock* bb) {
  std::vector<BasicBlock*> succs;
  if (!bb) return succs;
  
  auto& insts = bb->GetInstructions();
  if (insts.empty()) return succs;
  
  auto* term = insts.back().get();
  if (!term) return succs;
  
  switch (term->GetOpcode()) {
    case Opcode::Br: {
      auto* br = static_cast<BranchInst*>(term);
      succs.push_back(br->GetTarget());
      break;
    }
    case Opcode::CondBr: {
      auto* cbr = static_cast<CondBranchInst*>(term);
      succs.push_back(cbr->GetTrueTarget());
      succs.push_back(cbr->GetFalseTarget());
      break;
    }
    case Opcode::Ret:
      break;
    default:
      break;
  }
  
  return succs;
}

// 计算基本块的前驱（通过扫描所有块的终止指令，不依赖缓存）
std::vector<BasicBlock*> GetPredecessors(BasicBlock* bb, Function* func = nullptr) {
  std::vector<BasicBlock*> preds;
  if (!bb || !func) return preds;
  
  for (auto& block : func->GetBlocks()) {
    auto succs = GetSuccessors(block.get());
    for (auto* succ : succs) {
      if (succ == bb) {
        preds.push_back(block.get());
        break;
      }
    }
  }
  
  return preds;
}

// 支配树相关信息
struct DominatorInfo {
  std::unordered_map<BasicBlock*, BasicBlock*> idom;
  std::unordered_map<BasicBlock*, std::vector<BasicBlock*>> dom_tree;
  BasicBlock* entry = nullptr;
};

// 计算两个块的最近公共支配者
BasicBlock* Intersect(BasicBlock* b1, BasicBlock* b2,
                      const std::unordered_map<BasicBlock*, BasicBlock*>& idom) {
  auto* finger1 = b1;
  auto* finger2 = b2;
  
  while (finger1 != finger2) {
    while (finger1 && idom.count(finger1) && idom.at(finger1) && finger1 != idom.at(finger1)) {
      finger1 = idom.at(finger1);
    }
    while (finger2 && idom.count(finger2) && idom.at(finger2) && finger2 != idom.at(finger2)) {
      finger2 = idom.at(finger2);
    }
    
    if (finger1 == finger2) break;
    
    if (finger1 && idom.count(finger1) && idom.at(finger1)) finger1 = idom.at(finger1);
    if (finger2 && idom.count(finger2) && idom.at(finger2)) finger2 = idom.at(finger2);
  }
  
  return finger1;
}

// 预先计算所有块的支配者
DominatorInfo ComputeDominators(Function* func) {
  DominatorInfo dom_info;
  
  auto& blocks = func->GetBlocks();
  if (blocks.empty()) return dom_info;
  
  dom_info.entry = func->GetEntry();
  
  if (kDebugMem2Reg) {
    std::cerr << "[Mem2Reg] Computing dominators for function: " << func->GetName() << std::endl;
    std::cerr << "[Mem2Reg] Entry block: " << dom_info.entry->GetName() << std::endl;
    std::cerr << "[Mem2Reg] Total blocks: " << blocks.size() << std::endl;
    
    // 调试前驱信息
    for (auto& bb : blocks) {
      auto preds = GetPredecessors(bb.get(), func);
      std::cerr << "  Block " << bb->GetName() << " has " << preds.size() << " predecessors: {";
      bool first = true;
      for (auto* pred : preds) {
        if (!first) std::cerr << ", ";
        std::cerr << pred->GetName();
        first = false;
      }
      std::cerr << "}" << std::endl;
    }
  }
  
  // 使用标准的迭代数据流算法计算支配者
  // 初始化：entry的idom是自己，其他都是nullptr
  for (auto& bb : blocks) {
    dom_info.idom[bb.get()] = nullptr;
  }
  dom_info.idom[dom_info.entry] = dom_info.entry;
  
  // 迭代直到收敛
  bool changed = true;
  int iteration = 0;
  while (changed) {
    changed = false;
    ++iteration;
    
    if (iteration > 1000) {
      std::cerr << "[Mem2Reg] ERROR: Dominator computation did not converge after 1000 iterations!" << std::endl;
      break;
    }
    
    for (auto& bb : blocks) {
      if (bb.get() == dom_info.entry) continue;
      
      auto preds = GetPredecessors(bb.get(), func);
      if (preds.empty()) continue;
      
      // 找到第一个已经有idom的前驱
      BasicBlock* new_idom = nullptr;
      for (auto* pred : preds) {
        if (dom_info.idom[pred] != nullptr) {
          new_idom = pred;
          break;
        }
      }
      
      if (!new_idom) continue;
      
      // 与其他有idom的前驱求交集
      for (auto* pred : preds) {
        if (pred == new_idom) continue;
        if (dom_info.idom[pred] == nullptr) continue;
        
        new_idom = Intersect(new_idom, pred, dom_info.idom);
        if (!new_idom) break;
      }
      
      if (new_idom && dom_info.idom[bb.get()] != new_idom) {
        dom_info.idom[bb.get()] = new_idom;
        changed = true;
      }
    }
  }
  
  if (kDebugMem2Reg) {
    std::cerr << "[Mem2Reg] Dominator computation completed after " << iteration << " iterations" << std::endl;
  }
  
  // 构建支配树
  for (auto& bb : blocks) {
    if (bb.get() == dom_info.entry) continue;
    auto* idom = dom_info.idom[bb.get()];
    if (idom && idom != bb.get()) {
      dom_info.dom_tree[idom].push_back(bb.get());
    }
  }
  
  if (kDebugMem2Reg) {
    std::cerr << "[Mem2Reg] Dominator tree:" << std::endl;
    int non_null_count = 0;
    for (auto& bb : blocks) {
      auto* idom = dom_info.idom[bb.get()];
      if (idom) non_null_count++;
      std::cerr << "  " << bb->GetName() << " dominated by: " 
                << (idom ? idom->GetName() : "null") << std::endl;
    }
    std::cerr << "[Mem2Reg] Blocks with non-null idom: " << non_null_count << "/" << blocks.size() << std::endl;
  }
  
  return dom_info;
}

// 计算支配边界
std::unordered_map<BasicBlock*, std::set<BasicBlock*>> ComputeDominanceFrontier(
    Function* func, const DominatorInfo& dom_info) {
  std::unordered_map<BasicBlock*, std::set<BasicBlock*>> df;
  
  auto& blocks = func->GetBlocks();
  for (auto& bb : blocks) {
    df[bb.get()] = std::set<BasicBlock*>();
  }
  
  for (auto& bb : blocks) {
    auto preds = GetPredecessors(bb.get(), func);
    if (preds.size() < 2) continue;
    
    auto* bb_idom = dom_info.idom.count(bb.get()) ? dom_info.idom.at(bb.get()) : nullptr;
    if (!bb_idom) continue;
    
    for (auto* pred : preds) {
      auto* runner = pred;
      while (runner != bb_idom) {
        df[runner].insert(bb.get());
        if (!runner || !dom_info.idom.count(runner) || !dom_info.idom.at(runner)) {
          break;
        }
        auto* next = dom_info.idom.at(runner);
        if (next == runner) break;
        runner = next;
      }
    }
  }
  
  if (kDebugMem2Reg) {
    std::cerr << "[Mem2Reg] Dominance frontiers:" << std::endl;
    for (auto& bb : blocks) {
      if (!df[bb.get()].empty()) {
        std::cerr << "  " << bb->GetName() << " -> {";
        bool first = true;
        for (auto* df_bb : df[bb.get()]) {
          if (!first) std::cerr << ", ";
          std::cerr << df_bb->GetName();
          first = false;
        }
        std::cerr << "}" << std::endl;
      }
    }
  }
  
  return df;
}

// 检查是否是可以 SSA 化的 alloca
bool IsPromotableAlloca(AllocaInst* alloca) {
  if (alloca->IsArrayAlloca()) {
    return false;
  }
  
  auto uses = alloca->GetUses();
  
  if (uses.empty()) {
    return false;  // 没有使用，不需要提升
  }
  
  bool has_load = false;
  
  for (auto& use : uses) {
    auto* user = use.GetUser();
    if (!user) return false;
    
    auto* inst = dynamic_cast<Instruction*>(user);
    if (!inst) return false;
    
    auto opcode = inst->GetOpcode();
    if (opcode != Opcode::Load && opcode != Opcode::Store) {
      return false;
    }
    
    if (opcode == Opcode::Store) {
      auto* store = static_cast<StoreInst*>(inst);
      if (store->GetPtr() != alloca) {
        return false;
      }
    }
    
    if (opcode == Opcode::Load) {
      has_load = true;
    }
  }
  
  return has_load;  // 只有有 load 才需要提升
}

std::vector<AllocaInst*> FindPromotableAllocas(Function* func) {
  std::vector<AllocaInst*> promotable;
  
  // 扫描所有基本块中的 allocas（不仅限于 entry 块）
  for (auto& bb : func->GetBlocks()) {
    for (auto& inst_ptr : bb->GetInstructions()) {
      auto* inst = inst_ptr.get();
      if (auto* alloca = dynamic_cast<AllocaInst*>(inst)) {
        if (IsPromotableAlloca(alloca)) {
          promotable.push_back(alloca);
        }
      }
    }
  }
  
  if (kDebugMem2Reg && !promotable.empty()) {
    std::cerr << "[Mem2Reg] Found " << promotable.size() << " promotable allocas in " 
              << func->GetName() << std::endl;
    for (auto* alloca : promotable) {
      std::cerr << "  - " << alloca->GetName() << std::endl;
    }
  }
  
  return promotable;
}

// 插入 PHI 节点
void InsertPhiNodes(Function* func,
                    const std::vector<AllocaInst*>& allocas,
                    const DominatorInfo& dom_info,
                    const std::unordered_map<BasicBlock*, std::set<BasicBlock*>>& df,
                    std::unordered_map<AllocaInst*, std::unordered_map<BasicBlock*, PhiInst*>>& phi_nodes,
                    std::unordered_set<AllocaInst*>& skip_allocas) {
  for (auto* alloca : allocas) {
    std::set<BasicBlock*> store_blocks;
    int store_count = 0;
    for (auto& bb : func->GetBlocks()) {
      for (auto& inst_ptr : bb->GetInstructions()) {
        auto* inst = inst_ptr.get();
        if (auto* store = dynamic_cast<StoreInst*>(inst)) {
          if (store->GetPtr() == alloca) {
            store_blocks.insert(bb.get());
            store_count++;
          }
        }
      }
    }
    
    if (kDebugMem2Reg && !store_blocks.empty()) {
      std::cerr << "[Mem2Reg] Store blocks for " << alloca->GetName() << ": {";
      bool first = true;
      for (auto* bb : store_blocks) {
        if (!first) std::cerr << ", ";
        std::cerr << bb->GetName();
        first = false;
      }
      std::cerr << "}" << std::endl;
    }
    
    std::set<BasicBlock*> phi_blocks;
    std::vector<BasicBlock*> worklist(store_blocks.begin(), store_blocks.end());
    std::set<BasicBlock*> visited(store_blocks.begin(), store_blocks.end());
    
    while (!worklist.empty()) {
      auto* bb = worklist.back();
      worklist.pop_back();
      
      auto df_it = df.find(bb);
      if (df_it == df.end()) continue;
      
      for (auto* df_block : df_it->second) {
        if (phi_blocks.insert(df_block).second && visited.insert(df_block).second) {
          worklist.push_back(df_block);
        }
      }
    }
    
    if (kDebugMem2Reg && !phi_blocks.empty()) {
      std::cerr << "[Mem2Reg] PHI blocks for " << alloca->GetName() << ": {";
      bool first = true;
      for (auto* bb : phi_blocks) {
        if (!first) std::cerr << ", ";
        std::cerr << bb->GetName();
        first = false;
      }
      std::cerr << "}" << std::endl;
    }
    
    if (store_count > 1 && !phi_blocks.empty()) {
      if (kDebugMem2Reg) {
        std::cerr << "[Mem2Reg] Alloca " << alloca->GetName() 
                  << " has multiple stores (store_count=" << store_count 
                  << ", phi_blocks=" << phi_blocks.size() << "), promoting anyway" << std::endl;
      }
    }
    
    for (auto* bb : phi_blocks) {
      auto& insts = const_cast<std::vector<std::unique_ptr<Instruction>>&>(bb->GetInstructions());
      if (insts.empty()) continue;
      
      // 使用 alloca 的元素类型作为 PHI 的类型
      auto* phi = bb->Prepend<PhiInst>(alloca->GetElementType(), "");
      phi->SetAlloca(alloca);
      phi_nodes[alloca][bb] = phi;
    }
  }
}

// 全局 PHI 计数器
int gPhiCounter = 0;

// 重命名
void Rename(BasicBlock* bb,
            const DominatorInfo& dom_info,
            AllocaInst* alloca,
            std::unordered_map<AllocaInst*, std::vector<Value*>>& value_stack,
            std::unordered_map<AllocaInst*, std::unordered_map<BasicBlock*, PhiInst*>>& phi_nodes,
            std::unordered_map<AllocaInst*, Value*>& initial_values) {
  // 处理当前块的 PHI 节点
  auto phi_it = phi_nodes.find(alloca);
  if (phi_it != phi_nodes.end()) {
    auto phi_block_it = phi_it->second.find(bb);
    if (phi_block_it != phi_it->second.end()) {
      auto* phi = phi_block_it->second;
      // 给 PHI 节点一个名字
      if (phi->GetName().empty()) {
        phi->SetName("%phi" + std::to_string(gPhiCounter++));
      }
      value_stack[alloca].push_back(phi);
      if (kDebugMem2Reg) {
        std::cerr << "[Mem2Reg] Rename: Push PHI " << phi->GetName() << " for " << alloca->GetName() 
                  << " in block " << bb->GetName() << std::endl;
      }
    }
  }
  
  // 先处理指令：store 推栈，load 替换
  for (auto& inst_ptr : bb->GetInstructions()) {
    auto* inst = inst_ptr.get();
    
    if (dynamic_cast<PhiInst*>(inst)) continue;
    
    if (auto* store = dynamic_cast<StoreInst*>(inst)) {
      if (store->GetPtr() == alloca) {
        auto* val = store->GetValue();
        value_stack[alloca].push_back(val);
        // 记录初始值：第一个 store 的值作为初始值（覆盖默认值 0）
        // 这样可以正确处理循环不变量的情况
        if (initial_values.count(alloca)) {
          auto* current_init = initial_values[alloca];
          // 如果当前初始值是默认值 0（ConstantInt 0），则更新为 store 的值
          if (auto* const_int = dynamic_cast<ConstantInt*>(current_init)) {
            if (const_int->GetValue() == 0) {
              initial_values[alloca] = val;
              if (kDebugMem2Reg) {
                std::cerr << "[Mem2Reg] Rename: Update initial value for " << alloca->GetName() 
                          << " from 0 to " << val->GetName() << " in block " << bb->GetName() << std::endl;
              }
            }
          }
        } else {
          initial_values[alloca] = val;
          if (kDebugMem2Reg) {
            std::cerr << "[Mem2Reg] Rename: Set initial value for " << alloca->GetName() 
                      << " to " << val->GetName() << " in block " << bb->GetName() << std::endl;
          }
        }
        if (kDebugMem2Reg) {
          std::cerr << "[Mem2Reg] Rename: Push store value " << val->GetName() << " for " << alloca->GetName()
                    << " in block " << bb->GetName() << std::endl;
        }
      }
    }
    
    if (auto* load = dynamic_cast<LoadInst*>(inst)) {
      if (load->GetPtr() == alloca) {
        auto stack_it = value_stack.find(alloca);
        if (kDebugMem2Reg) {
          std::cerr << "[Mem2Reg]   Checking load " << load->GetName() << " in " << bb->GetName() 
                    << ", stack size: " << (stack_it != value_stack.end() ? stack_it->second.size() : 0) << std::endl;
        }
        if (stack_it != value_stack.end() && !stack_it->second.empty()) {
          auto* current_val = stack_it->second.back();
          if (kDebugMem2Reg) {
            std::cerr << "[Mem2Reg] Rename: Replacing load " << load->GetName()
                      << " of " << alloca->GetName()
                      << " in block " << bb->GetName() 
                      << " with value: " << (current_val ? current_val->GetName() : "null") << std::endl;
            if (current_val && current_val->GetType()) {
              std::cerr << "[Mem2Reg]   Replacement value type: " << (int)current_val->GetType()->GetKind() << std::endl;
            }
          }
          // 类型安全检查：确保替换值的类型与 load 的结果类型完全匹配
          if (current_val && current_val->GetType()) {
            auto* load_type = load->GetType().get();
            auto* val_type = current_val->GetType().get();
            if (load_type && val_type && load_type == val_type) {
              load->ReplaceAllUsesWith(current_val);
            } else {
              if (kDebugMem2Reg) {
                std::cerr << "[Mem2Reg]   WARNING: Type mismatch! Load type " 
                          << (load_type ? std::to_string((int)load_type->GetKind()) : "null")
                          << " vs value type " 
                          << (val_type ? std::to_string((int)val_type->GetKind()) : "null") << std::endl;
              }
            }
          }
          } else {
            if (kDebugMem2Reg) {
              std::cerr << "[Mem2Reg] Rename: ERROR - Empty stack for load " << load->GetName()
                        << " in block " << bb->GetName() << std::endl;
            }
          }
      }
    }
  }
  
  // 处理完指令后，填充后继块的 PHI 节点
  auto successors = GetSuccessors(bb);
  for (auto* succ : successors) {
    auto succ_phi_it = phi_nodes.find(alloca);
    if (succ_phi_it != phi_nodes.end()) {
      auto succ_block_it = succ_phi_it->second.find(succ);
      if (succ_block_it != succ_phi_it->second.end()) {
        auto* phi = succ_block_it->second;
        // 如果栈为空，使用初始值
        Value* incoming_val = nullptr;
        if (!value_stack[alloca].empty()) {
          incoming_val = value_stack[alloca].back();
          if (kDebugMem2Reg && incoming_val == phi) {
            std::cerr << "[Mem2Reg] Rename: Self-referencing PHI in " 
                      << succ->GetName() << " for alloca " << alloca->GetName()
                      << " (variable unchanged on this path)" << std::endl;
          }
        } else if (initial_values.count(alloca)) {
          incoming_val = initial_values[alloca];
        }
        
        // 检查是否已经为这个前驱添加过条目
        bool already_added = false;
        for (size_t i = 1; i < phi->GetNumOperands(); i += 2) {
          if (phi->GetOperand(i) == bb) {
            already_added = true;
            break;
          }
        }
        
        if (!already_added && incoming_val) {
          phi->AddOperand(incoming_val);
          phi->AddOperand(bb);
          if (kDebugMem2Reg) {
            std::cerr << "[Mem2Reg] Rename: Adding operand to PHI in " 
                      << succ->GetName() << " from " << bb->GetName() 
                      << ", value: " << incoming_val->GetName() << std::endl;
          }
        } else if (!already_added && !incoming_val && kDebugMem2Reg) {
          std::cerr << "[Mem2Reg] Rename: WARNING - No value for PHI in " 
                    << succ->GetName() << " from " << bb->GetName() << std::endl;
        }
      }
    }
  }
  
  // 递归处理支配树中的子块
  auto dom_tree_it = dom_info.dom_tree.find(bb);
  if (dom_tree_it != dom_info.dom_tree.end()) {
    for (auto* child : dom_tree_it->second) {
      Rename(child, dom_info, alloca, value_stack, phi_nodes, initial_values);
    }
  }
  
  // 回溯
  for (auto rit = bb->GetInstructions().rbegin(); rit != bb->GetInstructions().rend(); ++rit) {
    auto* inst = rit->get();
    if (dynamic_cast<PhiInst*>(inst)) continue;
    
    if (auto* store = dynamic_cast<StoreInst*>(inst)) {
      if (store->GetPtr() == alloca) {
        auto& stack = value_stack[alloca];
        if (!stack.empty()) {
          stack.pop_back();
          if (kDebugMem2Reg) {
            std::cerr << "[Mem2Reg] Rename: Pop store value for " << alloca->GetName()
                      << " in block " << bb->GetName() << std::endl;
          }
        }
      }
    }
  }
  
  // 弹出 PHI
  auto phi_it2 = phi_nodes.find(alloca);
  if (phi_it2 != phi_nodes.end()) {
    auto phi_block_it = phi_it2->second.find(bb);
    if (phi_block_it != phi_it2->second.end()) {
      auto& stack = value_stack[alloca];
      if (!stack.empty()) {
        stack.pop_back();
        if (kDebugMem2Reg) {
          std::cerr << "[Mem2Reg] Rename: Pop PHI for " << alloca->GetName()
                    << " in block " << bb->GetName() << std::endl;
        }
      }
    }
  }
}

// 删除冗余指令
void RemoveRedundantInstructions(Function* func, const std::vector<AllocaInst*>& allocas) {
  std::set<Instruction*> to_delete;
  std::set<AllocaInst*> alloca_set(allocas.begin(), allocas.end());
  
  for (auto* alloca : allocas) {
    to_delete.insert(alloca);
  }
  
  for (auto& bb : func->GetBlocks()) {
    for (auto& inst_ptr : bb->GetInstructions()) {
      auto* inst = inst_ptr.get();
      if (auto* load = dynamic_cast<LoadInst*>(inst)) {
        auto* ptr = load->GetPtr();
        if (auto* alloca = dynamic_cast<AllocaInst*>(ptr)) {
          if (alloca_set.count(alloca) && load->GetUses().empty()) {
            to_delete.insert(load);
          }
        }
      } else if (auto* store = dynamic_cast<StoreInst*>(inst)) {
        auto* ptr = store->GetPtr();
        if (auto* alloca = dynamic_cast<AllocaInst*>(ptr)) {
          if (alloca_set.count(alloca)) {
            to_delete.insert(store);
          }
        }
      }
    }
  }
  
  for (auto& bb : func->GetBlocks()) {
    auto& insts = const_cast<std::vector<std::unique_ptr<Instruction>>&>(bb->GetInstructions());
    for (auto it = insts.begin(); it != insts.end();) {
      auto* inst = it->get();
      if (to_delete.count(inst) && !dynamic_cast<AllocaInst*>(inst)) {
        for (size_t i = 0; i < inst->GetNumOperands(); ++i) {
          if (auto* op = inst->GetOperand(i)) {
            auto& uses = const_cast<std::vector<Use>&>(op->GetUses());
            uses.erase(
              std::remove_if(uses.begin(), uses.end(),
                [inst](const Use& use) { return use.GetUser() == inst; }),
              uses.end());
          }
        }
        it = insts.erase(it);
      } else {
        ++it;
      }
    }
  }
  
  for (auto& bb : func->GetBlocks()) {
    auto& insts = const_cast<std::vector<std::unique_ptr<Instruction>>&>(bb->GetInstructions());
    for (auto it = insts.begin(); it != insts.end();) {
      if (dynamic_cast<AllocaInst*>(it->get()) && to_delete.count(it->get())) {
        auto* alloca = static_cast<AllocaInst*>(it->get());
        auto& uses = const_cast<std::vector<Use>&>(alloca->GetUses());
        uses.clear();
        it = insts.erase(it);
      } else {
        ++it;
      }
    }
  }
}

} // anonymous namespace

// Mem2Reg 主函数
void RunMem2Reg(Module& module) {
  if (kDebugMem2Reg) {
    std::cerr << "[Mem2Reg] === Starting Mem2Reg Pass ===" << std::endl;
  }
  
  for (auto& func_ptr : module.GetFunctions()) {
    auto* func = func_ptr.get();
    
    if (func->IsExternal()) continue;
    
    // 跳过过大的函数以避免编译超时
    if (func->GetBlocks().size() > 2000) {
      if (kDebugMem2Reg) {
        std::cerr << "[Mem2Reg] Skipping large function " << func->GetName() 
                  << " with " << func->GetBlocks().size() << " blocks" << std::endl;
      }
      continue;
    }
    
    if (kDebugMem2Reg) {
      std::cerr << "[Mem2Reg] Processing function: " << func->GetName() << std::endl;
    }
    
    auto promotable_allocas = FindPromotableAllocas(func);
    if (promotable_allocas.empty()) {
      if (kDebugMem2Reg) {
        std::cerr << "[Mem2Reg] No promotable allocas in " << func->GetName() << std::endl;
      }
      continue;
    }
    
    auto dom_info = ComputeDominators(func);
    
    auto df = ComputeDominanceFrontier(func, dom_info);
    
    std::unordered_map<AllocaInst*, std::unordered_map<BasicBlock*, PhiInst*>> phi_nodes;
    std::unordered_set<AllocaInst*> skip_allocas;
    InsertPhiNodes(func, promotable_allocas, dom_info, df, phi_nodes, skip_allocas);
    
    // 过滤掉需要跳过的 alloca
    std::vector<AllocaInst*> filtered_allocas;
    for (auto* alloca : promotable_allocas) {
      if (skip_allocas.find(alloca) == skip_allocas.end()) {
        filtered_allocas.push_back(alloca);
      }
    }
    
    if (filtered_allocas.empty()) {
      if (kDebugMem2Reg) {
        std::cerr << "[Mem2Reg] No promotable allocas left after filtering in " << func->GetName() << std::endl;
      }
      continue;
    }
    
    // 计算总 PHI 节点数量
    int total_phi_count = 0;
    for (auto& pair : phi_nodes) {
      total_phi_count += pair.second.size();
    }
    
    // 启发式：如果 PHI 节点数量过多，跳过该函数
    // PHI 节点在 llc -O0 下会生成 StoreStack 操作，可能导致性能下降
    // 阈值设置：基本块数量的 1/4，最小 10，最大 30
    int block_count = func->GetBlocks().size();
    int phi_threshold = std::max(50, block_count);
    if (total_phi_count > phi_threshold) {
      if (kDebugMem2Reg) {
        std::cerr << "[Mem2Reg] Skipping function " << func->GetName() 
                  << " due to too many PHI nodes: " << total_phi_count 
                  << " (threshold: " << phi_threshold << ")" << std::endl;
      }
      continue;
    }
    
    for (auto* alloca : filtered_allocas) {
      std::unordered_map<AllocaInst*, std::vector<Value*>> value_stack;
      std::unordered_map<AllocaInst*, Value*> initial_values;
      
      // 根据 alloca 的元素类型创建对应类型的默认常量 0
      auto* elem_type = alloca->GetElementType().get();
      Value* default_val = nullptr;
      if (elem_type->IsInt32()) {
        default_val = new ConstantInt(Type::GetInt32Type(), 0);
      } else if (elem_type->IsInt1()) {
        default_val = new ConstantInt(Type::GetInt1Type(), 0);
      } else if (elem_type->IsFloat32()) {
        default_val = new ConstantFloat(Type::GetFloat32Type(), 0.0);
      }
      if (default_val) {
        initial_values[alloca] = default_val;
      }
      
      if (kDebugMem2Reg) {
        std::cerr << "[Mem2Reg] === Renaming for alloca: " << alloca->GetName() << " ===" << std::endl;
      }
      Rename(func->GetEntry(), dom_info, alloca, value_stack, phi_nodes, initial_values);
    }
    
    RemoveRedundantInstructions(func, filtered_allocas);
    
    if (kDebugMem2Reg) {
      std::cerr << "[Mem2Reg] === Finished " << func->GetName() << " ===" << std::endl;
    }
  }
  
  if (kDebugMem2Reg) {
    std::cerr << "[Mem2Reg] === Mem2Reg Pass Complete ===" << std::endl;
  }
}

} // namespace ir