Add trivial jump block elimination 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_trivial_jump_block_elimination_test.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/tests/ir_unreachable_block_elimination_test.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/tests/symbol_table_test.cpp
     )

src/ir/passes/TrivialJumpBlockEliminationPass.cpp

@@ -0,0 +1,221 @@
+#include "ir/passes/TrivialJumpBlockEliminationPass.hpp"
+
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "ir/BBlock.hpp"
+#include "ir/CFG.hpp"
+#include "ir/Tac.hpp"
+
+namespace {
+
+using BlockNameMap = std::unordered_map<std::string, BBlock *>;
+using ResolveCache = std::unordered_map<BBlock *, BBlock *>;
+
+[[nodiscard]] bool is_trampoline_block(const BBlock &block) {
+    const auto &instructions = block.getInstructions();
+    if (instructions.size() != 1U) {
+        return false;
+    }
+    if (dynamic_cast<const JumpTac *>(instructions.front().get()) == nullptr) {
+        return false;
+    }
+    return block.hasTrueBlock() && !block.hasFalseBlock();
+}
+
+void collect_method_blocks(BBlock *root, std::vector<BBlock *> &blocks,
+                           BlockNameMap &block_names) {
+    std::vector<BBlock *> stack{root};
+    std::unordered_set<BBlock *> visited;
+
+    while (!stack.empty()) {
+        auto *block = stack.back();
+        stack.pop_back();
+
+        if (!visited.insert(block).second) {
+            continue;
+        }
+
+        blocks.push_back(block);
+        block_names[block->getName()] = block;
+
+        if (block->hasTrueBlock()) {
+            stack.push_back(block->getTrueBlock());
+        }
+        if (block->hasFalseBlock()) {
+            stack.push_back(block->getFalseBlock());
+        }
+    }
+}
+
+[[nodiscard]] BBlock *resolve_target(BBlock *start, ResolveCache &cache) {
+    if (start == nullptr) {
+        return nullptr;
+    }
+
+    if (const auto it = cache.find(start); it != cache.end()) {
+        return it->second;
+    }
+
+    if (!is_trampoline_block(*start)) {
+        cache[start] = start;
+        return start;
+    }
+
+    std::unordered_set<BBlock *> seen;
+    seen.insert(start);
+
+    auto *current = start;
+    while (is_trampoline_block(*current)) {
+        if (!current->hasTrueBlock()) {
+            cache[start] = start;
+            return start;
+        }
+
+        auto *next = current->getTrueBlock();
+        if (next == nullptr || !seen.insert(next).second) {
+            cache[start] = start;
+            return start;
+        }
+
+        if (!is_trampoline_block(*next)) {
+            cache[start] = next;
+            return next;
+        }
+
+        current = next;
+    }
+
+    cache[start] = current;
+    return current;
+}
+
+bool rewrite_cfg_edges(const std::vector<BBlock *> &blocks, BBlock *root,
+                       ResolveCache &cache) {
+    bool changed = false;
+
+    for (auto *block : blocks) {
+        if (block == nullptr) {
+            continue;
+        }
+
+        if (block == root && is_trampoline_block(*block)) {
+            continue;
+        }
+
+        if (block->hasTrueBlock()) {
+            auto *target = block->getTrueBlock();
+            auto *resolved = resolve_target(target, cache);
+            if (resolved != nullptr && resolved != target) {
+                block->setTrueBlock(resolved);
+                changed = true;
+            }
+        }
+
+        if (block->hasFalseBlock()) {
+            auto *target = block->getFalseBlock();
+            auto *resolved = resolve_target(target, cache);
+            if (resolved != nullptr && resolved != target) {
+                block->setFalseBlock(resolved);
+                changed = true;
+            }
+        }
+    }
+
+    return changed;
+}
+
+bool rewrite_jump_labels(const std::vector<BBlock *> &blocks, BBlock *root,
+                         const BlockNameMap &block_names,
+                         ResolveCache &cache) {
+    bool changed = false;
+
+    for (auto *block : blocks) {
+        if (block == nullptr) {
+            continue;
+        }
+
+        if (block == root && is_trampoline_block(*block)) {
+            continue;
+        }
+
+        auto &instructions = block->getInstructions();
+        for (auto &instruction_ptr : instructions) {
+            if (instruction_ptr == nullptr) {
+                continue;
+            }
+
+            if (auto *jump = dynamic_cast<JumpTac *>(instruction_ptr.get());
+                jump != nullptr) {
+                const auto it = block_names.find(jump->getResult());
+                if (it == block_names.end()) {
+                    continue;
+                }
+
+                auto *resolved = resolve_target(it->second, cache);
+                if (resolved != nullptr && resolved != it->second) {
+                    jump->setResult(resolved->getName());
+                    changed = true;
+                }
+                continue;
+            }
+
+            auto *cond_jump =
+                dynamic_cast<CondJumpTac *>(instruction_ptr.get());
+            if (cond_jump == nullptr) {
+                continue;
+            }
+
+            const auto *label =
+                std::get_if<std::string>(&cond_jump->getRhsOperand());
+            if (label == nullptr) {
+                continue;
+            }
+
+            const auto it = block_names.find(*label);
+            if (it == block_names.end()) {
+                continue;
+            }
+
+            auto *resolved = resolve_target(it->second, cache);
+            if (resolved != nullptr && resolved != it->second) {
+                cond_jump->setRhsOperand(resolved->getName());
+                changed = true;
+            }
+        }
+    }
+
+    return changed;
+}
+
+bool process_method_root(BBlock *root) {
+    std::vector<BBlock *> blocks;
+    BlockNameMap block_names;
+    collect_method_blocks(root, blocks, block_names);
+
+    ResolveCache cache;
+    bool changed = false;
+    changed = rewrite_cfg_edges(blocks, root, cache) || changed;
+    changed = rewrite_jump_labels(blocks, root, block_names, cache) || changed;
+    return changed;
+}
+
+} // namespace
+
+bool TrivialJumpBlockEliminationPass::run(CFG &graph) {
+    bool changed = false;
+
+    for (auto *root : graph.getMethodRoots()) {
+        if (root == nullptr) {
+            continue;
+        }
+
+        changed = process_method_root(root) || changed;
+    }
+
+    const auto removed_blocks = graph.removeUnreachableBlocks();
+    return changed || removed_blocks;
+}
+

src/ir/passes/TrivialJumpBlockEliminationPass.hpp

@@ -0,0 +1,18 @@
+#ifndef TRIVIAL_JUMP_BLOCK_ELIMINATION_PASS_HPP
+#define TRIVIAL_JUMP_BLOCK_ELIMINATION_PASS_HPP
+
+#include <string_view>
+
+#include "ir/passes/IRPass.hpp"
+
+class TrivialJumpBlockEliminationPass final : public IRPass {
+  public:
+    [[nodiscard]] std::string_view name() const override {
+        return "trivial-jump-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/TrivialJumpBlockEliminationPass.hpp"
 #include "ir/passes/UnreachableBlockEliminationPass.hpp"
 #include "lexing/LegacyDiagnostics.hpp"
 #include "lexing/Lexer.hpp"
@@ -323,6 +324,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<TrivialJumpBlockEliminationPass>());
     pass_manager.addPass(std::make_unique<UnreachableBlockEliminationPass>());
     (void)pass_manager.run(graph);