Improve detection of struct dependencies

Author: timholy

src/codeedges.jl

@@ -192,18 +192,34 @@ function namedkeys(cl::CodeLinks)
 end
 
 function direct_links!(cl::CodeLinks, src::CodeInfo)
+    # Utility for when a stmt itself contains a CodeInfo
+    function add_inner!(cl::CodeLinks, icl::CodeLinks, idx)
+        for (name, _) in icl.nameassigns
+            assigns = get(cl.nameassigns, name, nothing)
+            if assigns === nothing
+                cl.nameassigns[name] = assigns = Int[]
+            end
+            push!(assigns, idx)
+        end
+        for (name, _) in icl.namesuccs
+            succs = get(cl.namesuccs, name, nothing)
+            if succs === nothing
+                cl.namesuccs[name] = succs = Links()
+            end
+            push!(succs.ssas, idx)
+        end
+    end
+
     for (i, stmt) in enumerate(src.code)
         if isexpr(stmt, :thunk) && isa(stmt.args[1], CodeInfo)
             icl = CodeLinks(stmt.args[1])
-            for (name, _) in icl.nameassigns
-                assign = get(cl.nameassigns, name, nothing)
-                if assign === nothing
-                    cl.nameassigns[name] = assign = Int[]
-                end
-                push!(assign, i)
-            end
+            add_inner!(cl, icl, i)
             continue
         elseif isa(stmt, Expr) && stmt.head ∈ trackedheads
+            if stmt.head === :method && length(stmt.args) === 3 && isa(stmt.args[3], CodeInfo)
+                icl = CodeLinks(stmt.args[3])
+                add_inner!(cl, icl, i)
+            end
             name = stmt.args[1]
             if isa(name, Symbol)
                 assign = get(cl.nameassigns, name, nothing)
