Focus running-name lookup on just the caller (#45)

Fixes timholy/Revise.jl#420

Author: timholy

Project.toml

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

docs/src/api.md

@@ -29,7 +29,7 @@ LoweredCodeUtils.MethodInfo
 LoweredCodeUtils.identify_framemethod_calls
 LoweredCodeUtils.iscallto
 LoweredCodeUtils.getcallee
-LoweredCodeUtils.find_corrected_name
+LoweredCodeUtils.find_name_caller_sig
 LoweredCodeUtils.replacename!
 LoweredCodeUtils.Variable
 ```

src/signatures.jl

@@ -121,7 +121,7 @@ to methods defined in `frame` that occur within `frame`.
 
 `methodinfos` is a Dict of `name=>info` pairs, where `info` is a [`MethodInfo`](@ref).
 
-`selfcalls` is a list of `(linetop, linebody, callee, caller)` tuples that holds the location of
+`selfcalls` is a list of `SelfCall(linetop, linebody, callee, caller)` that holds the location of
 calls the methods defined in `frame`. `linetop` is the line in `frame` (top meaning "top level"),
 which will correspond to a 3-argument `:method` expression containing a `CodeInfo` body.
 `linebody` is the line within the `CodeInfo` body from which the call is made.
@@ -223,26 +223,27 @@ function callchain(selfcalls)
     return calledby
 end
 
-function set_to_running_name!(@nospecialize(recurse), replacements, frame, methodinfos, calledby, callee, caller)
+function set_to_running_name!(@nospecialize(recurse), replacements, frame, methodinfos, selfcall, calledby, callee, caller)
     if isa(caller, Symbol) && startswith(String(caller), '#')
         rep = get(replacements, caller, nothing)
         if rep === nothing
             parentcaller = get(calledby, caller, nothing)
             if parentcaller !== nothing
-                set_to_running_name!(recurse, replacements, frame, methodinfos, calledby, caller, parentcaller)
+                set_to_running_name!(recurse, replacements, frame, methodinfos, selfcall, calledby, caller, parentcaller)
             end
         else
             caller = rep
         end
     end
-    if isa(caller, Symbol)
-        mi = methodinfos[caller]
-        cname, _pc, _ = get_running_name(recurse, frame, mi.start, callee, get(replacements, caller, caller))
-    else
-        # For generated constructors (which have no name), we just assume they immediately follow their callee
-        mi = methodinfos[callee]
-        cname, _pc, _ = get_running_name(recurse, frame, mi.stop+1, callee, get(replacements, caller, caller))
+    # Back up to the beginning of the signature
+    pc = selfcall.linetop
+    stmt = pc_expr(frame, pc)
+    while pc > 1 && !ismethod1(stmt)
+        pc -= 1
+        stmt = pc_expr(frame, pc)
     end
+    @assert ismethod1(stmt)
+    cname, _pc, _ = get_running_name(recurse, frame, pc+1, callee, get(replacements, caller, caller))
     replacements[callee] = cname
     mi = methodinfos[cname] = methodinfos[callee]
     src = frame.framecode.src
@@ -270,8 +271,10 @@ function rename_framemethods!(@nospecialize(recurse), frame::Frame, methodinfos,
     replacements = Dict{Symbol,Symbol}()
     for (callee, caller) in calledby
         (!startswith(String(callee), '#') || haskey(replacements, callee)) && continue
+        idx = findfirst(sc->sc.callee === callee && sc.caller === caller, selfcalls)
+        idx === nothing && continue
         try
-            set_to_running_name!(recurse, replacements, frame, methodinfos, calledby, callee, caller)
+            set_to_running_name!(recurse, replacements, frame, methodinfos, selfcalls[idx], calledby, callee, caller)
         catch err
             @warn "skipping callee $callee (called by $caller) due to $err"
         end
@@ -297,7 +300,7 @@ end
 rename_framemethods!(frame::Frame) = rename_framemethods!(finish_and_return!, frame)
 
 """
-    pctop, isgen = find_corrected_name(recurse, frame, pc, name, parentname)
+    pctop, isgen = find_name_caller_sig(recurse, frame, pc, name, parentname)
 
 Scans forward from `pc` in `frame` until a method is found that calls `name`.
 `pctop` points to the beginning of that method's signature.
@@ -306,7 +309,7 @@ Scans forward from `pc` in `frame` until a method is found that calls `name`.
 Alternatively, this returns `nothing` if `pc` does not appear to point to either
 a keyword or generated method.
 """
-function find_corrected_name(@nospecialize(recurse), frame, pc, name, parentname)
+function find_name_caller_sig(@nospecialize(recurse), frame, pc, name, parentname)
     stmt = pc_expr(frame, pc)
     while true
         pc0 = pc
@@ -316,18 +319,6 @@ function find_corrected_name(@nospecialize(recurse), frame, pc, name, parentname
             stmt = pc_expr(frame, pc)
         end
         body = stmt.args[3]
-        if stmt.args[1] !== name && isa(body, SSAValue)
-            # OK, we can't skip all the stuff that might define the body
-            # See https://github.com/timholy/Revise.jl/issues/398
-            pc = pc0
-            stmt = pc_expr(frame, pc)
-            while !ismethod3(stmt)
-                pc = step_expr!(recurse, frame, stmt, true)
-                pc === nothing && return nothing
-                stmt = pc_expr(frame, pc)
-            end
-            body = @lookup(frame, stmt.args[3])
-        end
         if stmt.args[1] !== name && isa(body, CodeInfo)
             # This might be the GeneratedFunctionStub for a @generated method
             for (i, bodystmt) in enumerate(body.code)
@@ -379,8 +370,7 @@ function replacename!(args::Vector{Any}, pr)
 end
 
 function get_running_name(@nospecialize(recurse), frame, pc, name, parentname)
-    # Get the correct name (the one that's actively running)
-    nameinfo = find_corrected_name(recurse, frame, pc, name, parentname)
+    nameinfo = find_name_caller_sig(recurse, frame, pc, name, parentname)
     if nameinfo === nothing
         pc = skip_until(stmt->isexpr(stmt, :method, 3), frame, pc)
         pc = next_or_nothing(frame, pc)

test/signatures.jl

@@ -350,6 +350,8 @@ bodymethtest5(x, y=Dict(1=>2)) = 5
     @test length(ks) == 2
     @test dct[ks[1]] == dct[ks[2]]
     @test isdefined(Lowering, ks[1]) || isdefined(Lowering, ks[2])
+    nms = filter(sym->occursin(r"#Items#\d+#\d+", String(sym)), names(Lowering; all=true))
+    @test length(nms) == 1
 
     # https://github.com/timholy/Revise.jl/issues/422
     ex = :(@generated function fneg(x::T) where T<:LT{<:FloatingTypes}