#include "ir/PassManager.h"
#include <iostream>
#include <unordered_set>
#include <queue>
#include <vector>

namespace ir {

bool RunDCE(Function* func) {
  std::unordered_set<Instruction*> live_instructions;
  std::queue<Instruction*> worklist;

  // 1. Mark inherently live instructions
  for (const auto& bbPtr : func->GetBlocks()) {
    for (const auto& instPtr : bbPtr->GetInstructions()) {
      auto* inst = instPtr.get();
      bool inherently_live = false;

      switch (inst->GetOpcode()) {
        case Opcode::Ret:
        case Opcode::Br:
        case Opcode::Store:
        case Opcode::Call:
          inherently_live = true;
          break;
        default:
          break;
      }

      if (inherently_live) {
        live_instructions.insert(inst);
        worklist.push(inst);
      }
    }
  }

  // 2. Propagate liveness along the def-use chains
  while (!worklist.empty()) {
    auto* inst = worklist.front();
    worklist.pop();

    for (size_t i = 0; i < inst->GetNumOperands(); ++i) {
      auto* operand = inst->GetOperand(i);
      if (auto* op_inst = dynamic_cast<Instruction*>(operand)) {
        if (live_instructions.find(op_inst) == live_instructions.end()) {
          live_instructions.insert(op_inst);
          worklist.push(op_inst);
        }
      }
    }
  }

  // 3. Sweep dead instructions
  bool changed = false;
  for (const auto& bbPtr : func->GetBlocks()) {
    std::vector<Instruction*> dead_instructions;
    for (const auto& instPtr : bbPtr->GetInstructions()) {
      auto* inst = instPtr.get();
      if (live_instructions.find(inst) == live_instructions.end()) {
        dead_instructions.push_back(inst);
      }
    }

    if (!dead_instructions.empty()) {
      changed = true;
      for (auto* inst : dead_instructions) {
        bbPtr->EraseInstruction(inst);
      }
    }
  }

  return changed;
}

} // namespace ir