@@ -511,6 +527,27 @@ function postprint_lineedges(io::IO, idx::Int, edges::CodeEdges, bbchanged::Bool
     return nothing
 end
 
+function terminal_preds(i::Int, edges::CodeEdges)
+    function terminal_preds!(s, j, edges, covered)
+        j ∈ covered && return s
+        push!(covered, j)
+        preds = edges.preds[j]
+        if isempty(preds)
+            push!(s, j)
+        else
+            for p in preds
+                terminal_preds!(s, p, edges, covered)
+            end
+        end
+        return s
+    end
+    s, covered = BitSet(), BitSet()
+    push!(covered, i)
+    for p in edges.preds[i]
+        terminal_preds!(s, p, edges, covered)
+    end
+    return s
+end
 
 """
     isrequired = lines_required(obj::Union{Symbol,GlobalRef}, src::CodeInfo, edges::CodeEdges)
@@ -551,95 +588,153 @@ end
 function lines_required!(isrequired::AbstractVector{Bool}, objs, src::CodeInfo, edges::CodeEdges)
     # Do a traveral of "numbered" predecessors
     function add_preds!(isrequired, idx, edges::CodeEdges)
+        changed = false
         preds = edges.preds[idx]
         for p in preds
             isrequired[p] && continue
             isrequired[p] = true
+            changed = true
             add_preds!(isrequired, p, edges)
         end
-        return isrequired
+        return changed
     end
     function add_succs!(isrequired, idx, edges::CodeEdges, succs)
+        changed = false
         for p in succs
             isrequired[p] && continue
             isrequired[p] = true
+            changed = true
             add_succs!(isrequired, p, edges, edges.succs[p])
         end
-        return isrequired
+        return changed
+    end
+    function add_obj!(isrequired, objs, obj, edges::CodeEdges)
+        changed = false
+        for d in edges.byname[obj].assigned
+            isrequired[d] || add_preds!(isrequired, d, edges)
+            isrequired[d] = true
+            changed = true
+        end
+        push!(objs, obj)
+        return changed
     end
 
+    objsnew = Set{Union{Symbol,GlobalRef}}()
+    for obj in objs
+        add_obj!(isrequired, objsnew, obj, edges)
+    end
+    objs = objsnew
     bbs = Core.Compiler.compute_basic_blocks(src.code)  # needed for control-flow analysis
     changed = true
     iter = 0
     while changed
         changed = false
-        # Add "named" object dependencies
-        for obj in objs
-            def = edges.byname[obj].assigned
-            if !all(i->isrequired[i], def)
-                changed = true
-                for d in def
-                    isrequired[d] = true
-                    add_preds!(isrequired, d, edges)
-                    if isexpr(src.code[d], :thunk) && startswith(String(obj), '#')
-                        # For anonymous types, we also want their associated methods
-                        add_succs!(isrequired, d, edges, edges.byname[obj].succs)
-                    end
-                end
-            end
-        end
-        # Add "numbered" dependencies
+        # Handle ssa predecessors
         for idx = 1:length(isrequired)
             if isrequired[idx]
-                preds = edges.preds[idx]
-                if !all(i->isrequired[i], preds)
-                    changed = true
-                    isrequired[preds] .= true
-                end
+                changed |= add_preds!(isrequired, idx, edges)
+            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)
             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.
-        for (i, bb) in enumerate(bbs.blocks)
+        for (ibb, bb) in enumerate(bbs.blocks)
             r = rng(bb)
             if any(view(isrequired, r))
                 # if !isempty(bb.succs)
-                if i != length(bbs.blocks)
+                if ibb != length(bbs.blocks)
                     idxlast = r[end]
                     changed |= !isrequired[idxlast]
                     isrequired[idxlast] = true
                 end
-                for ibb in bb.preds
-                    rpred = rng(bbs.blocks[ibb])
+                for ibbp in bb.preds
+                    rpred = rng(bbs.blocks[ibbp])
                     idxlast = rpred[end]
                     changed |= !isrequired[idxlast]
                     isrequired[idxlast] = true
                 end
-                for ibb in bb.succs
-                    ibb == length(bbs.blocks) && continue
-                    rpred = rng(bbs.blocks[ibb])
+                for ibbs in bb.succs
+                    ibbs == length(bbs.blocks) && continue
+                    rpred = rng(bbs.blocks[ibbs])
                     idxlast = rpred[end]
                     changed |= !isrequired[idxlast]
                     isrequired[idxlast] = true
                 end
             end
         end
-        # In preparation for the next round, add any new named objects
-        # required by these dependencies
-        for (obj, uses) in edges.byname
-            obj ∈ objs && continue
-            if any(i->isrequired[i], uses.succs)
-                push!(objs, obj)
-                changed = true
+        # 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.
+        for (idx, stmt) in enumerate(src.code)
+            isrequired[idx] || continue
+            if isexpr(stmt, :(=))
+                stmt = stmt.args[2]
+            end
+            # Is this a struct definition?
+            if (isa(stmt, Expr) && stmt.head ∈ structheads) ||                            # < Julia 1.5
+               (isexpr(stmt, :call) && callee_matches(stmt.args[1], Core, :_structtype))  # >= Julia 1.5
+                stmt = stmt::Expr
+                name = stmt.args[stmt.head === :call ? 3 : 1]
+                if isa(name, QuoteNode)
+                    name = name.value
+                end
+                name = name::NamedVar
+                # Some lines we need have been marked as successors of this name
+                for d in edges.byname[name].succs
+                    stmt2 = src.code[d]
+                    if isa(stmt2, Expr)
+                        head = stmt2.head
+                        if head === :method || head === :global || head === :const
+                            changed |= !isrequired[d]
+                            isrequired[d] = true
+                        end
+                    end
+                    # Julia 1.5+: others are successor of a slotnum->ssa load
+                    if isslotnum(stmt2)
+                        for s in edges.succs[d]
+                            stmt3 = src.code[s]
+                            if isexpr(stmt3, :call) && (callee_matches(stmt3.args[1], Core, :_setsuper!) ||
+                                                        callee_matches(stmt3.args[1], Core, :_typebody!))
+                                changed |= !isrequired[s]
+                                isrequired[s] = true
+                            end
+                        end
+                    end
+                    # Julia 1.5+: for non-parametric types, the Core._setsuper! call happens without the slotname->ssa load
+                    if isexpr(stmt2, :call) && callee_matches((stmt2::Expr).args[1], Core, :_setsuper!)
+                        changed |= !isrequired[d]
+                        isrequired[d] = true
+                    end
+                end
+            end
+            # Anonymous functions may not yet include the method definition
+            if isexpr(stmt, :thunk) && isanonymous_typedef(stmt.args[1])
+                i = idx + 1
+                while i <= length(src.code) && !ismethod3(src.code[i])
+                    i += 1
+                end
+                if i <= length(src.code) && src.code[i].args[1] == false
+                    tpreds = terminal_preds(i, edges)
+                    if minimum(tpreds) == idx
+                        changed |= !isrequired[i]
+                        isrequired[i] = true
+                    end
+                end
             end
         end
         iter += 1  # just for diagnostics
     end
     return isrequired
 end
 
-function caller_matches(f, mod, sym)
+function callee_matches(f, mod, sym)
     is_global_ref(f, mod, sym) && return true
     if isdefined(mod, sym)
         is_quotenode(f, getfield(mod, sym)) && return true

test/codeedges.jl

@@ -1,27 +1,17 @@
 using LoweredCodeUtils
 using LoweredCodeUtils.JuliaInterpreter
-using LoweredCodeUtils: caller_matches
+using LoweredCodeUtils: callee_matches
 using JuliaInterpreter: is_global_ref, is_quotenode
 using Test
 
-# # Utilities for finding statements corresponding to particular calls
-# function callpredicate(@nospecialize(stmt), fsym)
-#     isa(stmt, Expr) || return false
-#     head = stmt.head
-#     if head === :(=)
-#         return callpredicate(stmt.args[2], fsym)
-#     end
-#     return LoweredCodeUtils.iscallto(stmt, fsym)
-# end
-
 function hastrackedexpr(stmt; heads=LoweredCodeUtils.trackedheads)
     haseval = false
     if isa(stmt, Expr)
         if stmt.head === :call
             f = stmt.args[1]
-            haseval = f === :eval || (caller_matches(f, Base, :getproperty) && is_quotenode(stmt.args[2], :eval))
-            caller_matches(f, Core, :_typebody!) && return true, haseval
-            caller_matches(f, Core, :_setsuper!) && return true, haseval
+            haseval = f === :eval || (callee_matches(f, Base, :getproperty) && is_quotenode(stmt.args[2], :eval))
+            callee_matches(f, Core, :_typebody!) && return true, haseval
+            callee_matches(f, Core, :_setsuper!) && return true, haseval
             f === :include && return true, haseval
         elseif stmt.head === :thunk
             any(s->any(hastrackedexpr(s; heads=heads)), stmt.args[1].code) && return true, haseval
@@ -32,32 +22,19 @@ function hastrackedexpr(stmt; heads=LoweredCodeUtils.trackedheads)
     return false, haseval
 end
 
-function minimal_evaluation(src::Core.CodeInfo, edges::CodeEdges)
-    musteval = fill(false, length(src.code))
+function minimal_evaluation(predicate, src::Core.CodeInfo, edges::CodeEdges)
+    isrequired = fill(false, length(src.code))
     for (i, stmt) in enumerate(src.code)
-        if !musteval[i]
-            musteval[i], haseval = hastrackedexpr(stmt)
+        if !isrequired[i]
+            isrequired[i], haseval = predicate(stmt)
             if haseval
-                musteval[edges.succs[i]] .= true
+                isrequired[edges.succs[i]] .= true
             end
         end
     end
     # All tracked expressions are marked. Now add their dependencies.
-    lines_required!(musteval, src, edges)
-    # Struct definitions likely omitted const/global. Look for them via name.
-    for (name, var) in edges.byname
-        if !isempty(var.assigned) && any(i->musteval[i], var.succs)
-            foreach(var.succs) do i
-                stmt = src.code[i]
-                if isa(stmt, Expr)
-                    if stmt.head === :global || stmt.head === :const
-                        musteval[i] = true
-                    end
-                end
-            end
-        end
-    end
-    return musteval
+    lines_required!(isrequired, src, edges)
+    return isrequired
 end
 
 function allmissing(mod::Module, names)
@@ -209,18 +186,40 @@ end
     # Check that the StructParent name is discovered everywhere it is used
     var = edges.byname[:StructParent]
     @test var.preds[end] ∈ var.succs
-    isrequired = minimal_evaluation(src, edges)
+    isrequired = minimal_evaluation(hastrackedexpr, src, edges)
     selective_eval_fromstart!(frame, isrequired, true)
     @test supertype(ModSelective.StructParent) === AbstractArray
     # Also check redefinition (it's OK when the definition doesn't change)
     Core.eval(ModEval, ex)
     frame = JuliaInterpreter.prepare_thunk(ModEval, ex)
     src = frame.framecode.src
     edges = CodeEdges(src)
-    isrequired = minimal_evaluation(src, edges)
+    isrequired = minimal_evaluation(hastrackedexpr, src, edges)
     selective_eval_fromstart!(frame, isrequired, true)
     @test supertype(ModEval.StructParent) === AbstractArray
 
+    # Finding all dependencies in a struct definition
+    # Nonparametric
+    ex = :(struct NoParam end)
+    frame = JuliaInterpreter.prepare_thunk(ModSelective, ex)
+    src = frame.framecode.src
+    edges = CodeEdges(src)
+    isrequired = minimal_evaluation(stmt->(LoweredCodeUtils.ismethod3(stmt)&&stmt.args[1]===:NoParam,false), src, edges)  # initially mark only the constructor
+    selective_eval_fromstart!(frame, isrequired, true)
+    @test isa(ModSelective.NoParam(), ModSelective.NoParam)
+    # Parametric
+    ex = quote
+        struct Struct{T} <: StructParent{T,1}
+            x::Vector{T}
+        end
+    end
+    frame = JuliaInterpreter.prepare_thunk(ModSelective, ex)
+    src = frame.framecode.src
+    edges = CodeEdges(src)
+    isrequired = minimal_evaluation(stmt->(LoweredCodeUtils.ismethod3(stmt)&&stmt.args[1]===:Struct,false), src, edges)  # initially mark only the constructor
+    selective_eval_fromstart!(frame, isrequired, true)
+    @test isa(ModSelective.Struct([1,2,3]), ModSelective.Struct{Int})
+
     # Anonymous functions
     ex = :(max_values(T::Union{map(X -> Type{X}, Base.BitIntegerSmall_types)...}) = 1 << (8*sizeof(T)))
     frame = JuliaInterpreter.prepare_thunk(ModSelective, ex)