// CFG 简化：
// - 删除不可达块、合并空块、简化分支等
// - 改善 IR 结构，便于后续优化与后端生成

#include "ir/IR.h"

#include <algorithm>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace ir {

namespace {

// 从入口块开始进行 BFS/DFS，标记所有可达的基本块
std::unordered_set<BasicBlock*> FindReachableBlocks(Function* func) {
  std::unordered_set<BasicBlock*> reachable;
  std::vector<BasicBlock*> worklist;
  
  auto* entry = func->GetEntry();
  if (!entry) return reachable;
  
  reachable.insert(entry);
  worklist.push_back(entry);
  
  while (!worklist.empty()) {
    auto* bb = worklist.back();
    worklist.pop_back();
    
    // 遍历该块的所有后继块
    for (const auto& inst_ptr : bb->GetInstructions()) {
      auto* inst = inst_ptr.get();
      
      // 无条件分支：跳转到目标块
      if (auto* br = dynamic_cast<BranchInst*>(inst)) {
        auto* target = br->GetTarget();
        if (reachable.insert(target).second) {
          worklist.push_back(target);
        }
      }
      // 条件分支：跳转到 true 和 false 目标
      else if (auto* condbr = dynamic_cast<CondBranchInst*>(inst)) {
        auto* true_target = condbr->GetTrueTarget();
        auto* false_target = condbr->GetFalseTarget();
        if (reachable.insert(true_target).second) {
          worklist.push_back(true_target);
        }
        if (reachable.insert(false_target).second) {
          worklist.push_back(false_target);
        }
      }
      // 终止指令，没有后继
      if (inst->IsTerminator()) break;
    }
  }
  
  return reachable;
}

// 获取基本块的 terminator 指令
Instruction* GetTerminator(BasicBlock* bb) {
  const auto& insts = bb->GetInstructions();
  if (insts.empty()) return nullptr;
  return insts.back().get();
}

}  // namespace

