wip: implement a better statement selection logic

Specifically, this commit aims to review the implementation of
`add_control_flow!` and improves its accuracy. Ideally, it should pass
JET's existing test cases as well as the newly added ones, including the
test cases from JuliaDebug/LoweredCodeUtils.jl#99. The goal is to share
the same high-precision CFG selection logic between LoweredCodeUtils
and JET.

The current implementation is based on
[this paper](https://www.cse.msu.edu/~cse870/Public/Homework/SS2003/HW5/p439-weiser.pdf),
and it has been modified to use an algorithm that checks for liveness in
the reachable blocks up to the nearest common postdominator of the
successors of a conditional terminator.

Author: aviatesk

src/toplevel/virtualprocess.jl

@@ -1177,63 +1177,58 @@ end
 # and if there is an active successor and the terminator is not a fall-through, then request
 # the concretization of that terminator. Additionally, for conditional terminators, a simple
 # optimization using post-domination analysis is also performed.
-function add_control_flow!(concretize::BitVector, src::CodeInfo, cfg::CFG, postdomtree)
+function add_control_flow!(concretize::BitVector, src::CodeInfo, cfg::CFG, domtree, postdomtree)
     local changed::Bool = false
     function mark_concretize!(idx::Int)
         if !concretize[idx]
-            concretize[idx] = true
+            changed |= concretize[idx] = true
             return true
         end
         return false
     end
-    nblocks = length(cfg.blocks)
-    for bbidx = 1:nblocks
-        bb = cfg.blocks[bbidx] # forward traversal
+    for bbidx = 1:length(cfg.blocks) # forward traversal
+        bb = cfg.blocks[bbidx]
         nsuccs = length(bb.succs)
         if nsuccs == 0
             continue
         elseif nsuccs == 1
-            terminator_idx = bb.stmts[end]
-            if src.code[terminator_idx] isa GotoNode
-                # If the destination of this `GotoNode` is not active, it's fine to ignore
-                # the control flow caused by this `GotoNode` and treat it as a fall-through.
-                # If the block that is fallen through to is active and has a dependency on
-                # this goto block, then the concretization of this goto block should already
-                # be requested (at some point of the higher concretization convergence cycle
-                # of `select_dependencies`), and thus, this `GotoNode` will be concretized.
-                if any(@view concretize[cfg.blocks[only(bb.succs)].stmts])
-                    changed |= mark_concretize!(terminator_idx)
+            termidx = bb.stmts[end]
+            if src.code[termidx] isa GotoNode
+                succ = only(bb.succs)
+                if any(@view concretize[cfg.blocks[succ].stmts])
+                    dominator = CC.nearest_common_dominator(domtree, bbidx, succ)
+                    if dominator ≠ succ
+                        for blk in reachable_blocks(cfg, dominator, succ)
+                            if blk == dominator || blk == succ
+                                continue
+                            end
+                            if any(@view concretize[cfg.blocks[blk].stmts])
+                                mark_concretize!(termidx)
+                                break
+                            end
+                        end
+                    else
+                        mark_concretize!(termidx)
+                    end
                 end
             end
+            continue # otherwise we can just fall-through
         elseif nsuccs == 2
-            terminator_idx = bb.stmts[end]
-            @assert is_conditional_terminator(src.code[terminator_idx]) "invalid IR"
+            termidx = bb.stmts[end]
+            @assert is_conditional_terminator(src.code[termidx]) "invalid IR"
             succ1, succ2 = bb.succs
-            succ1_req = any(@view concretize[cfg.blocks[succ1].stmts])
-            succ2_req = any(@view concretize[cfg.blocks[succ2].stmts])
-            if succ1_req
-                if succ2_req
-                    changed |= mark_concretize!(terminator_idx)
-                else
-                    active_bb, inactive_bb = succ1, succ2
-                    @goto asymmetric_case
+            postdominator = CC.nearest_common_dominator(postdomtree, succ1, succ2)
+            for blk in (reachable_blocks(cfg, succ1, postdominator) ∪
+                        reachable_blocks(cfg, succ2, postdominator))
+                if blk == postdominator
+                    continue
                 end
-            elseif succ2_req
-                active_bb, inactive_bb = succ2, succ1
-                @label asymmetric_case
-                # We can ignore the control flow of this conditional terminator and treat
-                # it as a fall-through if only one of its successors is active and the
-                # active block post-dominates the inactive one, since the post-domination
-                # ensures that the active basic block will be reached regardless of the
-                # control flow.
-                if CC.postdominates(postdomtree, active_bb, inactive_bb)
-                    # fall through this block
-                else
-                    changed |= mark_concretize!(terminator_idx)
+                if any(@view concretize[cfg.blocks[blk].stmts])
+                    mark_concretize!(termidx)
+                    break
                 end
-            else
-                # both successors are inactive, just fall through this block
             end
+            # we can just fall-through to the post dominator block (by ignoring all statements between)
         end
     end
     return changed
@@ -1242,6 +1237,25 @@ end
 is_conditional_terminator(@nospecialize stmt) = stmt isa GotoIfNot ||
     (@static @isdefined(EnterNode) ? stmt isa EnterNode : isexpr(stmt, :enter))
 
+function reachable_blocks(cfg::CFG, from_bb::Int, to_bb::Int)
+    worklist = Int[from_bb]
+    visited = BitSet(from_bb)
+    if to_bb == from_bb
+        return visited
+    end
+    push!(visited, to_bb)
+    function visit!(bb::Int)
+        if bb ∉ visited
+            push!(visited, bb)
+            push!(worklist, bb)
+        end
+    end
+    while !isempty(worklist)
+        foreach(visit!, cfg.blocks[pop!(worklist)].succs)
+    end
+    return visited
+end
+
 function add_required_inplace!(concretize::BitVector, src::CodeInfo, edges, cl)
     changed = false
     for i = 1:length(src.code)
@@ -1275,6 +1289,7 @@ end
 function select_dependencies!(concretize::BitVector, src::CodeInfo, edges, cl)
     typedefs = LoweredCodeUtils.find_typedefs(src)
     cfg = CC.compute_basic_blocks(src.code)
+    domtree = CC.construct_domtree(cfg.blocks)
     postdomtree = CC.construct_postdomtree(cfg.blocks)
 
     while true
@@ -1292,7 +1307,7 @@ function select_dependencies!(concretize::BitVector, src::CodeInfo, edges, cl)
 
         # mark necessary control flows,
         # and propagate the definition requirements by tracking SSA precedessors
-        changed |= add_control_flow!(concretize, src, cfg, postdomtree)
+        changed |= add_control_flow!(concretize, src, cfg, domtree, postdomtree)
         changed |= add_ssa_preds!(concretize, src, edges, ())
 
         changed || break

test/toplevel/test_virtualprocess.jl

@@ -1748,7 +1748,7 @@ end
                 found_write = true
                 @test !slice[i]
             elseif (JET.isexpr(stmt, :call) && (arg1 = stmt.args[1]; arg1 isa Core.SSAValue) &&
-                src.code[arg1.id] === :write)
+                    src.code[arg1.id] === :write)
                 found_write = true
                 @test !slice[i]
             end
