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 new algorithm is based on what was proposed in [^Wei84]. If there is
even one active block in the blocks reachable from a conditional branch
up to its successors' nearest common post-dominator (referred to as
**INFL** in the paper), it is necessary to follow that conditional
branch and execute the code. Otherwise, execution can be short-circuited
from the conditional branch to the nearest common post-dominator.

COMBAK: It is important to note that in Julia's IR (`CodeInfo`),
"short-circuiting" a specific code region is not a simple task. Simply
ignoring the path to the post-dominator does not guarantee fall-through
to the post-dominator. Therefore, a more careful implementation is
required for this aspect.

[Wei84]: M. Weiser, "Program Slicing," IEEE Transactions on Software Engineering, 10, pages 352-357, July 1984.

Author: aviatesk

src/JET.jl

@@ -41,9 +41,10 @@ using .CC: @nospecs, ⊑,
     InvokeCallInfo, MethodCallResult, MethodMatchInfo, MethodMatches, NOT_FOUND,
     OptimizationState, OptimizationParams, OverlayMethodTable, StmtInfo, UnionSplitInfo,
     UnionSplitMethodMatches, VarState, VarTable, WorldRange, WorldView,
-    argextype, argtype_by_index, argtypes_to_type, hasintersect, ignorelimited,
-    instanceof_tfunc, istopfunction, singleton_type, slot_id, specialize_method,
-    tmeet, tmerge, typeinf_lattice, widenconst, widenlattice
+    argextype, argtype_by_index, argtypes_to_type, compute_basic_blocks, construct_domtree,
+    construct_postdomtree, hasintersect, ignorelimited, instanceof_tfunc, istopfunction,
+    nearest_common_dominator, singleton_type, slot_id, specialize_method, tmeet, tmerge,
+    typeinf_lattice, widenconst, widenlattice
 
 using Base: IdSet, get_world_counter
 

src/toplevel/virtualprocess.jl

@@ -1091,16 +1091,44 @@ end
 
 # select statements that should be concretized, and actually interpreted rather than abstracted
 function select_statements(mod::Module, src::CodeInfo)
-    stmts = src.code
     cl = LoweredCodeUtils.CodeLinks(mod, src) # make `CodeEdges` hold `CodeLinks`?
     edges = LoweredCodeUtils.CodeEdges(src, cl)
-
-    concretize = falses(length(stmts))
-
-    select_direct_requirement!(concretize, stmts, edges)
-
+    concretize = falses(length(src.code))
+    select_direct_requirement!(concretize, src.code, edges)
     select_dependencies!(concretize, src, edges, cl)
+    return concretize
+end
 
+# just for testing, and debugging
+function select_statements(mod::Module, src::CodeInfo, names::Symbol...)
+    cl = LoweredCodeUtils.CodeLinks(mod, src) # make `CodeEdges` hold `CodeLinks`?
+    edges = LoweredCodeUtils.CodeEdges(src, cl)
+    concretize = falses(length(src.code))
+    objs = Set{GlobalRef}(GlobalRef(mod, name) for name in names)
+    LoweredCodeUtils.add_requests!(concretize, objs, edges, ())
+    select_dependencies!(concretize, src, edges, cl)
+    return concretize
+end
+function select_statements(mod::Module, src::CodeInfo, slots::SlotNumber...)
+    cl = LoweredCodeUtils.CodeLinks(mod, src) # make `CodeEdges` hold `CodeLinks`?
+    edges = LoweredCodeUtils.CodeEdges(src, cl)
+    concretize = falses(length(src.code))
+    for slot in slots
+        for d in cl.slotassigns[slot.id]
+            concretize[d] = true
+        end
+    end
+    select_dependencies!(concretize, src, edges, cl)
+    return concretize
+end
+function select_statements(mod::Module, src::CodeInfo, idxs::Int...)
+    cl = LoweredCodeUtils.CodeLinks(mod, src) # make `CodeEdges` hold `CodeLinks`?
+    edges = LoweredCodeUtils.CodeEdges(src, cl)
+    concretize = falses(length(src.code))
+    for idx = idxs
+        concretize[idx] |= true
+    end
+    select_dependencies!(concretize, src, edges, cl)
     return concretize
 end
 
