make norequire configurable

This PR makes `lines_required!`'s `norequire` logic more configurable.
That means the `exclude_named_typedefs` option is now got abstracted,
and each consumer can implement its own strategy to escape from the
required statement completion by control flow traversal.

The motivation for this change is we usually want to respect a control
flow in  general context, but some consumer may not want that.
Especially, JET doesn't want to interpret all the statements within a
`try/catch` block, but just select those involved with a method
definition.
(issue: <aviatesk/JET.jl#150>)

For example, `lines_required!` selects statements in the snippet below
as JET expects:
```julia
for fname in (:foo, :bar, :baz)
    @eval begin
        @inline ($(Symbol("is", fname)))(a) = a === $(QuoteNode(fname))
    end
end
```

, but in the example below, `lines_required` selects "too much"
statements and we need a customized `norequire`:
```julia
try
    foo(a) = sum(a) # should be selected (selected initially)

    foo("julia") # shouldn't be selected, but `lines_required` will 
select this
catch err
    err # shouldn't be selected, but `lines_required` will select this
end
```

Here is an example usage of this customizable `norequire` logic: 
<aviatesk/JET.jl#152>

---

One downside of this change is that now we need to performa the basic
block traversal twice when using `norequire = 
exclude_named_typedefs(src, edges)`.
As far as I confirmed, this computation would never be a performance
bottleneck, and thus this change hopefully won't hurt the performance.
I tried to profile the time with the following snippet:
```julia
function select_statements(n, src)
    for _ in 1:n
        stmts = src.code

        isrq = rand(Bool, length(stmts))

        edges = CodeEdges(src)
        norequire = LoweredCodeUtils.exclude_named_typedefs(src, edges)
        lines_required!(isrq, src, edges, norequire)
    end
end

src = code_lowered(...)
@profiler select_statements(100, src)
```

Author: aviatesk

Project.toml

@@ -1,7 +1,7 @@
 name = "LoweredCodeUtils"
 uuid = "6f1432cf-f94c-5a45-995e-cdbf5db27b0b"
 authors = ["Tim Holy <[email protected]>"]
-version = "1.2.9"
+version = "1.3.0"
 
 [deps]
 JuliaInterpreter = "aa1ae85d-cabe-5617-a682-6adf51b2e16a"

src/codeedges.jl

@@ -567,44 +567,61 @@ 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::Union{Symbol,GlobalRef}, src::CodeInfo, edges::CodeEdges; kwargs...)
+function lines_required(obj::Union{Symbol,GlobalRef}, src::CodeInfo, edges::CodeEdges, args...)
     isrequired = falses(length(edges.preds))
     objs = Set{Union{Symbol,GlobalRef}}([obj])
-    return lines_required!(isrequired, objs, src, edges; kwargs...)
+    return lines_required!(isrequired, objs, src, edges, args...)
 end
 
-function lines_required(idx::Int, src::CodeInfo, edges::CodeEdges; kwargs...)
+function lines_required(idx::Int, src::CodeInfo, edges::CodeEdges, args...)
     isrequired = falses(length(edges.preds))
     isrequired[idx] = true
     objs = Set{Union{Symbol,GlobalRef}}()
