From 6a9fb505bef5629dea466210d48df60ad3a6941b Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Tue, 16 Jan 2024 20:21:00 -0600
Subject: [PATCH 1/8] Add failing test for proper termination

Currently we use a heuristic to avoid marking the exit of the last
basic block, expecting it to be a `return` and so that execution
will cease whether we evaluate it or not. However, it is possible
to write code for which the final statement is a `GotoNode` and
in this case we can get incorrect answers if we fail to evaluate it.

This example illustrates a tricky point: `return` both terminates
execution but may also return a value. If we don't require the
value, we still might need to terminate execution. This example
seems to illustrate that having `isrequired[i]` be either `true` or
`false` may be insufficiently expressive; we might need it to be
three states, `:no`, `:yes`, `:exit`. During marking, encountering
`:exit` would not force one to evaluate the returned SSAValue.
---
 test/codeedges.jl | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/test/codeedges.jl b/test/codeedges.jl
index 85c44c1..995fade 100644
--- a/test/codeedges.jl
+++ b/test/codeedges.jl
@@ -216,6 +216,23 @@ module ModSelective end
     @test ModSelective.k11 == 0
     @test 3 <= ModSelective.s11 <= 15
 
+    # Final block is not a `return`
+    ex = quote
+        x = 1
+        y = 7
+        @label loop
+        x += 1
+        x < 5 || return y
+        @goto loop
+    end
+    frame = Frame(ModSelective, ex)
+    src = frame.framecode.src
+    edges = CodeEdges(src)
+    isrequired = lines_required(:x, src, edges)
+    selective_eval_fromstart!(frame, isrequired, true)
+    @test ModSelective.x == 5
+    @test !isdefined(ModSelective, :y)
+
     # Control-flow in an abstract type definition
     ex = :(abstract type StructParent{T, N} <: AbstractArray{T, N} end)
     frame = Frame(ModSelective, ex)

From e5db6000103e80fe1075274c5af3956287c6b084 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Sat, 20 Jul 2024 07:26:26 -0500
Subject: [PATCH 2/8] Possible fix

---
 src/codeedges.jl  | 121 ++++++++++++++++++++++++++++------------------
 src/packagedef.jl |   7 ++-
 test/codeedges.jl |  17 ++++---
 3 files changed, 87 insertions(+), 58 deletions(-)

diff --git a/src/codeedges.jl b/src/codeedges.jl
index d6697f3..ab0b253 100644
--- a/src/codeedges.jl
+++ b/src/codeedges.jl
@@ -561,6 +561,8 @@ function terminal_preds(i::Int, edges::CodeEdges)
     return s
 end
 
+initialize_isrequired(n) = fill!(Vector{Union{Bool,Symbol}}(undef, n), false)
+
 """
     isrequired = lines_required(obj::GlobalRef, src::CodeInfo, edges::CodeEdges)
     isrequired = lines_required(idx::Int,                     src::CodeInfo, edges::CodeEdges)
@@ -573,20 +575,20 @@ will end up skipping a subset of such statements, perhaps while repeating others
 See also [`lines_required!`](@ref) and [`selective_eval!`](@ref).
 """
 function lines_required(obj::GlobalRef, src::CodeInfo, edges::CodeEdges; kwargs...)
-    isrequired = falses(length(edges.preds))
+    isrequired = initialize_isrequired(length(src.code))
     objs = Set{GlobalRef}([obj])
     return lines_required!(isrequired, objs, src, edges; kwargs...)
 end
 
 function lines_required(idx::Int, src::CodeInfo, edges::CodeEdges; kwargs...)
-    isrequired = falses(length(edges.preds))
+    isrequired = initialize_isrequired(length(edges.preds))
     isrequired[idx] = true
     objs = Set{GlobalRef}()
     return lines_required!(isrequired, objs, src, edges; kwargs...)
 end
 
 """
-    lines_required!(isrequired::AbstractVector{Bool}, src::CodeInfo, edges::CodeEdges;
+    lines_required!(isrequired::AbstractVector{Union{Bool,Symbol}}, src::CodeInfo, edges::CodeEdges;
                     norequire = ())
 
 Like `lines_required`, but where `isrequired[idx]` has already been set to `true` for all statements
@@ -598,7 +600,7 @@ should _not_ be marked as a requirement.
 For example, use `norequire = LoweredCodeUtils.exclude_named_typedefs(src, edges)` if you're
 extracting method signatures and not evaluating new definitions.
 """
-function lines_required!(isrequired::AbstractVector{Bool}, src::CodeInfo, edges::CodeEdges; kwargs...)
+function lines_required!(isrequired::AbstractVector{Union{Bool,Symbol}}, src::CodeInfo, edges::CodeEdges; kwargs...)
     objs = Set{GlobalRef}()
     return lines_required!(isrequired, objs, src, edges; kwargs...)
 end
@@ -620,7 +622,7 @@ function exclude_named_typedefs(src::CodeInfo, edges::CodeEdges)
     return norequire
 end
 
-function lines_required!(isrequired::AbstractVector{Bool}, objs, src::CodeInfo, edges::CodeEdges; norequire = ())
+function lines_required!(isrequired::AbstractVector{Union{Bool,Symbol}}, objs, src::CodeInfo, edges::CodeEdges; norequire = ())
     # Mark any requested objects (their lines of assignment)
     objs = add_requests!(isrequired, objs, edges, norequire)
 