@@ -1173,67 +1201,70 @@ end
 
 # The goal of this function is to request concretization of the minimal necessary control
 # flow to evaluate statements whose concretization have already been requested.
-# The basic approach is to check if there are any active successors for each basic block,
-# 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)
+# The basic algorithm is based on what was proposed in [^Wei84]. If there is even one active
+# block in the blocks reachable from a conditional branch up to its successors' nearest
+# common post-dominator (referred to as **INFL** in the paper), it is necessary to follow
+# that conditional branch and execute the code. Otherwise, execution can be short-circuited
+# from the conditional branch to the nearest common post-dominator.
+#
+# COMBAK: It is important to note that in Julia's intermediate code representation (`CodeInfo`),
+# "short-circuiting" a specific code region is not a simple task. Simply ignoring the path
+# to the post-dominator does not guarantee fall-through to the post-dominator. Therefore,
+# a more careful implementation is required for this aspect.
+#
+# [Wei84]: M. Weiser, "Program Slicing," IEEE Transactions on Software Engineering, 10, pages 352-357, July 1984.
+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 = 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 # 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 = nearest_common_dominator(postdomtree, succ1, succ2)
+            inflblks = reachable_blocks(cfg, succ1, postdominator) ∪ reachable_blocks(cfg, succ2, postdominator)
+            for blk in inflblks
+                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 +1273,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)
@@ -1272,27 +1322,42 @@ function is_arg_requested(@nospecialize(arg), concretize, edges, cl)
     return false
 end
 
+# The purpose of this function is to find other statements that affect the execution of the
+# statements choosen by `select_direct_dependencies!`. In other words, it extracts the
+# minimal amount of code necessary to realize the required concretization.
+# This technique is generally referred to as "program slicing," and specifically as
+# "static program slicing". The basic algorithm implemented here is an extension of the one
+# proposed in https://dl.acm.org/doi/10.5555/800078.802557, which is especially tuned for
+# Julia's intermediate code representation.
 function select_dependencies!(concretize::BitVector, src::CodeInfo, edges, cl)
     typedefs = LoweredCodeUtils.find_typedefs(src)
-    cfg = CC.compute_basic_blocks(src.code)
-    postdomtree = CC.construct_postdomtree(cfg.blocks)
+    cfg = compute_basic_blocks(src.code)
+    domtree = construct_domtree(cfg.blocks)
+    postdomtree = construct_postdomtree(cfg.blocks)
 
     while true
         changed = false
 
-        # discover struct/method definitions at the beginning,
-        # and propagate the definition requirements by tracking SSA precedessors
+        # Discover Dtruct/method definitions at the beginning,
+        # and propagate the definition requirements by tracking SSA precedessors.
+        # (TODO maybe hoist this out of the loop?)
         changed |= LoweredCodeUtils.add_typedefs!(concretize, src, edges, typedefs, ())
         changed |= add_ssa_preds!(concretize, src, edges, ())
 
-        # mark some common inplace operations like `push!(x, ...)` and `setindex!(x, ...)`
-        # when `x` has been marked already: otherwise we may end up using it with invalid state
+        # Mark some common inplace operations like `push!(x, ...)` and `setindex!(x, ...)`
+        # when `x` has been marked already: otherwise we may end up using it with invalid state.
+        # However, note that this is an incomplete approach, and note that the slice created
+        # by this routine will not be sound because of this. This is because
+        # `add_required_inplace!` only requires certain special-cased function calls and
+        # does not take into account the possibility that arguments may be mutated in
+        # arbitrary function calls. Ideally, function calls should be required while
+        # considering the effects of these statements, or by some sort of an
+        # inter-procedural program slicing
         changed |= add_required_inplace!(concretize, src, edges, cl)
         changed |= add_ssa_preds!(concretize, src, edges, ())
 
-        # mark necessary control flows,
-        # and propagate the definition requirements by tracking SSA precedessors
-        changed |= add_control_flow!(concretize, src, cfg, postdomtree)
+        # Mark necessary control flows.
+        changed |= add_control_flow!(concretize, src, cfg, domtree, postdomtree)
         changed |= add_ssa_preds!(concretize, src, edges, ())
 
         changed || break