retry #62, make lines_required!'s norequire logic more configurable: (#64)

Author: aviatesk

Project.toml

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

src/codeedges.jl

@@ -577,34 +577,52 @@ function lines_required(idx::Int, src::CodeInfo, edges::CodeEdges; kwargs...)
     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; kwargs...)
 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` keyword argument 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...)
     objs = Set{Union{Symbol,GlobalRef}}()
     return lines_required!(isrequired, objs, src, edges; kwargs...)
 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 +636,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 +656,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 +702,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 +852,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
 
@@ -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; norequire=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; norequire=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)