Add unreachable block elimination optimization pass

CMakeLists.txt

@@ -65,6 +65,7 @@ if(BUILD_TESTING)
         ${CMAKE_CURRENT_SOURCE_DIR}/tests/parser_exact_test.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/tests/parser_syntax_error_files_test.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/tests/ir_constant_folding_test.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/tests/ir_unreachable_block_elimination_test.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/tests/symbol_table_test.cpp
     )
     target_link_libraries(minijava_tests PRIVATE GTest::gtest_main minijava_core)

src/ir/CFG.cpp

@@ -103,6 +103,54 @@ BBlock *CFG::addMethodRootBlock(const std::string &className,
     return ptr;
 }
 
+bool CFG::removeUnreachableBlocks() {
+    std::vector<BBlock *> stack;
+    std::unordered_set<BBlock *> reachable;
+    stack.reserve(methodRoots.size());
+
+    for (auto *root : methodRoots) {
+        if (root != nullptr) {
+            stack.push_back(root);
+        }
+    }
+
+    while (!stack.empty()) {
+        auto *block = stack.back();
+        stack.pop_back();
+
+        if (!reachable.insert(block).second) {
+            continue;
+        }
+
+        if (block->hasTrueBlock()) {
+            stack.push_back(block->getTrueBlock());
+        }
+        if (block->hasFalseBlock()) {
+            stack.push_back(block->getFalseBlock());
+        }
+    }
+
+    methodRoots.erase(
+        std::remove_if(methodRoots.begin(), methodRoots.end(),
+                       [&reachable](BBlock *root) {
+                           return root == nullptr ||
+                                  !reachable.contains(root);
+                       }),
+        methodRoots.end());
+
+    if (currentBlock != nullptr && !reachable.contains(currentBlock)) {
+        currentBlock = nullptr;
+    }
+
+    const auto old_size = allBlocks.size();
+    allBlocks.erase(std::remove_if(allBlocks.begin(), allBlocks.end(),
+                                   [&reachable](const auto &block) {
+                                       return !reachable.contains(block.get());
+                                   }),
+                    allBlocks.end());
+    return allBlocks.size() != old_size;
+}
+
 const std::string *CFG::typeOf(const Node &node) const {
     if (type_info_ == nullptr) {
         return nullptr;

src/ir/CFG.hpp

@@ -1,6 +1,7 @@
 #ifndef CFG_HPP
 #define CFG_HPP
 
+#include <cstddef>
 #include <memory>
 #include <vector>
 
@@ -40,6 +41,8 @@ class CFG {
     [[nodiscard]] BBlock *addMethodRootBlock(const std::string &className,
                                              const std::string &methodName);
     [[nodiscard]] const auto &getMethodRoots() const { return methodRoots; }
+    [[nodiscard]] std::size_t getBlockCount() const { return allBlocks.size(); }
+    [[nodiscard]] bool removeUnreachableBlocks();
 
     void setTypeInfo(const TypeInfo *info) { type_info_ = info; }
     [[nodiscard]] const std::string *typeOf(const Node &node) const;

src/ir/passes/UnreachableBlockEliminationPass.cpp

@@ -0,0 +1,7 @@
+#include "ir/passes/UnreachableBlockEliminationPass.hpp"
+
+#include "ir/CFG.hpp"
+
+bool UnreachableBlockEliminationPass::run(CFG &graph) {
+    return graph.removeUnreachableBlocks();
+}

src/ir/passes/UnreachableBlockEliminationPass.hpp

@@ -0,0 +1,16 @@
+#ifndef UNREACHABLE_BLOCK_ELIMINATION_PASS_HPP
+#define UNREACHABLE_BLOCK_ELIMINATION_PASS_HPP
+
+#include <string_view>
+
+#include "ir/passes/IRPass.hpp"
+
+class UnreachableBlockEliminationPass final : public IRPass {
+  public:
+    [[nodiscard]] std::string_view name() const override {
+        return "unreachable-block-elimination";
+    }
+    bool run(CFG &graph) override;
+};
+
+#endif

src/main.cpp

@@ -14,6 +14,7 @@ namespace fs = std::filesystem;
 #include "ir/passes/ConditionalJumpFoldingPass.hpp"
 #include "ir/passes/ConstantFoldingPass.hpp"
 #include "ir/passes/IRPassManager.hpp"
+#include "ir/passes/UnreachableBlockEliminationPass.hpp"
 #include "lexing/LegacyDiagnostics.hpp"
 #include "lexing/Lexer.hpp"
 #include "lexing/SourceBuffer.hpp"
@@ -322,6 +323,7 @@ int main(int argc, char **argv) {
     IRPassManager pass_manager;
     pass_manager.addPass(std::make_unique<ConstantFoldingPass>());
     pass_manager.addPass(std::make_unique<ConditionalJumpFoldingPass>());
+    pass_manager.addPass(std::make_unique<UnreachableBlockEliminationPass>());
     (void)pass_manager.run(graph);
 
     graph.printGraphviz(controlFlowGraph);

tests/ir_unreachable_block_elimination_test.cpp

@@ -0,0 +1,69 @@
+#include <gtest/gtest.h>
+
+#include "ir/CFG.hpp"
+#include "ir/passes/UnreachableBlockEliminationPass.hpp"
+
+TEST(IRUnreachableBlockElimination, RemovesDetachedSubgraph) {
+    CFG graph;
+
+    auto *root = graph.addMethodRootBlock("Main", "main");
+    auto *reachable = graph.newBlock();
+    root->setTrueBlock(reachable);
+
+    auto *detached_a = graph.newBlock();
+    auto *detached_b = graph.newBlock();
+    detached_a->setTrueBlock(detached_b);
+
+    graph.setCurrentBlock(detached_b);
+    EXPECT_EQ(graph.getBlockCount(), 4U);
+
+    UnreachableBlockEliminationPass pass;
+    EXPECT_TRUE(pass.run(graph));
+    EXPECT_EQ(graph.getBlockCount(), 2U);
+    EXPECT_EQ(graph.getMethodRoots().size(), 1U);
+    EXPECT_EQ(graph.getCurrentBlock(), nullptr);
+
+    EXPECT_FALSE(pass.run(graph));
+    EXPECT_EQ(graph.getBlockCount(), 2U);
+}
+
+TEST(IRUnreachableBlockElimination, NoChangeWhenAllBlocksReachable) {
+    CFG graph;
+
+    auto *root = graph.addMethodRootBlock("Main", "main");
+    auto *then_block = graph.newBlock();
+    auto *else_block = graph.newBlock();
+    root->setTrueBlock(then_block);
+    root->setFalseBlock(else_block);
+    else_block->setTrueBlock(then_block);
+
+    graph.setCurrentBlock(then_block);
+    EXPECT_EQ(graph.getBlockCount(), 3U);
+
+    UnreachableBlockEliminationPass pass;
+    EXPECT_FALSE(pass.run(graph));
+    EXPECT_EQ(graph.getBlockCount(), 3U);
+    EXPECT_EQ(graph.getMethodRoots().size(), 1U);
+    EXPECT_EQ(graph.getCurrentBlock(), then_block);
+}
+
+TEST(IRUnreachableBlockElimination, KeepsMultipleMethodRoots) {
+    CFG graph;
+
+    auto *root_one = graph.addMethodRootBlock("Main", "main");
+    auto *root_two = graph.addMethodRootBlock("Foo", "run");
+    auto *reachable_from_root_one = graph.newBlock();
+    root_one->setTrueBlock(reachable_from_root_one);
+
+    (void)graph.newBlock();
+
+    graph.setCurrentBlock(root_two);
+    EXPECT_EQ(graph.getBlockCount(), 4U);
+    EXPECT_EQ(graph.getMethodRoots().size(), 2U);
+
+    UnreachableBlockEliminationPass pass;
+    EXPECT_TRUE(pass.run(graph));
+    EXPECT_EQ(graph.getBlockCount(), 3U);
+    EXPECT_EQ(graph.getMethodRoots().size(), 2U);
+    EXPECT_EQ(graph.getCurrentBlock(), root_two);
+}