From ce07bbb64fe99b496d259f14e13501de535b4630 Mon Sep 17 00:00:00 2001 From: writemorecode Date: Fri, 13 Feb 2026 23:07:30 +0100 Subject: [PATCH] Add conditional jump folding IR pass --- src/ir/CFG.cpp | 5 + src/ir/passes/ConditionalJumpFoldingPass.cpp | 188 +++++++++++++++++++ src/ir/passes/ConditionalJumpFoldingPass.hpp | 16 ++ src/main.cpp | 2 + tests/ir_constant_folding_test.cpp | 116 ++++++++++++ 5 files changed, 327 insertions(+) create mode 100644 src/ir/passes/ConditionalJumpFoldingPass.cpp create mode 100644 src/ir/passes/ConditionalJumpFoldingPass.hpp diff --git a/src/ir/CFG.cpp b/src/ir/CFG.cpp index c183ab2..8c4fd6b 100644 --- a/src/ir/CFG.cpp +++ b/src/ir/CFG.cpp @@ -58,6 +58,11 @@ void CFG::printGraphviz(std::ostream &os) const { for (auto *el : methodRoots) { el->printBlockGraphviz(os); } + for (const auto &block : allBlocks) { + if (!block->isVisited()) { + block->printBlockGraphviz(os); + } + } os << "}\n"; } diff --git a/src/ir/passes/ConditionalJumpFoldingPass.cpp b/src/ir/passes/ConditionalJumpFoldingPass.cpp new file mode 100644 index 0000000..d168690 --- /dev/null +++ b/src/ir/passes/ConditionalJumpFoldingPass.cpp @@ -0,0 +1,188 @@ +#include "ir/passes/ConditionalJumpFoldingPass.hpp" + +#include +#include +#include +#include +#include +#include + +#include "ir/BBlock.hpp" +#include "ir/CFG.hpp" +#include "ir/Tac.hpp" + +namespace { + +using BlockNameMap = std::unordered_map; + +[[nodiscard]] std::optional as_label(const Operand &operand) { + const auto *label = std::get_if(&operand); + if (label == nullptr) { + return std::nullopt; + } + return *label; +} + +void collect_method_blocks(BBlock *root, std::vector &blocks, + BlockNameMap &block_names) { + std::vector stack{root}; + std::unordered_set 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()); + } + } +} + +struct FoldDecision { + std::string target_label; + BBlock *target_block = nullptr; +}; + +[[nodiscard]] const JumpTac * +next_jump_instruction(const std::vector> &instructions, + std::size_t index) { + if (index + 1 >= instructions.size()) { + return nullptr; + } + return dynamic_cast(instructions[index + 1].get()); +} + +[[nodiscard]] std::optional +resolve_target_label(const CondJumpTac &conditional_jump, int condition_value, + BBlock &block, + const std::vector> &instructions, + std::size_t index) { + if (condition_value == 0) { + return as_label(conditional_jump.getRhsOperand()); + } + + if (const auto *next_jump = next_jump_instruction(instructions, index); + next_jump != nullptr) { + return next_jump->getResult(); + } + if (block.hasTrueBlock()) { + return block.getTrueBlock()->getName(); + } + return std::nullopt; +} + +[[nodiscard]] BBlock *resolve_target_block(BBlock &block, + int condition_value, + const std::string &target_label, + const BlockNameMap &block_names) { + if (condition_value == 0 && block.hasFalseBlock()) { + return block.getFalseBlock(); + } + if (condition_value != 0 && block.hasTrueBlock()) { + return block.getTrueBlock(); + } + + if (const auto it = block_names.find(target_label); it != block_names.end()) { + return it->second; + } + return nullptr; +} + +[[nodiscard]] std::optional +build_fold_decision(const CondJumpTac &conditional_jump, int condition_value, + BBlock &block, + const std::vector> &instructions, + std::size_t index, const BlockNameMap &block_names) { + const auto target_label = resolve_target_label( + conditional_jump, condition_value, block, instructions, index); + if (!target_label.has_value()) { + return std::nullopt; + } + + BBlock *target_block = resolve_target_block( + block, condition_value, *target_label, block_names); + if (target_block == nullptr) { + return std::nullopt; + } + + return FoldDecision{.target_label = *target_label, + .target_block = target_block}; +} + +bool apply_fold(BBlock &block, std::vector> &instructions, + std::size_t index, const FoldDecision &decision) { + instructions[index] = std::make_unique(decision.target_label); + + if (next_jump_instruction(instructions, index) != nullptr) { + instructions.erase(instructions.begin() + static_cast(index + 1)); + } + + block.setTrueBlock(decision.target_block); + block.setFalseBlock(nullptr); + return true; +} + +bool process_block(BBlock &block, const BlockNameMap &block_names) { + auto &instructions = block.getInstructions(); + bool changed = false; + + for (std::size_t index = 0; index < instructions.size(); ++index) { + const auto *conditional_jump = + dynamic_cast(instructions[index].get()); + if (conditional_jump == nullptr) { + continue; + } + + const auto *condition_value = + std::get_if(&conditional_jump->getLhsOperand()); + if (condition_value == nullptr) { + continue; + } + + const auto decision = build_fold_decision( + *conditional_jump, *condition_value, block, instructions, index, + block_names); + if (!decision.has_value()) { + continue; + } + + changed = apply_fold(block, instructions, index, *decision) || changed; + } + + return changed; +} + +bool process_method_root(BBlock *root) { + std::vector blocks; + BlockNameMap block_names; + collect_method_blocks(root, blocks, block_names); + + bool changed = false; + for (auto *block : blocks) { + changed = process_block(*block, block_names) || changed; + } + return changed; +} + +} // namespace + +bool ConditionalJumpFoldingPass::run(CFG &graph) { + bool changed = false; + for (auto *root : graph.getMethodRoots()) { + if (root == nullptr) { + continue; + } + changed = process_method_root(root) || changed; + } + return changed; +} diff --git a/src/ir/passes/ConditionalJumpFoldingPass.hpp b/src/ir/passes/ConditionalJumpFoldingPass.hpp new file mode 100644 index 0000000..f720638 --- /dev/null +++ b/src/ir/passes/ConditionalJumpFoldingPass.hpp @@ -0,0 +1,16 @@ +#ifndef CONDITIONAL_JUMP_FOLDING_PASS_HPP +#define CONDITIONAL_JUMP_FOLDING_PASS_HPP + +#include + +#include "ir/passes/IRPass.hpp" + +class ConditionalJumpFoldingPass final : public IRPass { + public: + [[nodiscard]] std::string_view name() const override { + return "conditional-jump-folding"; + } + bool run(CFG &graph) override; +}; + +#endif diff --git a/src/main.cpp b/src/main.cpp index 4ff8dc1..a496810 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -11,6 +11,7 @@ namespace fs = std::filesystem; #include "bytecode/BytecodeProgram.hpp" #include "ir/CFG.hpp" #include "ir/IRGenerationVisitor.hpp" +#include "ir/passes/ConditionalJumpFoldingPass.hpp" #include "ir/passes/ConstantFoldingPass.hpp" #include "ir/passes/IRPassManager.hpp" #include "lexing/LegacyDiagnostics.hpp" @@ -320,6 +321,7 @@ int main(int argc, char **argv) { IRPassManager pass_manager; pass_manager.addPass(std::make_unique()); + pass_manager.addPass(std::make_unique()); (void)pass_manager.run(graph); graph.printGraphviz(controlFlowGraph); diff --git a/tests/ir_constant_folding_test.cpp b/tests/ir_constant_folding_test.cpp index 3534a18..cecd829 100644 --- a/tests/ir_constant_folding_test.cpp +++ b/tests/ir_constant_folding_test.cpp @@ -16,6 +16,7 @@ #include "bytecode/Opcode.hpp" #include "ir/CFG.hpp" #include "ir/IRGenerationVisitor.hpp" +#include "ir/passes/ConditionalJumpFoldingPass.hpp" #include "ir/passes/ConstantFoldingPass.hpp" #include "ir/passes/IRPassManager.hpp" #include "lexing/Diagnostics.hpp" @@ -90,6 +91,7 @@ compile_with_constant_folding(std::string_view source) { IRPassManager pass_manager; pass_manager.addPass(std::make_unique()); + pass_manager.addPass(std::make_unique()); (void)pass_manager.run(graph); auto program = std::make_unique(); @@ -233,6 +235,120 @@ class Foo { EXPECT_TRUE(contains_instruction(instructions, Opcode::DIV)); } +TEST(IRConstantFolding, + IfElseConstantTrue_RemovesCjmp_ElseBecomesUnreachable) { + constexpr std::string_view source = + R"(public class Main { + public static void main(String[] args) { + System.out.println(new Foo().run()); + } +} + +class Foo { + public int run() { + int x; + if (true) { + x = 111; + } else { + x = 222; + } + return x; + } +} +)"; + + const auto program = compile_with_constant_folding(source); + ASSERT_NE(program, nullptr); + const auto instructions = collect_instructions(*program); + EXPECT_FALSE(contains_instruction(instructions, Opcode::CJMP)); + EXPECT_TRUE(contains_const(instructions, 111)); + EXPECT_FALSE(contains_const(instructions, 222)); +} + +TEST(IRConstantFolding, + IfElseConstantFalse_RemovesCjmp_ThenBecomesUnreachable) { + constexpr std::string_view source = + R"(public class Main { + public static void main(String[] args) { + System.out.println(new Foo().run()); + } +} + +class Foo { + public int run() { + int x; + if (false) { + x = 333; + } else { + x = 444; + } + return x; + } +} +)"; + + const auto program = compile_with_constant_folding(source); + ASSERT_NE(program, nullptr); + const auto instructions = collect_instructions(*program); + EXPECT_FALSE(contains_instruction(instructions, Opcode::CJMP)); + EXPECT_FALSE(contains_const(instructions, 333)); + EXPECT_TRUE(contains_const(instructions, 444)); +} + +TEST(IRConstantFolding, WhileConstantFalse_RemovesCjmp_BodyBecomesUnreachable) { + constexpr std::string_view source = + R"(public class Main { + public static void main(String[] args) { + System.out.println(new Foo().run()); + } +} + +class Foo { + public int run() { + int x; + x = 5; + while (false) { + x = 777; + } + return x; + } +} +)"; + + const auto program = compile_with_constant_folding(source); + ASSERT_NE(program, nullptr); + const auto instructions = collect_instructions(*program); + EXPECT_FALSE(contains_instruction(instructions, Opcode::CJMP)); + EXPECT_TRUE(contains_const(instructions, 5)); + EXPECT_FALSE(contains_const(instructions, 777)); +} + +TEST(IRConstantFolding, NonConstantCondition_KeepsCjmp) { + constexpr std::string_view source = + R"(public class Main { + public static void main(String[] args) { + System.out.println(new Foo().run()); + } +} + +class Foo { + public int run() { + int x; + x = 1; + while (x < 3) { + x = x + 1; + } + return x; + } +} +)"; + + const auto program = compile_with_constant_folding(source); + ASSERT_NE(program, nullptr); + const auto instructions = collect_instructions(*program); + EXPECT_TRUE(contains_instruction(instructions, Opcode::CJMP)); +} + TEST(IRConstantFolding, SignedSerializerRoundTrip) { const auto unique_suffix = std::to_string( std::chrono::steady_clock::now().time_since_epoch().count());