@@ -1778,6 +1778,39 @@ end
         @test isempty(s)
     end
 
+    # more complex case
+    let src = @src begin
+            x2 = 5
+            a2 = 1
+            @eval global geta2() = $a2 # concretization is forced
+        end
+        slice = JET.select_statements(@__MODULE__, src)
+
+        found_a2 = found_a2_get_binding_type = found_x2 = found_x2_get_binding_type = false
+        for (i, stmt) in enumerate(src.code)
+            if JET.isexpr(stmt, :(=))
+                lhs, rhs = stmt.args
+                if lhs isa GlobalRef
+                    lhs = lhs.name
+                end
+                if lhs === :a2
+                    found_a2 = true
+                    @test slice[i]
+                elseif lhs === :x2
+                    found_x2 = true
+                    @test !slice[i] # this is easy to meet
+                end
+            elseif JET.@capture(stmt, $(GlobalRef(Core, :get_binding_type))(_, :a2))
+                found_a2_get_binding_type = true
+                @test slice[i]
+            elseif JET.@capture(stmt, $(GlobalRef(Core, :get_binding_type))(_, :x2))
+                found_x2_get_binding_type = true
+                @test !slice[i] # this is difficult to meet
+            end
+        end
+        @test found_a2; @test found_a2_get_binding_type; @test found_x2; @test found_x2_get_binding_type
+    end
+
     @testset "captured variables" begin
         let (vmod, res) = @analyze_toplevel2 begin
                 begin