@@ -638,6 +640,9 @@ function lines_required!(isrequired::AbstractVector{Bool}, objs, src::CodeInfo,
     iter = 0
     while changed
         changed = false
+        @show iter
+        print_with_code(stdout, src, isrequired)
+        println()
 
         # Handle ssa predecessors
         changed |= add_ssa_preds!(isrequired, src, edges, norequire)
@@ -646,8 +651,8 @@ function lines_required!(isrequired::AbstractVector{Bool}, objs, src::CodeInfo,
         changed |= add_named_dependencies!(isrequired, edges, objs, norequire)
 
         # Add control-flow
-        changed |= add_loops!(isrequired, cfg)
-        changed |= add_control_flow!(isrequired, cfg, domtree, postdomtree)
+        changed |= add_loops!(isrequired, cfg, domtree, postdomtree)
+        changed |= add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
 
         # So far, everything is generic graph traversal. Now we add some domain-specific information
         changed |= add_typedefs!(isrequired, src, edges, typedefs, norequire)
@@ -669,7 +674,7 @@ end
 function add_ssa_preds!(isrequired, src::CodeInfo, edges::CodeEdges, norequire)
     changed = false
     for idx = 1:length(src.code)
-        if isrequired[idx]
+        if isrequired[idx] == true
             changed |= add_preds!(isrequired, idx, edges, norequire)
         end
     end
@@ -680,7 +685,7 @@ function add_named_dependencies!(isrequired, edges::CodeEdges, objs, norequire)
     changed = false
     for (obj, uses) in edges.byname
         obj ∈ objs && continue
-        if any(view(isrequired, uses.succs))
+        if any(==(true), view(isrequired, uses.succs))
             changed |= add_obj!(isrequired, objs, obj, edges, norequire)
         end
     end
@@ -691,7 +696,7 @@ function add_preds!(isrequired, idx, edges::CodeEdges, norequire)
     chngd = false
     preds = edges.preds[idx]
     for p in preds
-        isrequired[p] && continue
+        isrequired[p] == true && continue
         p ∈ norequire && continue
         isrequired[p] = true
         chngd = true
@@ -702,7 +707,7 @@ end
 function add_succs!(isrequired, idx, edges::CodeEdges, succs, norequire)
     chngd = false
     for p in succs
-        isrequired[p] && continue
+        isrequired[p] == true && continue
         p ∈ norequire && continue
         isrequired[p] = true
         chngd = true
@@ -713,8 +718,9 @@ end
 function add_obj!(isrequired, objs, obj, edges::CodeEdges, norequire)
     chngd = false
     for d in edges.byname[obj].assigned
+        isrequired[d] == true && continue
         d ∈ norequire && continue
-        isrequired[d] || add_preds!(isrequired, d, edges, norequire)
+        add_preds!(isrequired, d, edges, norequire)
         isrequired[d] = true
         chngd = true
     end
@@ -725,30 +731,25 @@ end
 ## Add control-flow
 
 # Mark loops that contain evaluated statements
-function add_loops!(isrequired, cfg)
+function add_loops!(isrequired, cfg, domtree, postdomtree)
     changed = false
     for (ibb, bb) in enumerate(cfg.blocks)
-        needed = false
         for ibbp in bb.preds
             # Is there a backwards-pointing predecessor, and if so are there any required statements between the two?
             ibbp > ibb || continue   # not a loop-block predecessor
-            r, rp = rng(bb), rng(cfg.blocks[ibbp])
-            r = first(r):first(rp)-1
-            needed |= any(view(isrequired, r))
-        end
-        if needed
-            # Mark the final statement of all predecessors
-            for ibbp in bb.preds
-                rp = rng(cfg.blocks[ibbp])
-                changed |= !isrequired[last(rp)]
-                isrequired[last(rp)] = true
+            if postdominates(postdomtree, ibb, ibbp)
+                r = rng(cfg.blocks[ibbp])
+                if isrequired[r[end]] != true
+                    isrequired[r[end]] = true
+                    changed = true
+                end
             end
         end
     end
     return changed
 end
 
-function add_control_flow!(isrequired, cfg, domtree, postdomtree)
+function add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
     changed, _changed = false, true
     blocks = cfg.blocks
     nblocks = length(blocks)
@@ -756,7 +757,7 @@ function add_control_flow!(isrequired, cfg, domtree, postdomtree)
         _changed = false
         for (ibb, bb) in enumerate(blocks)
             r = rng(bb)
-            if any(view(isrequired, r))
+            if any(==(true), view(isrequired, r))
                 # Walk up the dominators
                 jbb = ibb
                 while jbb != 1
@@ -766,8 +767,11 @@ function add_control_flow!(isrequired, cfg, domtree, postdomtree)
                     for s in dbb.succs
                         if !postdominates(postdomtree, jbb, s)
                             idxlast = rng(dbb)[end]
-                            _changed |= !isrequired[idxlast]
-                            isrequired[idxlast] = true
+                            if isrequired[idxlast] != true
+                                println("add 1: ", idxlast)
+                                _changed = true
+                                isrequired[idxlast] = true
+                            end
                             break
                         end
                     end
@@ -780,11 +784,27 @@ function add_control_flow!(isrequired, cfg, domtree, postdomtree)
                     # Check if the exit of this block is a GotoNode or `return`
                     if length(pdbb.succs) < 2
                         idxlast = rng(pdbb)[end]
-                        _changed |= !isrequired[idxlast]
-                        isrequired[idxlast] = true
+                        stmt = src.code[idxlast]
+                        if isa(stmt, GotoNode) || isa(stmt, Core.ReturnNode)
+                            if isrequired[idxlast] == false
+                                println("add 2: ", idxlast)
+                                _changed = true
+                                isrequired[idxlast] = :exit
+                            end
+                        end
                     end
                     jbb = postdomtree.idoms_bb[jbb]
                 end
+            elseif length(r) == 1
+                # pdbb = blocks[ibb]
+                # if length(pdbb.succs) < 2 && isa(src.code[r[1]], GotoNode)
+                #     idxlast = r[end]
+                #     if isrequired[idxlast] == false
+                #         println("add 3: ", idxlast)
+                #         _changed = true
+                #         isrequired[idxlast] = true
+                #     end
+                # end
             end
         end
         changed |= _changed
@@ -825,11 +845,11 @@ function add_typedefs!(isrequired, src::CodeInfo, edges::CodeEdges, (typedef_blo
     idx = 1
     while idx < length(stmts)
         stmt = stmts[idx]
-        isrequired[idx] || (idx += 1; continue)
+        isrequired[idx] == true || (idx += 1; continue)
         for (typedefr, typedefn) in zip(typedef_blocks, typedef_names)
             if idx ∈ typedefr
                 ireq = view(isrequired, typedefr)
-                if !all(ireq)
+                if !all(==(true), ireq)
                     changed = true
                     ireq .= true
                     # Also mark any by-type constructor(s) associated with this typedef
@@ -857,8 +877,10 @@ function add_typedefs!(isrequired, src::CodeInfo, edges::CodeEdges, (typedef_blo
             if i <= length(stmts) && (stmts[i]::Expr).args[1] == false
                 tpreds = terminal_preds(i, edges)
                 if minimum(tpreds) == idx && i ∉ norequire
-                    changed |= !isrequired[i]
-                    isrequired[i] = true
+                    if isrequired[i] != true
+                        changed = true
+                        isrequired[i] = true
+                    end
                 end
             end
         end
@@ -877,15 +899,17 @@ function add_inplace!(isrequired, src, edges, norequire)
             callee_matches(fname, Base, :pop!) ||
             callee_matches(fname, Base, :empty!) ||
             callee_matches(fname, Base, :setindex!))
-            _changed = !isrequired[j]
-            isrequired[j] = true
+            if isrequired[j] != true
+                _changed = true
+               isrequired[j] = true
+            end
         end
         return _changed
     end
 
     changed = false
     for (i, isreq) in pairs(isrequired)
-        isreq || continue
+        isreq == true || continue
         for j in edges.succs[i]
             j ∈ norequire && continue
             stmt = src.code[j]
@@ -930,14 +954,16 @@ This will return either a `BreakpointRef`, the value obtained from the last exec
 (if stored to `frame.framedata.ssavlues`), or `nothing`.
 Typically, assignment to a variable binding does not result in an ssa store by JuliaInterpreter.
 """
-function selective_eval!(@nospecialize(recurse), frame::Frame, isrequired::AbstractVector{Bool}, istoplevel::Bool=false)
+function selective_eval!(@nospecialize(recurse), frame::Frame, isrequired, istoplevel::Bool=false)
     pc = pcexec = pclast = frame.pc
     while isa(pc, Int)
         frame.pc = pc
         te = isrequired[pc]
         pclast = pcexec::Int
-        if te
+        if te == true
             pcexec = pc = step_expr!(recurse, frame, istoplevel)
+        elseif te == :exit
+            pc = nothing
         else
             pc = next_or_nothing!(frame)
         end
@@ -946,11 +972,11 @@ function selective_eval!(@nospecialize(recurse), frame::Frame, isrequired::Abstr
     pcexec = (pcexec === nothing ? pclast : pcexec)::Int
     frame.pc = pcexec
     node = pc_expr(frame)
-    is_return(node) && return isrequired[pcexec] ? lookup_return(frame, node) : nothing
+    is_return(node) && return isrequired[pcexec] == true ? lookup_return(frame, node) : nothing
     isassigned(frame.framedata.ssavalues, pcexec) && return frame.framedata.ssavalues[pcexec]
     return nothing
 end
-function selective_eval!(frame::Frame, isrequired::AbstractVector{Bool}, istoplevel::Bool=false)
+function selective_eval!(frame::Frame, isrequired, istoplevel::Bool=false)
     selective_eval!(finish_and_return!, frame, isrequired, istoplevel)
 end
 
@@ -959,18 +985,18 @@ end
 
 Like [`selective_eval!`](@ref), except it sets `frame.pc` to the first `true` statement in `isrequired`.
 """
-function selective_eval_fromstart!(@nospecialize(recurse), frame, isrequired, istoplevel::Bool=false)
+function selective_eval_fromstart!(@nospecialize(recurse), frame::Frame, isrequired, istoplevel::Bool=false)
     pc = findfirst(isrequired)
     pc === nothing && return nothing
     frame.pc = pc
     return selective_eval!(recurse, frame, isrequired, istoplevel)
 end
-function selective_eval_fromstart!(frame::Frame, isrequired::AbstractVector{Bool}, istoplevel::Bool=false)
+function selective_eval_fromstart!(frame::Frame, isrequired, istoplevel::Bool=false)
     selective_eval_fromstart!(finish_and_return!, frame, isrequired, istoplevel)
 end
 
 """
-    print_with_code(io, src::CodeInfo, isrequired::AbstractVector{Bool})
+    print_with_code(io, src::CodeInfo, isrequired::AbstractVector{Union{Bool,Symbol}})
 
 Mark each line of code with its requirement status.
 
@@ -978,10 +1004,13 @@ Mark each line of code with its requirement status.
     This function produces dummy output if suitable support is missing in your version of Julia.
 
 """
-function print_with_code(io::IO, src::CodeInfo, isrequired::AbstractVector{Bool})
+function print_with_code(io::IO, src::CodeInfo, isrequired::AbstractVector{Union{Bool,Symbol}})
+    function markchar(c)
+        return c === true ? 't' : (c === false ? 'f' : (c === :exit ? 'e' : 'x'))
+    end
     nd = ndigits(length(isrequired))
     preprint(::IO) = nothing
-    preprint(io::IO, idx::Int) = (c = isrequired[idx]; printstyled(io, lpad(idx, nd), ' ', c ? "t " : "f "; color = c ? :cyan : :plain))
+    preprint(io::IO, idx::Int) = (c = markchar(isrequired[idx]); printstyled(io, lpad(idx, nd), ' ', c; color = c ∈ ('t', 'e') ? :cyan : :plain))
     postprint(::IO) = nothing
     postprint(io::IO, idx::Int, bbchanged::Bool) = nothing
 
diff --git a/src/packagedef.jl b/src/packagedef.jl
index 75833e8..2558477 100644
--- a/src/packagedef.jl
+++ b/src/packagedef.jl
@@ -21,6 +21,7 @@ if Base.VERSION < v"1.10"
 else
     const construct_domtree = Core.Compiler.construct_domtree
     const construct_postdomtree = Core.Compiler.construct_postdomtree
+    const dominates = Core.Compiler.dominates
     const postdominates = Core.Compiler.postdominates
 end
 
@@ -46,10 +47,8 @@ if ccall(:jl_generating_output, Cint, ()) == 1
     isrequired = lines_required(GlobalRef(@__MODULE__, :s), src, edges)
     lines_required(GlobalRef(@__MODULE__, :s), src, edges; norequire=())
     lines_required(GlobalRef(@__MODULE__, :s), src, edges; norequire=exclude_named_typedefs(src, edges))
-    for isreq in (isrequired, convert(Vector{Bool}, isrequired))
-        lines_required!(isreq, src, edges; norequire=())
-        lines_required!(isreq, src, edges; norequire=exclude_named_typedefs(src, edges))
-    end
+    lines_required!(isrequired, src, edges; norequire=())
+    lines_required!(isrequired, src, edges; norequire=exclude_named_typedefs(src, edges))
     frame = Frame(@__MODULE__, src)
     # selective_eval_fromstart!(frame, isrequired, true)
     precompile(selective_eval_fromstart!, (typeof(frame), typeof(isrequired), Bool))  # can't @eval during precompilation
diff --git a/test/codeedges.jl b/test/codeedges.jl
index 995fade..09b1e42 100644
--- a/test/codeedges.jl
+++ b/test/codeedges.jl
@@ -25,7 +25,7 @@ function hastrackedexpr(stmt; heads=LoweredCodeUtils.trackedheads)
 end
 
 function minimal_evaluation(predicate, src::Core.CodeInfo, edges::CodeEdges; kwargs...)
-    isrequired = fill(false, length(src.code))
+    isrequired = LoweredCodeUtils.initialize_isrequired(length(src.code))
     for (i, stmt) in enumerate(src.code)
         if !isrequired[i]
             isrequired[i], haseval = predicate(stmt)
@@ -65,7 +65,7 @@ module ModSelective end
     # Check that the result of direct evaluation agrees with selective evaluation
     Core.eval(ModEval, ex)
     isrequired = lines_required(GlobalRef(ModSelective, :x), src, edges)
-    # theere is too much diversity in lowering across Julia versions to make it useful to test `sum(isrequired)`
+    # there is too much diversity in lowering across Julia versions to make it useful to test `sum(isrequired)`
     selective_eval_fromstart!(frame, isrequired)
     @test ModSelective.x === ModEval.x
     @test allmissing(ModSelective, (:y, :z, :a, :b, :k))
@@ -160,7 +160,7 @@ module ModSelective end
     src = frame.framecode.src
     edges = CodeEdges(ModSelective, src)
     isrequired = lines_required(GlobalRef(ModSelective, :c_os), src, edges)
-    @test sum(isrequired) >= length(isrequired) - 3
+    @test sum(∈((true, :exit)), isrequired) >= length(isrequired) - 3
     selective_eval_fromstart!(frame, isrequired)
     Core.eval(ModEval, ex)
     @test ModSelective.c_os === ModEval.c_os == Sys.iswindows()
@@ -183,7 +183,7 @@ module ModSelective end
     # Mark just the load of Core.eval
     haseval(stmt) = (isa(stmt, Expr) && JuliaInterpreter.hasarg(isequal(:eval), stmt.args)) ||
                     (isa(stmt, Expr) && stmt.head === :call && is_quotenode(stmt.args[1], Core.eval))
-    isrequired = map(haseval, src.code)
+    isrequired = Union{Bool,Symbol}[haseval(stmt) for stmt in src.code]
     @test sum(isrequired) == 1
     isrequired[edges.succs[findfirst(isrequired)]] .= true   # add lines that use Core.eval
     lines_required!(isrequired, src, edges)
@@ -227,8 +227,8 @@ module ModSelective end
     end
     frame = Frame(ModSelective, ex)
     src = frame.framecode.src
-    edges = CodeEdges(src)
-    isrequired = lines_required(:x, src, edges)
+    edges = CodeEdges(ModSelective, src)
+    isrequired = lines_required(GlobalRef(ModSelective, :x), src, edges)
     selective_eval_fromstart!(frame, isrequired, true)
     @test ModSelective.x == 5
     @test !isdefined(ModSelective, :y)
@@ -297,7 +297,7 @@ module ModSelective end
     frame = Frame(ModSelective, ex)
     src = frame.framecode.src
     edges = CodeEdges(ModSelective, src)
-    isrequired = fill(false, length(src.code))
+    isrequired = LoweredCodeUtils.initialize_isrequired(length(src.code))
     j = length(src.code) - 1
     if !Meta.isexpr(src.code[end-1], :method, 3)
         j -= 1
@@ -471,7 +471,8 @@ module ModSelective end
         end)
         lwr = Meta.lower(Main, ex)
         src = lwr.args[1]
-        LoweredCodeUtils.print_with_code(io, src, trues(length(src.code)))
+        isrq = fill!(LoweredCodeUtils.initialize_isrequired(length(src.code)), true)
+        LoweredCodeUtils.print_with_code(io, src, isrq)
         str = String(take!(io))
         @test count("s = ", str) == 2
         @test count("i = ", str) == 1

From a1d2a28dae4093eefb937902c353d4c46a1484a1 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Sat, 20 Jul 2024 18:26:59 -0500
Subject: [PATCH 3/8] Rewrite

---
 src/codeedges.jl | 138 +++++++++++++++++++++++++++++------------------
 1 file changed, 87 insertions(+), 51 deletions(-)

diff --git a/src/codeedges.jl b/src/codeedges.jl
index ab0b253..5c26164 100644
--- a/src/codeedges.jl
+++ b/src/codeedges.jl
@@ -640,9 +640,6 @@ function lines_required!(isrequired::AbstractVector{Union{Bool,Symbol}}, objs, s
     iter = 0
     while changed
         changed = false
-        @show iter
-        print_with_code(stdout, src, isrequired)
-        println()
 
         # Handle ssa predecessors
         changed |= add_ssa_preds!(isrequired, src, edges, norequire)
@@ -651,7 +648,6 @@ function lines_required!(isrequired::AbstractVector{Union{Bool,Symbol}}, objs, s
         changed |= add_named_dependencies!(isrequired, edges, objs, norequire)
 
         # Add control-flow
-        changed |= add_loops!(isrequired, cfg, domtree, postdomtree)
         changed |= add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
 
         # So far, everything is generic graph traversal. Now we add some domain-specific information
@@ -730,85 +726,125 @@ end
 
 ## Add control-flow
 
-# Mark loops that contain evaluated statements
-function add_loops!(isrequired, cfg, domtree, postdomtree)
-    changed = false
-    for (ibb, bb) in enumerate(cfg.blocks)
-        for ibbp in bb.preds
-            # Is there a backwards-pointing predecessor, and if so are there any required statements between the two?
-            ibbp > ibb || continue   # not a loop-block predecessor
-            if postdominates(postdomtree, ibb, ibbp)
-                r = rng(cfg.blocks[ibbp])
-                if isrequired[r[end]] != true
-                    isrequired[r[end]] = true
-                    changed = true
-                end
-            end
+iscf(stmt) = isa(stmt, Core.GotoNode) || isa(stmt, Core.GotoIfNot) || isa(stmt, Core.ReturnNode)
+
+"""
+    ispredecessor(blocks, i, j)
+
+Determine whether block `i` is a predecessor of block `j` in the control-flow graph `blocks`.
+"""
+function ispredecessor(blocks, i, j, cache=Set{Int}())
+    for p in blocks[j].preds  # avoid putting `j` in the cache unless it loops back
+        getpreds!(cache, blocks, p)
+    end
+    return i ∈ cache
+end
+function getpreds!(cache, blocks, j)
+    if j ∈ cache
+        return cache
+    end
+    push!(cache, j)
+    for p in blocks[j].preds
+        getpreds!(cache, blocks, p)
+    end
+    return cache
+end
+
+function block_internals_needed(isrequired, src, r)
+    needed = false
+    for i in r
+        if isrequired[i] == true
+            iscf(src.code[i]) && continue
+            needed = true
+            break
         end
     end
-    return changed
+    return needed
 end
 
 function add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
     changed, _changed = false, true
     blocks = cfg.blocks
-    nblocks = length(blocks)
+    needed = falses(length(blocks))
+    cache = Set{Int}()
     while _changed
         _changed = false
         for (ibb, bb) in enumerate(blocks)
             r = rng(bb)
-            if any(==(true), view(isrequired, r))
-                # Walk up the dominators
+            if block_internals_needed(isrequired, src, r)
+                needed[ibb] = true
+                # Check control flow that's needed to reach this block by walking up the dominators
                 jbb = ibb
                 while jbb != 1
-                    jdbb = domtree.idoms_bb[jbb]
+                    jdbb = domtree.idoms_bb[jbb]   # immediate dominator of jbb
                     dbb = blocks[jdbb]
-                    # Check the successors; if jbb doesn't post-dominate, mark the last statement
-                    for s in dbb.succs
-                        if !postdominates(postdomtree, jbb, s)
-                            idxlast = rng(dbb)[end]
-                            if isrequired[idxlast] != true
-                                println("add 1: ", idxlast)
-                                _changed = true
-                                isrequired[idxlast] = true
+                    idxlast = rng(dbb)[end]
+                    if iscf(src.code[idxlast])
+                        # Check the idom's successors; if jbb doesn't post-dominate, mark the last statement
+                        for s in dbb.succs
+                            if !postdominates(postdomtree, jbb, s)
+                                if isrequired[idxlast] != true
+                                    _changed = true
+                                    isrequired[idxlast] = true
+                                    break
+                                end
                             end
-                            break
                         end
                     end
                     jbb = jdbb
                 end
-                # Walk down the post-dominators, including self
+                # Walk down the post-dominators, starting with self
                 jbb = ibb
-                while jbb != 0 && jbb < nblocks
-                    pdbb = blocks[jbb]
-                    # Check if the exit of this block is a GotoNode or `return`
-                    if length(pdbb.succs) < 2
+                while jbb != 0
+                    empty!(cache)
+                    if ispredecessor(blocks, jbb, ibb, cache)  # is post-dominator jbb also a predecessor of ibb? If so we have a loop.
+                        pdbb = blocks[jbb]
                         idxlast = rng(pdbb)[end]
                         stmt = src.code[idxlast]
-                        if isa(stmt, GotoNode) || isa(stmt, Core.ReturnNode)
-                            if isrequired[idxlast] == false
-                                println("add 2: ", idxlast)
+                        if iscf(stmt)
+                            if isrequired[idxlast] != true
                                 _changed = true
-                                isrequired[idxlast] = :exit
+                                if isa(stmt, Core.ReturnNode) && isrequired[idxlast] != :exit
+                                    isrequired[idxlast] = :exit
+                                else
+                                    isrequired[idxlast] = true
+                                    if isa(stmt, Core.GotoIfNot) && idxlast < length(isrequired) && isrequired[idxlast+1] != true && iscf(src.code[idxlast+1])
+                                        isrequired[idxlast+1] = true
+                                    end
+                                end
                             end
                         end
                     end
                     jbb = postdomtree.idoms_bb[jbb]
                 end
-            elseif length(r) == 1
-                # pdbb = blocks[ibb]
-                # if length(pdbb.succs) < 2 && isa(src.code[r[1]], GotoNode)
-                #     idxlast = r[end]
-                #     if isrequired[idxlast] == false
-                #         println("add 3: ", idxlast)
-                #         _changed = true
-                #         isrequired[idxlast] = true
-                #     end
-                # end
             end
         end
         changed |= _changed
     end
+    # Now handle "exclusions": in code that would fall through during selective evaluation, find a post-dominator between the two
+    # that is marked, or mark the end block
+    marked = findall(needed)
+    for k in Iterators.drop(eachindex(marked), 1)
+        ibb, jbb = marked[k-1], marked[k]
+        ok = false
+        ipbb = ibb
+        while ipbb < jbb
+            ipbb = postdomtree.idoms_bb[ipbb]
+            ipbb == 0 && break
+            idxlast = rng(blocks[ipbb])[end]
+            if isrequired[idxlast] != false
+                ok = true
+                break
+            end
+        end
+        if !ok
+            idxlast = rng(blocks[ibb])[end]
+            if isrequired[idxlast] != true
+                isrequired[idxlast] = true
+                changed = true
+            end
+        end
+    end
     return changed
 end
 

From 047a2733312216645030b54b6dc008d72bf7c510 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Sun, 21 Jul 2024 05:53:06 -0500
Subject: [PATCH 4/8] fix test

---
 test/codeedges.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/codeedges.jl b/test/codeedges.jl
index 09b1e42..4a4cae4 100644
--- a/test/codeedges.jl
+++ b/test/codeedges.jl
@@ -324,9 +324,9 @@ module ModSelective end
     for (iblock, block) in enumerate(bbs.blocks)
         r = LoweredCodeUtils.rng(block)
         if iblock == length(bbs.blocks)
-            @test any(idx->isrequired[idx], r)
+            @test any(idx->isrequired[idx]==true, r)
         else
-            @test !any(idx->isrequired[idx], r)
+            @test !any(idx->isrequired[idx]==true, r)
         end
     end
 

From e049fa966dc080368a04e84490821e6b731b6c18 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Sun, 21 Jul 2024 05:53:36 -0500
Subject: [PATCH 5/8] More sophisticated cf marking

---
 src/codeedges.jl | 94 +++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 82 insertions(+), 12 deletions(-)

diff --git a/src/codeedges.jl b/src/codeedges.jl
index 5c26164..bae9afe 100644
--- a/src/codeedges.jl
+++ b/src/codeedges.jl
@@ -640,6 +640,9 @@ function lines_required!(isrequired::AbstractVector{Union{Bool,Symbol}}, objs, s
     iter = 0
     while changed
         changed = false
+        @show iter
+        print_with_code(stdout, src, isrequired)
+        println()
 
         # Handle ssa predecessors
         changed |= add_ssa_preds!(isrequired, src, edges, norequire)
@@ -727,6 +730,13 @@ end
 ## Add control-flow
 
 iscf(stmt) = isa(stmt, Core.GotoNode) || isa(stmt, Core.GotoIfNot) || isa(stmt, Core.ReturnNode)
+function markcf!(isrequired, src, i)
+    stmt = src.code[i]
+    @assert iscf(stmt)
+    isrequired[i] ∈ (true, :exit) && return false
+    isrequired[i] = isa(stmt, Core.ReturnNode) ? :exit : true
+    return true
+end
 
 """
     ispredecessor(blocks, i, j)
@@ -773,6 +783,13 @@ function add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
             r = rng(bb)
             if block_internals_needed(isrequired, src, r)
                 needed[ibb] = true
+                # Check this blocks precessors and mark any that are just control-flow
+                for p in bb.preds
+                    r = rng(blocks[p])
+                    if length(r) == 1 && iscf(src.code[r[1]])
+                        _changed |= markcf!(isrequired, src, r[1])
+                    end
+                end
                 # Check control flow that's needed to reach this block by walking up the dominators
                 jbb = ibb
                 while jbb != 1
@@ -783,9 +800,8 @@ function add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
                         # Check the idom's successors; if jbb doesn't post-dominate, mark the last statement
                         for s in dbb.succs
                             if !postdominates(postdomtree, jbb, s)
-                                if isrequired[idxlast] != true
+                                if markcf!(isrequired, src, idxlast)
                                     _changed = true
-                                    isrequired[idxlast] = true
                                     break
                                 end
                             end
@@ -796,8 +812,7 @@ function add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
                 # Walk down the post-dominators, starting with self
                 jbb = ibb
                 while jbb != 0
-                    empty!(cache)
-                    if ispredecessor(blocks, jbb, ibb, cache)  # is post-dominator jbb also a predecessor of ibb? If so we have a loop.
+                    if ispredecessor(blocks, jbb, ibb, empty!(cache))  # is post-dominator jbb also a predecessor of ibb? If so we have a loop.
                         pdbb = blocks[jbb]
                         idxlast = rng(pdbb)[end]
                         stmt = src.code[idxlast]
@@ -821,29 +836,84 @@ function add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
         end
         changed |= _changed
     end
-    # Now handle "exclusions": in code that would fall through during selective evaluation, find a post-dominator between the two
-    # that is marked, or mark the end block
-    marked = findall(needed)
+    # Now handle "exclusions": in code that would inappropriately fall through
+    # during selective evaluation, find a post-dominator between the two that is
+    # marked, or mark one if absent.
+    marked = push!(findall(needed), length(blocks)+1)
     for k in Iterators.drop(eachindex(marked), 1)
         ibb, jbb = marked[k-1], marked[k]
+        if jbb <= length(blocks)
+            # are ibb and jbb exclusive?
+            ispredecessor(blocks, ibb, jbb, empty!(cache)) && continue
+        end
+        # Is there a required control-flow statement between ibb and jbb?
         ok = false
         ipbb = ibb
         while ipbb < jbb
             ipbb = postdomtree.idoms_bb[ipbb]
             ipbb == 0 && break
             idxlast = rng(blocks[ipbb])[end]
-            if isrequired[idxlast] != false
+            if iscf(src.code[idxlast]) && isrequired[idxlast] != false
                 ok = true
                 break
             end
         end
         if !ok
-            idxlast = rng(blocks[ibb])[end]
-            if isrequired[idxlast] != true
-                isrequired[idxlast] = true
-                changed = true
+            # Mark a control-flow statement between ibb and jbb
+            ipbb = ibb
+            while ipbb < jbb
+                ipbb = postdomtree.idoms_bb[ipbb]
+                ipbb == 0 && break
+                # Mark the ipostdom's predecessors...
+                for k in blocks[ipbb].preds
+                    idxlast = rng(blocks[k])[end]
+                    if iscf(src.code[idxlast])
+                        if markcf!(isrequired, src, idxlast)
+                            changed = true
+                            ok = true
+                        end
+                    end
+                end
+                # ...or the ipostdom itself
+                if !ok
+                    idxlast = rng(blocks[ipbb])[end]
+                    stmt = src.code[idxlast]
+                    if isa(stmt, Core.GotoNode) || isa(stmt, Core.ReturnNode)  # unconditional jump
+                        if markcf!(isrequired, src, idxlast)
+                            changed = true
+                            ok = true
+                        end
+                    end
+                    # r = rng(blocks[ipbb])
+                    # if length(r) == 1 && iscf(src.code[r[1]])
+                    #     if markcf!(isrequired, src, r[1])
+                    #         changed = true
+                    #         ok = true
+                    #     end
+                    # end
+                end
+                # idxlast = rng(blocks[ipbb])[end]
+                # if iscf(src.code[idxlast])    # ideally we might find an unconditional jump to prevent unnecessary evaluation of the conditional
+                #     if markcf!(isrequired, src, idxlast)
+                #         changed = true
+                #         ok = true
+                #         break
+                #     end
+                # end
+            end
+        end
+        if !ok
+            print_with_code(stdout, src, isrequired)
+            @show ibb jbb ipbb
+            ipbb = postdomtree.idoms_bb[ibb]
+            print("ibb postdoms: ")
+            while ipbb != 0
+                print(ipbb, " ")
+                ipbb = postdomtree.idoms_bb[ipbb]
             end
+            println()
         end
+        jbb <= length(blocks) && @assert ok
     end
     return changed
 end

From e2e780bc45a8de99065de53b2d94458a5491f229 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Sun, 21 Jul 2024 06:58:09 -0500
Subject: [PATCH 6/8] More fix

---
 src/codeedges.jl  | 13 +------------
 test/codeedges.jl |  8 ++++----
 2 files changed, 5 insertions(+), 16 deletions(-)

diff --git a/src/codeedges.jl b/src/codeedges.jl
index bae9afe..e277101 100644
--- a/src/codeedges.jl
+++ b/src/codeedges.jl
@@ -902,17 +902,6 @@ function add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
                 # end
             end
         end
-        if !ok
-            print_with_code(stdout, src, isrequired)
-            @show ibb jbb ipbb
-            ipbb = postdomtree.idoms_bb[ibb]
-            print("ibb postdoms: ")
-            while ipbb != 0
-                print(ipbb, " ")
-                ipbb = postdomtree.idoms_bb[ipbb]
-            end
-            println()
-        end
         jbb <= length(blocks) && @assert ok
     end
     return changed
@@ -953,7 +942,7 @@ function add_typedefs!(isrequired, src::CodeInfo, edges::CodeEdges, (typedef_blo
         stmt = stmts[idx]
         isrequired[idx] == true || (idx += 1; continue)
         for (typedefr, typedefn) in zip(typedef_blocks, typedef_names)
-            if idx ∈ typedefr
+            if idx ∈ typedefr && !iscf(stmt)  # exclude control-flow nodes since they may be needed for other purposes
                 ireq = view(isrequired, typedefr)
                 if !all(==(true), ireq)
                     changed = true
diff --git a/test/codeedges.jl b/test/codeedges.jl
index 4a4cae4..23e6c3e 100644
--- a/test/codeedges.jl
+++ b/test/codeedges.jl
@@ -219,10 +219,10 @@ module ModSelective end
     # Final block is not a `return`
     ex = quote
         x = 1
-        y = 7
+        yy = 7
         @label loop
         x += 1
-        x < 5 || return y
+        x < 5 || return yy
         @goto loop
     end
     frame = Frame(ModSelective, ex)
@@ -231,7 +231,7 @@ module ModSelective end
     isrequired = lines_required(GlobalRef(ModSelective, :x), src, edges)
     selective_eval_fromstart!(frame, isrequired, true)
     @test ModSelective.x == 5
-    @test !isdefined(ModSelective, :y)
+    @test !isdefined(ModSelective, :yy)
 
     # Control-flow in an abstract type definition
     ex = :(abstract type StructParent{T, N} <: AbstractArray{T, N} end)
@@ -326,7 +326,7 @@ module ModSelective end
         if iblock == length(bbs.blocks)
             @test any(idx->isrequired[idx]==true, r)
         else
-            @test !any(idx->isrequired[idx]==true, r)
+            @test !any(idx->isrequired[idx]==true && !LoweredCodeUtils.iscf(src.code[idx]), r)
         end
     end
 

From b64850aa5d2bb9b2daab8fc946951ca1a3ca06f7 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Sun, 21 Jul 2024 12:04:18 -0500
Subject: [PATCH 7/8] working

---
 src/codeedges.jl  | 63 ++++++++++++++++++++++++++++++++---------------
 test/codeedges.jl |  4 ++-
 2 files changed, 46 insertions(+), 21 deletions(-)

diff --git a/src/codeedges.jl b/src/codeedges.jl
index e277101..65e488b 100644
--- a/src/codeedges.jl
+++ b/src/codeedges.jl
@@ -369,8 +369,9 @@ struct CodeEdges
     preds::Vector{Vector{Int}}
     succs::Vector{Vector{Int}}
     byname::Dict{GlobalRef,Variable}
+    slotassigns::Vector{Vector{Int}}
 end
-CodeEdges(n::Integer) = CodeEdges([Int[] for i = 1:n], [Int[] for i = 1:n], Dict{GlobalRef,Variable}())
+CodeEdges(nstmts::Integer, nslots::Integer) = CodeEdges([Int[] for i = 1:nstmts], [Int[] for i = 1:nstmts], Dict{GlobalRef,Variable}(), [Int[] for i = 1:nslots])
 
 function Base.show(io::IO, edges::CodeEdges)
     println(io, "CodeEdges:")
@@ -389,7 +390,7 @@ end
 
 
 """
-    edges = CodeEdges(src::CodeInfo)
+    edges = CodeEdges(mod::Module, src::CodeInfo)
 
 Analyze `src` and determine the chain of dependencies.
 
@@ -410,7 +411,10 @@ function CodeEdges(src::CodeInfo, cl::CodeLinks)
     # Replace/add named intermediates (slot & named-variable references) with statement numbers
     nstmts, nslots = length(src.code), length(src.slotnames)
     marked, slothandled = BitSet(), fill(false, nslots)  # working storage during resolution
-    edges = CodeEdges(nstmts)
+    edges = CodeEdges(nstmts, nslots)
+    for (edge_s, link_s) in zip(edges.slotassigns, cl.slotassigns)
+        append!(edge_s, link_s)
+    end
     emptylink = Links()
     emptylist = Int[]
     for (i, stmt) in enumerate(src.code)
@@ -640,9 +644,9 @@ function lines_required!(isrequired::AbstractVector{Union{Bool,Symbol}}, objs, s
     iter = 0
     while changed
         changed = false
-        @show iter
-        print_with_code(stdout, src, isrequired)
-        println()
+        # @show iter
+        # print_with_code(stdout, src, isrequired)
+        # println()
 
         # Handle ssa predecessors
         changed |= add_ssa_preds!(isrequired, src, edges, norequire)
@@ -730,6 +734,17 @@ end
 ## Add control-flow
 
 iscf(stmt) = isa(stmt, Core.GotoNode) || isa(stmt, Core.GotoIfNot) || isa(stmt, Core.ReturnNode)
+function jumps_back(src, i)
+    stmt = src.code[i]
+    (isa(stmt, Core.GotoNode) && stmt.label < i ||
+        isa(stmt, Core.GotoIfNot) && stmt.dest < i) && return true
+    if isa(stmt, Core.GotoIfNot) && i < length(src.code)
+        stmt = src.code[i+1]
+        isa(stmt, Core.GotoNode) && return stmt.label < i
+    end
+    return false
+end
+
 function markcf!(isrequired, src, i)
     stmt = src.code[i]
     @assert iscf(stmt)
@@ -783,7 +798,7 @@ function add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
             r = rng(bb)
             if block_internals_needed(isrequired, src, r)
                 needed[ibb] = true
-                # Check this blocks precessors and mark any that are just control-flow
+                # Check this block's precessors and mark any that are just control-flow
                 for p in bb.preds
                     r = rng(blocks[p])
                     if length(r) == 1 && iscf(src.code[r[1]])
@@ -814,21 +829,25 @@ function add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
                 while jbb != 0
                     if ispredecessor(blocks, jbb, ibb, empty!(cache))  # is post-dominator jbb also a predecessor of ibb? If so we have a loop.
                         pdbb = blocks[jbb]
-                        idxlast = rng(pdbb)[end]
-                        stmt = src.code[idxlast]
-                        if iscf(stmt)
-                            if isrequired[idxlast] != true
-                                _changed = true
-                                if isa(stmt, Core.ReturnNode) && isrequired[idxlast] != :exit
-                                    isrequired[idxlast] = :exit
-                                else
-                                    isrequired[idxlast] = true
-                                    if isa(stmt, Core.GotoIfNot) && idxlast < length(isrequired) && isrequired[idxlast+1] != true && iscf(src.code[idxlast+1])
-                                        isrequired[idxlast+1] = true
+                        r = rng(pdbb)
+                        # if block_internals_needed(isrequired, src, r)
+                            idxlast = rng(pdbb)[end]
+                            stmt = src.code[idxlast]
+                            if iscf(stmt) && jumps_back(src, idxlast)
+                                if isrequired[idxlast] != true
+                                    _changed = true
+                                    if isa(stmt, Core.ReturnNode) && isrequired[idxlast] != :exit
+                                        isrequired[idxlast] = :exit
+                                    else
+                                        isrequired[idxlast] = true
+                                        if isa(stmt, Core.GotoIfNot) && idxlast < length(isrequired) && isrequired[idxlast+1] != true && iscf(src.code[idxlast+1])
+                                            isrequired[idxlast+1] = true
+                                        end
                                     end
+                                    break
                                 end
                             end
-                        end
+                        # end
                     end
                     jbb = postdomtree.idoms_bb[jbb]
                 end
@@ -867,10 +886,14 @@ function add_control_flow!(isrequired, src, cfg, domtree, postdomtree)
                 # Mark the ipostdom's predecessors...
                 for k in blocks[ipbb].preds
                     idxlast = rng(blocks[k])[end]
-                    if iscf(src.code[idxlast])
+                    stmt = src.code[idxlast]
+                    if iscf(stmt)
                         if markcf!(isrequired, src, idxlast)
                             changed = true
                             ok = true
+                            if isa(stmt, Core.GotoIfNot) && idxlast < length(isrequired) && isrequired[idxlast+1] != true && iscf(src.code[idxlast+1])
+                                isrequired[idxlast+1] = true
+                            end
                         end
                     end
                 end
diff --git a/test/codeedges.jl b/test/codeedges.jl
index 23e6c3e..80fbc6f 100644
--- a/test/codeedges.jl
+++ b/test/codeedges.jl
@@ -488,7 +488,9 @@ end
         stmts = src.code
         edges = CodeEdges(m, src)
 
-        isrq = lines_required!(istypedef.(stmts), src, edges)
+        isrq = LoweredCodeUtils.initialize_isrequired(length(stmts))
+        copyto!(isrq, istypedef.(stmts))
+        lines_required!(isrq, src, edges)
         frame = Frame(m, src)
         selective_eval_fromstart!(frame, isrq, #=toplevel=#true)
 

From 459a02d1bd2897e61e4614f98ee75786926f61df Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Sun, 21 Jul 2024 12:04:30 -0500
Subject: [PATCH 8/8] explain

---
 test/signatures.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/signatures.jl b/test/signatures.jl
index 7beba46..901443e 100644
--- a/test/signatures.jl
+++ b/test/signatures.jl
@@ -414,9 +414,10 @@ bodymethtest5(x, y=Dict(1=>2)) = 5
     oldenv = Pkg.project().path
     try
         # we test with the old version of CBinding, let's do it in an isolated environment
+        # so we don't cause package conflicts with everything else
         Pkg.activate(; temp=true, io=devnull)
 
-        @info "Adding CBinding to the environment for test purposes"
+        @info "Adding CBinding 0.9.4 to the environment for test purposes"
         Pkg.add(; name="CBinding", version="0.9.4", io=devnull) # `@cstruct` isn't defined for v1.0 and above
 
         m = Module()