Merge pull request #34 from JuliaDebug/teh/apply_iterate

Fix handling of Core._apply_iterate

Author: timholy

Project.toml

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

src/LoweredCodeUtils.jl

@@ -47,6 +47,9 @@ if ccall(:jl_generating_output, Cint, ()) == 1
     @assert precompile(CodeEdges, (CodeInfo,))
     @assert precompile(add_links!, (Pair{Union{SSAValue,SlotNumber,NamedVar},Links}, Any, CodeLinks))
     @assert precompile(lines_required!, (Vector{Bool}, Set{NamedVar}, CodeInfo, CodeEdges))
+
+    precompile(Tuple{typeof(setindex!),Dict{Union{GlobalRef, Symbol},Links},Links,Symbol})
+    precompile(Tuple{typeof(setindex!),Dict{Union{GlobalRef, Symbol},Variable},Variable,Symbol})
 end
 
 end # module

src/signatures.jl

@@ -164,6 +164,16 @@ function identify_framemethod_calls(frame)
                         if isa(mkey, Symbol)
                             # Could be a GlobalRef but then it's outside frame
                             haskey(methodinfos, mkey) && push!(selfcalls, (linetop=i, linebody=j, callee=mkey, caller=key))
+                        elseif is_global_ref(mkey, Core, isdefined(Core, :_apply_iterate) ? :_apply_iterate : :_apply)
+                            ssaref = mstmt.args[end-1]
+                            if isa(ssaref, JuliaInterpreter.SSAValue)
+                                id = ssaref.id
+                                has_self_call(msrc, msrc.code[id]) || continue
+                            end
+                            mkey = mstmt.args[end-2]
+                            if isa(mkey, Symbol)
+                                haskey(methodinfos, mkey) && push!(selfcalls, (linetop=i, linebody=j, callee=mkey, caller=key))
+                            end
                         end
                     elseif isexpr(mstmt, :meta) && mstmt.args[1] == :generated
                         newex = mstmt.args[2]
@@ -529,6 +539,15 @@ function is_self_call(@nospecialize(stmt), slotnames, argno=1)
     return false
 end
 
+function has_self_call(src, stmt::Expr)
+    # Check that it has a #self# call
+    hasself = false
+    for i = 2:length(stmt.args)
+        hasself |= is_self_call(stmt, src.slotnames, i)
+    end
+    return hasself
+end
+
 """
     mbody = bodymethod(m::Method)
 
@@ -556,12 +575,16 @@ function bodymethod(mkw::Method)
     length(src.code) > 1 || return m
     stmt = src.code[end-1]
     if isexpr(stmt, :call) && (f = stmt.args[1]; isa(f, QuoteNode))
-        # Check that it has a #self# call
-        hasself = false
-        for i = 2:length(stmt.args)
-            hasself |= is_self_call(stmt, src.slotnames, i)
+        if f.value === (isdefined(Core, :_apply_iterate) ? Core._apply_iterate : Core._apply)
+            ssaref = stmt.args[end-1]
+            if isa(ssaref, JuliaInterpreter.SSAValue)
+                id = ssaref.id
+                has_self_call(src, src.code[id]) || return m
+            end
+            f = stmt.args[end-2]
+        else
+            has_self_call(src, stmt) || return m
         end
-        hasself || return m
         f = f.value
         mths = methods(f)
         if length(mths) == 1
@@ -570,4 +593,3 @@ function bodymethod(mkw::Method)
     end
     return m
 end
-

src/utils.jl

@@ -13,7 +13,8 @@ function iscallto(@nospecialize(stmt), name)
     if isexpr(stmt, :call)
         a = stmt.args[1]
         a === name && return true
-        return is_global_ref(a, Core, :_apply) && stmt.args[2] === name
+        is_global_ref(a, Core, :_apply) && stmt.args[2] === name && return true
+        is_global_ref(a, Core, :_apply_iterate) && stmt.args[3] === name && return true
     end
     return false
 end
@@ -27,6 +28,7 @@ function getcallee(@nospecialize(stmt))
     if isexpr(stmt, :call)
         a = stmt.args[1]
         is_global_ref(a, Core, :_apply) && return stmt.args[2]
+        is_global_ref(a, Core, :_apply_iterate) && return stmt.args[3]
         return a
     end
     error(stmt, " is not a call expression")

test/signatures.jl

@@ -260,6 +260,26 @@ bodymethtest5(x, y=Dict(1=>2)) = 5
     # Issue in https://github.com/timholy/CodeTracking.jl/pull/48
     mbody = bodymethod(m)
     @test mbody != m && mbody.file != :none
+    # varargs keyword methods
+    m = which(Base.print_shell_escaped, (IO, AbstractString))
+    mbody = bodymethod(m)
+    @test isa(mbody, Method) && mbody != m
+
+    ex = quote
+        function print_shell_escaped(io::IO, cmd::AbstractString, args::AbstractString...;
+                                     special::AbstractString="")
+            print_shell_word(io, cmd, special)
+            for arg in args
+                print(io, ' ')
+                print_shell_word(io, arg, special)
+            end
+        end
+    end
+    frame = JuliaInterpreter.prepare_thunk(Base, ex)
+    rename_framemethods!(frame)
+    empty!(signatures)
+    methoddefs!(signatures, frame; define=false)
+    @test length(signatures) >= 3
 
     ex = :(typedsig(x) = 1)
     frame = JuliaInterpreter.prepare_thunk(Lowering, ex)