From c7470f2d45d1a85899e627da7adb198dc5e4354c Mon Sep 17 00:00:00 2001
From: ZhuangYumin <zhuangyumin@sjtu.edu.cn>
Date: Thu, 17 Oct 2024 15:20:27 +0000
Subject: [PATCH] built dom

---
 include/opt/cfg.h   |  2 ++
 src/opt/cfg.cpp     | 16 +++++++++++
 src/opt/mem2reg.cpp | 67 ++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/include/opt/cfg.h b/include/opt/cfg.h
index 430025f..3a6f47a 100644
--- a/include/opt/cfg.h
+++ b/include/opt/cfg.h
@@ -11,6 +11,7 @@ class CFGNodeType {
   std::vector<CFGNodeType *> successors, predecessors;
   BlockItem *corresponding_block;
   CFGNodeCollection dom;
+  bool visited;
   CFGNodeType *idom;
   std::vector<CFGNodeType *> successors_in_dom_tree;
   CFGNodeCollection dom_frontier;
@@ -27,5 +28,6 @@ class CFGType {
 CFGNodeCollection GetCFGNodeCollectionsIntersection(const CFGNodeCollection &a, const CFGNodeCollection &b);
 CFGNodeCollection GetCFGNodeCollectionsUnion(const CFGNodeCollection &a, const CFGNodeCollection &b);
 CFGNodeCollection GetCFGNodeCollectionsDifference(const CFGNodeCollection &a, const CFGNodeCollection &b);
+bool CFGNodeCollectionIsSame(const CFGNodeCollection &a, const CFGNodeCollection &b);
 
 CFGType BuildCFGForFunction(const std::shared_ptr<FunctionDefItem> &func);
\ No newline at end of file
diff --git a/src/opt/cfg.cpp b/src/opt/cfg.cpp
index 0450705..0c16ef9 100644
--- a/src/opt/cfg.cpp
+++ b/src/opt/cfg.cpp
@@ -71,6 +71,19 @@ CFGNodeCollection GetCFGNodeCollectionsDifference(const CFGNodeCollection &a, co
   return res;
 }
 
+bool CFGNodeCollectionIsSame(const CFGNodeCollection &a, const CFGNodeCollection &b) {
+  auto ita = a.begin();
+  auto itb = b.begin();
+  while (ita != a.end() && itb != b.end()) {
+    if (*ita != *itb) {
+      return false;
+    }
+    ita++;
+    itb++;
+  }
+  return ita == a.end() && itb == b.end();
+}
+
 CFGType BuildCFGForFunction(const std::shared_ptr<FunctionDefItem> &func) {
   CFGType res;
   if (!func->init_block) {
@@ -94,8 +107,11 @@ CFGType BuildCFGForFunction(const std::shared_ptr<FunctionDefItem> &func) {
       if (auto br = std::dynamic_pointer_cast<BRAction>(block->exit_action)) {
         node->successors.push_back(res.block_to_node[res.label_to_block[br->true_label_full]]);
         node->successors.push_back(res.block_to_node[res.label_to_block[br->false_label_full]]);
+        res.block_to_node[res.label_to_block[br->true_label_full]]->predecessors.push_back(node.get());
+        res.block_to_node[res.label_to_block[br->false_label_full]]->predecessors.push_back(node.get());
       } else if (auto uncond = std::dynamic_pointer_cast<UNConditionJMPAction>(block->exit_action)) {
         node->successors.push_back(res.block_to_node[res.label_to_block[uncond->label_full]]);
+        res.block_to_node[res.label_to_block[uncond->label_full]]->predecessors.push_back(node.get());
       } else if (auto ret = std::dynamic_pointer_cast<RETAction>(block->exit_action)) {
         // do nothing
       } else {
diff --git a/src/opt/mem2reg.cpp b/src/opt/mem2reg.cpp
index fbc8ca5..21f3afe 100644
--- a/src/opt/mem2reg.cpp
+++ b/src/opt/mem2reg.cpp
@@ -1,7 +1,72 @@
 #include "mem2reg.h"
+#include <queue>
 #include "cfg.h"
 
-void ConductMem2RegForFunction(const std::shared_ptr<FunctionDefItem> &func, const CFGType &cfg) {}
+void ConductMem2RegForFunction(const std::shared_ptr<FunctionDefItem> &func, const CFGType &cfg) {
+  bool all_dom_unchanged;
+  CFGNodeCollection all_nodes;
+  for (auto &node : cfg.nodes) {
+    all_nodes.push_back(node.get());
+  }
+  all_nodes.sort();
+  for (auto cfg_node : all_nodes) {
+    cfg_node->dom = all_nodes;
+  }
+  cfg.entry->dom = {cfg.entry};
+  do {
+    all_dom_unchanged = true;
+    for (auto cfg_node : all_nodes) {
+      cfg_node->visited = false;
+    }
+    std::queue<CFGNodeType *> Q;
+    Q.push(cfg.entry);
+    cfg.entry->visited = true;
+    while (Q.size() > 0) {
+      auto cur = Q.front();
+      Q.pop();
+      for (auto succ : cur->successors) {
+        if (!succ->visited) {
+          succ->visited = true;
+          Q.push(succ);
+        }
+      }
+      CFGNodeCollection new_dom = {cur};
+      if (cur->predecessors.size() > 0) {
+        CFGNodeCollection tmp = cur->predecessors[0]->dom;
+        for (size_t i = 1; i < cur->predecessors.size(); i++) {
+          tmp = GetCFGNodeCollectionsIntersection(tmp, cur->predecessors[i]->dom);
+        }
+        new_dom = GetCFGNodeCollectionsUnion(new_dom, tmp);
+      }
+      if (!CFGNodeCollectionIsSame(new_dom, cur->dom)) {
+        all_dom_unchanged = false;
+        cur->dom = new_dom;
+      }
+    }
+  } while (all_dom_unchanged);
+  for (auto node : cfg.nodes) {
+    if (node.get() == cfg.entry) continue;
+    for (auto potential_predecessor : node->predecessors) {
+      if (potential_predecessor->dom.size() + 1 == node->dom.size()) {
+        node->idom = potential_predecessor;
+        node->idom->successors_in_dom_tree.push_back(node.get());
+        break;
+      }
+    }
+  }
+  for (auto node : cfg.nodes) {
+    CFGNodeCollection is_frontier_of;
+    CFGNodeCollection tmp1 = {node.get()};
+    tmp1 = GetCFGNodeCollectionsDifference(node->dom, tmp1);
+    for (auto pred : node->predecessors) {
+      CFGNodeCollection tmp2 = GetCFGNodeCollectionsDifference(pred->dom, tmp1);
+      is_frontier_of = GetCFGNodeCollectionsUnion(is_frontier_of, tmp2);
+    }
+    for (auto frontier_node : is_frontier_of) {
+      frontier_node->dom_frontier.push_back(node.get());
+    }
+  }
+}
 std::shared_ptr<ModuleItem> Mem2Reg(std::shared_ptr<ModuleItem> src) {
   auto res = std::make_shared<ModuleItem>(*src);
   for (auto &func : res->function_defs) {