feat(DeadCodeElimination - While loop and Function wrapper): Add WhileStatementElimination and UnusedFunctionElimination in DCE pass

maldoca/js/ir/transforms/dead_code_elimination/BUILD

@@ -29,6 +29,8 @@ cc_library(
     hdrs = ["pass.h"],
     deps = [
         "//maldoca/js/ir",
+        "//maldoca/js/ast",
+        "@abseil-cpp//absl/container:flat_hash_map",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",

maldoca/js/ir/transforms/dead_code_elimination/pass.cc

@@ -14,13 +14,19 @@
 
 #include "maldoca/js/ir/transforms/dead_code_elimination/pass.h"
 
+#include <tuple>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
 #include "llvm/ADT/STLExtras.h"
+#include "maldoca/js/ast/ast.generated.h"
+#include "maldoca/js/ir/ir.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"
-#include "maldoca/js/ir/ir.h"
 
 namespace maldoca {
+
 void IfStatementElimination(mlir::Operation* root_op) {
   root_op->walk([&](JshirIfStatementOp op) {
     auto condition = op.getTest().getDefiningOp<JsirBooleanLiteralOp>();
@@ -56,7 +62,90 @@ void IfStatementElimination(mlir::Operation* root_op) {
   });
 }
 
+void WhileStatementElimination(mlir::Operation* root_op) {
+  root_op->walk([&](JshirWhileStatementOp op) {
+    mlir::Region& test_region = op.getTest();
+    if (test_region.empty()) {
+      return;
+    }
+
+    auto expr_region_end_op =
+        llvm::dyn_cast<JsirExprRegionEndOp>(&test_region.front().back());
+    if (expr_region_end_op == nullptr) {
+      return;
+    }
+
+    auto condition_op =
+        expr_region_end_op.getOperand().getDefiningOp<JsirBooleanLiteralOp>();
+    if (condition_op == nullptr) {
+      return;
+    }
+
+    if (!condition_op.getValue()) {
+      // Condition is constantly false, eliminate the entire loop.
+      op.erase();
+    }
+
+    // If condition is true, it's an infinite loop. For now, we leave it.
+  });
+}
+
+struct SymbolInfo {
+  std::vector<mlir::Operation*> definitions;
+  std::vector<mlir::Operation*> references;
+};
+
+bool operator==(const JsSymbolId& lhs, const JsSymbolId& rhs) {
+  return std::forward_as_tuple(lhs.name(), lhs.def_scope_uid()) ==
+         std::forward_as_tuple(rhs.name(), rhs.def_scope_uid());
+}
+
+template <typename H>
+H AbslHashValue(H h, const JsSymbolId& m) {
+  return H::combine(std::move(h), m.name(), m.def_scope_uid());
+}
+
+JsSymbolId GetSymbolIdFromAttr(JsirSymbolIdAttr symbol_attr) {
+  std::string name = symbol_attr.getName().str();
+  std::optional<int64_t> scope_uid = symbol_attr.getDefScopeId();
+  return JsSymbolId{name, scope_uid};
+}
+
+void UnusedFunctionElimination(mlir::Operation* root_op) {
+  absl::flat_hash_map<JsSymbolId, SymbolInfo> symbol_infos;
+
+  root_op->walk([&](mlir::Operation* op) {
+    auto trivia = llvm::dyn_cast<JsirTriviaAttr>(op->getLoc());
+    if (trivia == nullptr) {
+      return;
+    }
+    JsirSymbolIdAttr symbol = trivia.getReferencedSymbol();
+    if (symbol != nullptr) {
+      symbol_infos[GetSymbolIdFromAttr(symbol)].references.push_back(op);
+    }
+
+    llvm::ArrayRef<JsirSymbolIdAttr> mlir_defined_symbols =
+        trivia.getDefinedSymbols();
+    for (JsirSymbolIdAttr defined_symbol : mlir_defined_symbols) {
+      symbol_infos[GetSymbolIdFromAttr(defined_symbol)].definitions.push_back(
+          op);
+    }
+  });
+
+  for (const auto& [symbol, info] : symbol_infos) {
+    if (info.references.empty()) {
+      for (mlir::Operation* def_op : info.definitions) {
+        if (llvm::isa<JsirFunctionDeclarationOp>(def_op)) {
+          def_op->erase();
+        }
+      }
+    }
+  }
+}
+
 void DeadCodeElimination(mlir::Operation* root_op) {
   IfStatementElimination(root_op);
+  WhileStatementElimination(root_op);
+  UnusedFunctionElimination(root_op);
 }
 }  // namespace maldoca