void RunCFGSimplify(Module& module) {
  for (auto& func_ptr : module.GetFunctions()) {
    auto* func = func_ptr.get();
    if (func->IsExternal()) continue;
    
    auto& blocks = const_cast<std::vector<std::unique_ptr<BasicBlock>>&>(func->GetBlocks());
    
    bool changed = true;
    while (changed) {
      changed = false;
      
      auto reachable = FindReachableBlocks(func);
      
      std::unordered_set<BasicBlock*> unreachable;
      for (auto& bb : blocks) {
        if (reachable.find(bb.get()) == reachable.end()) {
          unreachable.insert(bb.get());
        }
      }
      
      for (auto& bb_ptr : blocks) {
        auto* bb = bb_ptr.get();
        if (unreachable.find(bb) != unreachable.end()) continue;
        
        std::vector<std::pair<PhiInst*, Value*>> phi_replacements;
        std::vector<PhiInst*> phi_to_delete;
        
        for (auto& inst_ptr : bb->GetInstructions()) {
          auto* phi = dynamic_cast<PhiInst*>(inst_ptr.get());
          if (!phi) break;
          
          Value* valid_val = nullptr;
          size_t valid_count = 0;
          
          for (size_t i = 0; i < phi->GetNumOperands(); i += 2) {
            auto* val = phi->GetOperand(i);
            auto* pred = static_cast<BasicBlock*>(phi->GetOperand(i + 1));
            if (unreachable.find(pred) == unreachable.end()) {
              valid_val = val;
              valid_count++;
            }
          }
          
          if (valid_count == 1 && valid_val) {
            phi_replacements.push_back({phi, valid_val});
          } else if (valid_count == 0) {
            phi_to_delete.push_back(phi);
          }
        }
        
        for (auto& [phi, val] : phi_replacements) {
          phi->ReplaceAllUsesWith(val);
          phi_to_delete.push_back(phi);
        }
        
        for (auto* phi : phi_to_delete) {
          for (size_t i = 0; i < phi->GetNumOperands(); ++i) {
            if (auto* op = phi->GetOperand(i)) {
              op->RemoveUse(phi, i);
            }
          }
        }
        
        auto& insts = const_cast<std::vector<std::unique_ptr<Instruction>>&>(bb->GetInstructions());
        auto new_end = std::remove_if(insts.begin(), insts.end(),
          [&phi_to_delete](const std::unique_ptr<Instruction>& inst_ptr) {
            return std::find(phi_to_delete.begin(), phi_to_delete.end(), inst_ptr.get()) != phi_to_delete.end();
          }
        );
        insts.erase(new_end, insts.end());
      }
      
      for (auto& bb_ptr : blocks) {
        if (unreachable.find(bb_ptr.get()) != unreachable.end()) {
          for (auto& inst_ptr : bb_ptr->GetInstructions()) {
            auto* inst = inst_ptr.get();
            for (size_t i = 0; i < inst->GetNumOperands(); ++i) {
              if (auto* op = inst->GetOperand(i)) {
                op->RemoveUse(inst, i);
              }
            }
          }
        }
      }
      
      size_t old_size = blocks.size();
      blocks.erase(
        std::remove_if(blocks.begin(), blocks.end(),
          [&reachable](const std::unique_ptr<BasicBlock>& bb_ptr) {
            return reachable.find(bb_ptr.get()) == reachable.end();
          }
        ),
        blocks.end()
      );
      if (blocks.size() != old_size) {
        changed = true;
      }
      
      for (auto& bb_ptr : blocks) {
        auto* bb = bb_ptr.get();
        auto* term = GetTerminator(bb);
        if (!term) continue;
        
        if (auto* condbr = dynamic_cast<CondBranchInst*>(term)) {
          auto* cond = condbr->GetCond();
          
          if (auto* const_int = dynamic_cast<ConstantInt*>(cond)) {
            auto* true_target = condbr->GetTrueTarget();
            auto* false_target = condbr->GetFalseTarget();
            
            auto* target = (const_int->GetValue() != 0) ? true_target : false_target;
            
            bb->RemoveInstruction(condbr);
            bb->Append<BranchInst>(Type::GetVoidType(), target);
            changed = true;
          }
        }
      }
    }
    
    for (auto& bb_ptr : blocks) {
      auto* bb = bb_ptr.get();
      const auto& preds = bb->GetPredecessors();
      
      if (preds.size() == 1) {
        std::vector<std::pair<PhiInst*, Value*>> phi_replacements;
        
        for (auto& inst_ptr : bb->GetInstructions()) {
          auto* phi = dynamic_cast<PhiInst*>(inst_ptr.get());
          if (!phi) break;
          
          if (phi->GetNumOperands() >= 2) {
            Value* incoming_val = phi->GetOperand(0);
            phi_replacements.push_back({phi, incoming_val});
          }
        }
        
        for (auto& [phi, val] : phi_replacements) {
          phi->ReplaceAllUsesWith(val);
        }
        
        std::vector<PhiInst*> phis_to_clean;
        for (auto& inst_ptr : bb->GetInstructions()) {
          auto* phi = dynamic_cast<PhiInst*>(inst_ptr.get());
          if (!phi) break;
          phis_to_clean.push_back(phi);
        }
        for (auto* phi : phis_to_clean) {
          for (size_t i = 0; i < phi->GetNumOperands(); ++i) {
            if (auto* op = phi->GetOperand(i)) {
              op->RemoveUse(phi, i);
            }
          }
        }
        
        auto& insts = const_cast<std::vector<std::unique_ptr<Instruction>>&>(bb->GetInstructions());
        auto new_end = std::remove_if(insts.begin(), insts.end(),
          [](const std::unique_ptr<Instruction>& inst_ptr) {
            return dynamic_cast<const PhiInst*>(inst_ptr.get()) != nullptr;
          }
        );
        insts.erase(new_end, insts.end());
      }
    }
  }
}

}  // namespace ir