-    return lines_required!(isrequired, src, edges; kwargs...)
+    return lines_required!(isrequired, objs, src, edges, args...)
 end
 
 """
-    lines_required!(isrequired::AbstractVector{Bool}, src::CodeInfo, edges::CodeEdges; exclude_named_typedefs::Bool=false)
+    lines_required!(isrequired::AbstractVector{Bool}, src::CodeInfo, edges::CodeEdges, norequire = ())
 
 Like `lines_required`, but where `isrequired[idx]` has already been set to `true` for all statements
 that you know you need to evaluate. All other statements should be marked `false` at entry.
 On return, the complete set of required statements will be marked `true`.
 
-Use `exclude_named_typedefs=true` if you're extracting method signatures and not evaluating new definitions.
+`norequire` specifies statements (represented as iterator of `Int`s) that 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{Bool}, src::CodeInfo, edges::CodeEdges, norequire = ())
     objs = Set{Union{Symbol,GlobalRef}}()
-    return lines_required!(isrequired, objs, src, edges; kwargs...)
+    return lines_required!(isrequired, objs, src, edges, norequire)
 end
 
-function lines_required!(isrequired::AbstractVector{Bool}, objs, src::CodeInfo, edges::CodeEdges; exclude_named_typedefs::Bool=false)
+function exclude_named_typedefs(src::CodeInfo, edges::CodeEdges)
+    norequire = BitSet()
+    i = 1
+    nstmts = length(src.code)
+    while i <= nstmts
+        stmt = rhs(src.code[i])
+        if istypedef(stmt) && !isanonymous_typedef(stmt::Expr)
+            r = typedef_range(src, i)
+            pushall!(norequire, r)
+            i = last(r)+1
+        else
+            i += 1
+        end
+    end
+    return norequire
+end
+
+function lines_required!(isrequired::AbstractVector{Bool}, objs, src::CodeInfo, edges::CodeEdges, norequire = ())
     # Do a traveral of "numbered" predecessors
     # We'll mostly use generic graph traversal to discover all the lines we need,
     # but structs are in a bit of a different category (especially on Julia 1.5+).
     # It's easiest to discover these at the beginning.
-    # Moreover, if we're excluding named type definitions, we'll add them to `norequire`
-    # to prevent them from being marked.
     typedef_blocks, typedef_names = UnitRange{Int}[], Symbol[]
-    norequire = BitSet()
-    nstmts = length(src.code)
     i = 1
+    nstmts = length(src.code)
     while i <= nstmts
         stmt = rhs(src.code[i])
         if istypedef(stmt) && !isanonymous_typedef(stmt::Expr)
@@ -618,9 +635,6 @@ function lines_required!(isrequired::AbstractVector{Bool}, objs, src::CodeInfo,
             isa(name, Symbol) || @show src i r stmt
             push!(typedef_names, name::Symbol)
             i = last(r)+1
-            if exclude_named_typedefs && !isanonymous_typedef(stmt)
-                pushall!(norequire, r)
-            end
         else
             i += 1
         end
@@ -641,19 +655,22 @@ function lines_required!(isrequired::AbstractVector{Bool}, objs, src::CodeInfo,
     iter = 0
     while changed
         changed = false
+
         # Handle ssa predecessors
         for idx = 1:nstmts
             if isrequired[idx]
                 changed |= add_preds!(isrequired, idx, edges, norequire)
             end
         end
+
         # Handle named dependencies
         for (obj, uses) in edges.byname
             obj ∈ objs && continue
             if any(view(isrequired, uses.succs))
                 changed |= add_obj!(isrequired, objs, obj, edges, norequire)
             end
         end
+
         # Add control-flow. For any basic block with an evaluated statement inside it,
         # check to see if the block has any successors, and if so mark that block's exit statement.
         # Likewise, any preceding blocks should have *their* exit statement marked.
@@ -684,6 +701,7 @@ function lines_required!(isrequired::AbstractVector{Bool}, objs, src::CodeInfo,
                 end
             end
         end
+
         # So far, everything is generic graph traversal. Now we add some domain-specific information.
         # New struct definitions, including their constructors, get spread out over many
         # statements. If we're evaluating any of them, it's important to evaluate *all* of them.
@@ -833,7 +851,7 @@ Mark each line of code with its requirement status.
 function print_with_code(io::IO, src::CodeInfo, isrequired::AbstractVector{Bool})
     nd = ndigits(length(isrequired))
     preprint(::IO) = nothing
-    preprint(io::IO, idx::Int) = print(io, lpad(idx, nd), ' ', isrequired[idx] ? "t " : "f ")
+    preprint(io::IO, idx::Int) = (c = isrequired[idx]; printstyled(io, lpad(idx, nd), ' ', c ? "t " : "f "; color = c ? :cyan : :plain))
     postprint(::IO) = nothing
     postprint(io::IO, idx::Int, bbchanged::Bool) = nothing
 

test/codeedges.jl

@@ -1,6 +1,6 @@
 using LoweredCodeUtils
 using LoweredCodeUtils.JuliaInterpreter
-using LoweredCodeUtils: callee_matches, istypedef
+using LoweredCodeUtils: callee_matches, istypedef, exclude_named_typedefs
 using JuliaInterpreter: is_global_ref, is_quotenode
 using Test
 
@@ -22,7 +22,7 @@ function hastrackedexpr(stmt; heads=LoweredCodeUtils.trackedheads)
     return false, haseval
 end
 
-function minimal_evaluation(predicate, src::Core.CodeInfo, edges::CodeEdges; kwargs...)
+function minimal_evaluation(predicate, src::Core.CodeInfo, edges::CodeEdges, args...)
     isrequired = fill(false, length(src.code))
     for (i, stmt) in enumerate(src.code)
         if !isrequired[i]
@@ -33,7 +33,7 @@ function minimal_evaluation(predicate, src::Core.CodeInfo, edges::CodeEdges; kwa
         end
     end
     # All tracked expressions are marked. Now add their dependencies.
-    lines_required!(isrequired, src, edges; kwargs...)
+    lines_required!(isrequired, src, edges, args...)
     return isrequired
 end
 
@@ -262,7 +262,7 @@ end
     frame = Frame(ModEval, ex)
     src = frame.framecode.src
     edges = CodeEdges(src)
-    isrequired = minimal_evaluation(stmt->(LoweredCodeUtils.ismethod3(stmt),false), src, edges; exclude_named_typedefs=true)  # initially mark only the constructor
+    isrequired = minimal_evaluation(stmt->(LoweredCodeUtils.ismethod3(stmt),false), src, edges, exclude_named_typedefs(src, edges))  # initially mark only the constructor
     bbs = Core.Compiler.compute_basic_blocks(src.code)
     for (iblock, block) in enumerate(bbs.blocks)
         r = LoweredCodeUtils.rng(block)
@@ -301,7 +301,7 @@ end
     src = thk.args[1]
     edges = CodeEdges(src)
     idx = findfirst(stmt->Meta.isexpr(stmt, :method), src.code)
-    lr = lines_required(idx, src, edges; exclude_named_typedefs=true)
+    lr = lines_required(idx, src, edges, exclude_named_typedefs(src, edges))
     idx = findfirst(stmt->Meta.isexpr(stmt, :(=)) && Meta.isexpr(stmt.args[2], :call) && is_global_ref(stmt.args[2].args[1], Core, :Box), src.code)
     @test lr[idx]
     # but make sure we don't break primitivetype & abstracttype (https://github.com/timholy/Revise.jl/pull/611)