IRStructurizer: handle merge phis for if-then regions (#53)

Author: AntonOresten

IRStructurizer/src/control_tree.jl

@@ -17,6 +17,7 @@ Acyclic regions:
 - `REGION_IF_THEN_ELSE`: Conditional with two symmetric branches `u` ─→ `v` and `u` ─→ `w` and a single merge block reachable by `v` ─→ x or `w` ─→ x
 - `REGION_SWITCH`: Conditional with any number of branches [`u` ─→ `vᵢ`, `u` ─→ `vᵢ₊₁`, ...] and a single merge block reachable by [`vᵢ` ─→ `w`, `vᵢ₊₁` ─→ `w`, ...]
 - `REGION_TERMINATION`: Acyclic region which contains a block `v` with multiple branches, including one or multiple branches to blocks `wᵢ` which end with a function termination instruction. The region is composed of `v` and all the `wᵢ`.
+- `REGION_PROPER`: Generic proper region that doesn't match simpler patterns. Single-entry acyclic subgraph where the entry dominates all nodes.
 
 Cyclic regions:
 - `REGION_WHILE_LOOP`: Simple cyclic region made of a condition block `u`, a loop body block `v` and a merge block `w` such that `v` ⇆ `u` ─→ `w`
@@ -32,6 +33,7 @@ Cyclic regions:
     REGION_FOR_LOOP
     REGION_WHILE_LOOP
     REGION_NATURAL_LOOP
+    REGION_PROPER
 end
 
 @active block_region(args) begin

IRStructurizer/src/structure.jl

@@ -23,6 +23,48 @@ function get_loop_blocks(tree::ControlTree, ir::IRCode)
     return loop_blocks
 end
 
+"""
+    get_region_blocks(tree::ControlTree, ir::IRCode) -> Set{Int}
+
+Get all block indices contained in a control tree region.
+"""
+function get_region_blocks(tree::ControlTree, ir::IRCode)
+    blocks = Set{Int}()
+    nblocks = length(ir.cfg.blocks)
+    for subtree in PreOrderDFS(tree)
+        idx = node_index(subtree)
+        if 1 <= idx <= nblocks
+            push!(blocks, idx)
+        end
+    end
+    return blocks
+end
+
+"""
+    get_exit_block(tree::ControlTree, ir::IRCode) -> Int
+
+Get the exit block index of a control tree region.
+For single-block regions, this is the block itself.
+For multi-block regions, this is the block that has successors outside the region.
+"""
+function get_exit_block(tree::ControlTree, ir::IRCode)
+    blocks = get_region_blocks(tree, ir)
+    nblocks = length(ir.cfg.blocks)
+
+    # Find block(s) with successors outside the region
+    for block_idx in blocks
+        1 <= block_idx <= nblocks || continue
+        for succ in ir.cfg.blocks[block_idx].succs
+            if !(succ in blocks)
+                return block_idx
+            end
+        end
+    end
+
+    # Fallback to entry block
+    return node_index(tree)
+end
+
 """
     convert_phi_value(val) -> IRValue
 
@@ -89,6 +131,9 @@ function tree_to_block(tree::ControlTree, ir::IRCode, ctx::StructurizationContex
         handle_loop!(block, tree, ir, ctx)
     elseif rtype == REGION_SWITCH
         handle_switch!(block, tree, ir, ctx)
+    elseif rtype == REGION_PROPER
+        # Generic proper region - treat like REGION_BLOCK
+        handle_block_region!(block, tree, ir, ctx)
     else
         error("Unknown region type: $rtype")
     end
@@ -150,6 +195,9 @@ function handle_nested_region!(block::Block, tree::ControlTree, ir::IRCode,
         handle_loop!(block, tree, ir, ctx)
     elseif rtype == REGION_SWITCH
         handle_switch!(block, tree, ir, ctx)
+    elseif rtype == REGION_PROPER
+        # Generic proper region - treat like REGION_BLOCK
+        handle_block_region!(block, tree, ir, ctx)
     else
         error("Unknown region type in nested region: $rtype")
     end
@@ -196,8 +244,9 @@ function handle_if_then_else!(block::Block, tree::ControlTree, ir::IRCode,
     else_blk = tree_to_block(else_tree, ir, ctx)
 
     # Find merge block and detect merge phis
-    then_block_idx = node_index(then_tree)
-    else_block_idx = node_index(else_tree)
+    # Use exit blocks (not entry blocks) for multi-block regions
+    then_block_idx = get_exit_block(then_tree, ir)
+    else_block_idx = get_exit_block(else_tree, ir)
     merge_phis = find_merge_phis(ir, then_block_idx, else_block_idx)
 
     # Add YieldOp terminators with phi values
@@ -321,10 +370,43 @@ function handle_if_then!(block::Block, tree::ControlTree, ir::IRCode,
     # Empty else block
     else_blk = Block()
 
+    # Find merge block and detect merge phis
+    # For if-then, the merge block is the common successor of cond_idx and then_block_idx.
+    # cond_idx acts as the "else" path since it has a direct edge to merge when condition is false.
+    # Use exit block (not entry block) for multi-block then regions
+    then_block_idx = get_exit_block(then_tree, ir)
+    merge_phis = find_merge_phis(ir, then_block_idx, cond_idx)
+
+    # Add YieldOp terminators with phi values
+    if !isempty(merge_phis)
+        then_values = [phi.then_val for phi in merge_phis]
+        else_values = [phi.else_val for phi in merge_phis]
+        then_blk.terminator = YieldOp(then_values)
+        else_blk.terminator = YieldOp(else_values)
+    end
+
     if_op = IfOp(cond_value, then_blk, else_blk)
-    bb = ir.cfg.blocks[cond_idx]
-    result_idx = gotoifnot_idx !== nothing ? gotoifnot_idx : last(bb.stmts)
-    push!(block, result_idx, if_op, Nothing)
+
+    # Allocate synthesized SSA index for IfOp
+    if_result_idx = ctx.next_ssa_idx
+    ctx.next_ssa_idx += 1
+
+    # Use Tuple type for IfOp results (uniform handling in codegen)
+    if !isempty(merge_phis)
+        phi_types = [ctx.ssavaluetypes[phi.ssa_idx] for phi in merge_phis]
+        result_type = Tuple{phi_types...}
+        push!(block, if_result_idx, if_op, result_type)
+
+        # Generate getfield statements at original phi indices
+        # This preserves SSA reference semantics: `return %7` still works because getfield is at %7
+        for (i, phi) in enumerate(merge_phis)
+            getfield_expr = Expr(:call, Core.getfield, SSAValue(if_result_idx), i)
+            push!(block, phi.ssa_idx, getfield_expr, phi_types[i])
+        end
+    else
+        # No merge phis - use Tuple{} for uniformity
+        push!(block, if_result_idx, if_op, Tuple{})
+    end
 end
 
 """

IRStructurizer/test/runtests.jl

@@ -644,6 +644,62 @@ end
     validate_scf(sci)
 end
 
+# If-then (no else) must yield phi values, not return Nothing
+@testset "if-then yields phi values" begin
+    @test @filecheck begin
+        code_structured(Tuple{Bool}) do flag::Bool
+            x = 0
+            @check "if"
+            if flag
+                x = 1
+            end
+            @check "yield"
+            @check "else"
+            @check "yield"
+            @check "getfield"
+            return x
+        end
+    end
+end
+
+@testset "if-then phi inside loop" begin
+    @test @filecheck begin
+        code_structured(Tuple{Int, Bool}) do n::Int, flag::Bool
+            acc = 0
+            j = 1
+            @check "for"
+            while j <= n
+                x = 0
+                @check "if"
+                if flag && j >= 2
+                    x = 1
+                end
+                @check "yield"
+                @check "getfield"
+                acc += x
+                j += 1
+            end
+            return acc
+        end
+    end
+end
+
+@testset "if-then with multiple phis" begin
+    @test @filecheck begin
+        code_structured(Tuple{Bool}) do flag::Bool
+            x, y = 0, 0
+            @check "if"
+            if flag
+                x, y = 1, 2
+            end
+            @check "yield"
+            @check "getfield"
+            @check "getfield"
+            return x + y
+        end
+    end
+end
+
 end  # regression
 
 #=============================================================================