Handle Core._expr method bodies (#22)

Fixes timholy/Revise.jl#398

Author: timholy

Project.toml

@@ -12,7 +12,8 @@ julia = "1"
 
 [extras]
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["InteractiveUtils", "Test"]
+test = ["InteractiveUtils", "Parameters", "Test"]

src/LoweredCodeUtils.jl

@@ -365,7 +365,7 @@ end
 rename_framemethods!(frame::Frame) = rename_framemethods!(finish_and_return!, frame)
 
 """
-    pctop, isgen = find_corrected_name(frame, pc, name, parentname)
+    pctop, isgen = find_corrected_name(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.
@@ -374,16 +374,29 @@ 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(frame, pc, name, parentname)
+function find_corrected_name(@nospecialize(recurse), frame, pc, name, parentname)
     stmt = pc_expr(frame, pc)
     while true
+        pc0 = pc
         while !ismethod3(stmt)
             pc = next_or_nothing(frame, pc)
             pc === nothing && return nothing
             stmt = pc_expr(frame, pc)
         end
         body = stmt.args[3]
-        if stmt.args[1] != name
+        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)
                 if isexpr(bodystmt, :meta) && bodystmt.args[1] == :generated
@@ -435,7 +448,7 @@ 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(frame, pc, name, parentname)
+    nameinfo = find_corrected_name(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/runtests.jl

@@ -8,6 +8,7 @@ using Test
 @test isempty(detect_ambiguities(LoweredCodeUtils, JuliaInterpreter, Base, Core))
 
 module Lowering
+using Parameters
 struct Caller end
 struct Gen{T} end
 end
@@ -308,4 +309,19 @@ bodymethtest5(x, y=Dict(1=>2)) = 5
     empty!(signatures)
     methoddefs!(signatures, frame; define=false)
     @test Tuple{typeof(Lowering.CustomMS)} ∈ signatures
+
+    # https://github.com/timholy/Revise.jl/issues/398
+    ex = quote
+        @with_kw struct Items
+            n::Int
+            items::Vector{Int} = [i for i=1:n]
+        end
+    end
+    Core.eval(Lowering, ex)
+    frame = JuliaInterpreter.prepare_thunk(Lowering, ex)
+    dct = rename_framemethods!(frame)
+    ks = collect(filter(k->startswith(String(k), "#Items#"), keys(dct)))
+    @test length(ks) == 2
+    @test dct[ks[1]] == dct[ks[2]]
+    @test isdefined(Lowering, ks[1]) || isdefined(Lowering, ks[2])
 end