Lab4: Implement basic scalar optimizations and lower Phi nodes to assembly

src/ir/analysis/DominatorTree.cpp

@@ -1,4 +1,192 @@
-// 支配树分析：
-// - 构建/查询 Dominator Tree 及相关关系
-// - 为 mem2reg、CFG 优化与循环分析提供基础能力
+#include "ir/PassManager.h"
+#include <algorithm>
+#include <iostream>
+#include <queue>
+#include <unordered_set>
 
+namespace ir {
+
+// Helper to rebuild CFG predecessors and successors.
+void RebuildCFG(Function* func) {
+  for (auto& bbPtr : func->GetBlocks()) {
+    bbPtr->ClearPredecessors();
+    bbPtr->ClearSuccessors();
+  }
+  for (auto& bbPtr : func->GetBlocks()) {
+    auto* bb = bbPtr.get();
+    const auto& insts = bb->GetInstructions();
+    if (insts.empty()) continue;
+    auto* term = insts.back().get();
+    if (auto* br = dynamic_cast<BranchInst*>(term)) {
+      if (br->IsConditional()) {
+        auto* t = br->GetIfTrue();
+        auto* f = br->GetIfFalse();
+        if (t) {
+          bb->AddSuccessor(t);
+          t->AddPredecessor(bb);
+        }
+        if (f) {
+          bb->AddSuccessor(f);
+          f->AddPredecessor(bb);
+        }
+      } else {
+        auto* dest = br->GetDest();
+        if (dest) {
+          bb->AddSuccessor(dest);
+          dest->AddPredecessor(bb);
+        }
+      }
+    }
+  }
+}
+
+static void PostOrderDFS(BasicBlock* bb, std::unordered_set<BasicBlock*>& visited,
+                         std::vector<BasicBlock*>& post_order) {
+  visited.insert(bb);
+  for (auto* succ : bb->GetSuccessors()) {
+    if (visited.find(succ) == visited.end()) {
+      PostOrderDFS(succ, visited, post_order);
+    }
+  }
+  post_order.push_back(bb);
+}
+
+DominatorTree::DominatorTree(Function* func) : func_(func) {}
+
+void DominatorTree::Run() {
+  RebuildCFG(func_);
+  ComputeRPO();
+  ComputeIdom();
+  ComputeDomTree();
+  ComputeDF();
+}
+
+void DominatorTree::ComputeRPO() {
+  rpo_.clear();
+  if (func_->GetBlocks().empty()) return;
+  std::unordered_set<BasicBlock*> visited;
+  std::vector<BasicBlock*> post_order;
+  PostOrderDFS(func_->GetEntry(), visited, post_order);
+  rpo_ = std::vector<BasicBlock*>(post_order.rbegin(), post_order.rend());
+}
+
+void DominatorTree::ComputeIdom() {
+  idom_.clear();
+  if (rpo_.empty()) return;
+
+  BasicBlock* entry = rpo_.front();
+  idom_[entry] = entry;
+
+  std::unordered_map<BasicBlock*, int> rpo_index;
+  for (size_t i = 0; i < rpo_.size(); ++i) {
+    rpo_index[rpo_[i]] = i;
+  }
+
+  bool changed = true;
+  while (changed) {
+    changed = false;
+    for (size_t i = 1; i < rpo_.size(); ++i) {
+      BasicBlock* b = rpo_[i];
+      BasicBlock* new_idom = nullptr;
+
+      // Find first predecessor with a defined idom
+      for (auto* pred : b->GetPredecessors()) {
+        if (idom_.find(pred) != idom_.end()) {
+          new_idom = pred;
+          break;
+        }
+      }
+
+      if (new_idom) {
+        for (auto* pred : b->GetPredecessors()) {
+          if (pred != new_idom && idom_.find(pred) != idom_.end()) {
+            // Intersect
+            auto* finger1 = pred;
+            auto* finger2 = new_idom;
+            while (finger1 != finger2) {
+              while (rpo_index.at(finger1) > rpo_index.at(finger2)) {
+                finger1 = idom_.at(finger1);
+              }
+              while (rpo_index.at(finger2) > rpo_index.at(finger1)) {
+                finger2 = idom_.at(finger2);
+              }
+            }
+            new_idom = finger1;
+          }
+        }
+
+        if (idom_.find(b) == idom_.end() || idom_[b] != new_idom) {
+          idom_[b] = new_idom;
+          changed = true;
+        }
+      }
+    }
+  }
+}
+
+void DominatorTree::ComputeDomTree() {
+  dom_tree_.clear();
+  for (auto* b : rpo_) {
+    dom_tree_[b] = {};
+  }
+  for (auto* b : rpo_) {
+    if (b != rpo_.front()) {
+      auto* parent = idom_[b];
+      dom_tree_[parent].push_back(b);
+    }
+  }
+}
+
+void DominatorTree::ComputeDF() {
+  df_.clear();
+  for (auto* b : rpo_) {
+    df_[b] = {};
+  }
+  for (auto* b : rpo_) {
+    if (b->GetPredecessors().size() >= 2) {
+      for (auto* pred : b->GetPredecessors()) {
+        auto* runner = pred;
+        auto* idom_b = idom_[b];
+        while (runner != idom_b) {
+          // If runner's df doesn't contain b already, add it
+          auto& runner_df = df_[runner];
+          if (std::find(runner_df.begin(), runner_df.end(), b) == runner_df.end()) {
+            runner_df.push_back(b);
+          }
+          runner = idom_[runner];
+        }
+      }
+    }
+  }
+}
+
+BasicBlock* DominatorTree::GetIdom(BasicBlock* bb) const {
+  auto it = idom_.find(bb);
+  return it != idom_.end() ? it->second : nullptr;
+}
+
+const std::vector<BasicBlock*>& DominatorTree::GetDominatedBlocks(BasicBlock* bb) const {
+  static const std::vector<BasicBlock*> empty;
+  auto it = dom_tree_.find(bb);
+  return it != dom_tree_.end() ? it->second : empty;
+}
+
+const std::vector<BasicBlock*>& DominatorTree::GetDominanceFrontier(BasicBlock* bb) const {
+  static const std::vector<BasicBlock*> empty;
+  auto it = df_.find(bb);
+  return it != df_.end() ? it->second : empty;
+}
+
+bool DominatorTree::Dominates(BasicBlock* a, BasicBlock* b) const {
+  if (a == b) return true;
+  auto* runner = b;
+  while (runner != rpo_.front()) {
+    auto it = idom_.find(runner);
+    if (it == idom_.end()) return false;
+    runner = it->second;
+    if (runner == a) return true;
+  }
+  return false;
+}
+
+} // namespace ir