Add bodymethod

This allows you to find the body method for any input method.

Author: timholy

src/LoweredCodeUtils.jl

@@ -6,7 +6,7 @@ using JuliaInterpreter
 using JuliaInterpreter: @lookup, moduleof, pc_expr, step_expr!, is_global_ref, whichtt,
                         next_until!, finish_and_return!, nstatements, codelocation
 
-export signature, methoddef!, methoddefs!
+export signature, methoddef!, methoddefs!, bodymethod
 
 """
     iscallto(stmt, name)
@@ -441,6 +441,68 @@ function _methoddefs!(@nospecialize(recurse), signatures, frame::Frame, pc; defi
     return pc
 end
 
+"""
+    mbody = bodymethod(m::Method)
+
+Return the "body method" for a method `m`. `mbody` contains the code of the function body
+when `m` was defined.
+"""
+function bodymethod(mkw::Method)
+    function is_self_call(stmt, slotnames, argno=1)
+        if isa(stmt, Expr)
+            if stmt.head == :call
+                a = stmt.args[argno]
+                if isa(a, Core.SlotNumber)
+                    if slotnames[a.id] == Symbol("#self#")
+                        return true
+                    end
+                end
+            end
+        end
+        return false
+    end
+    m = mkw
+    local src
+    while true
+        framecode = JuliaInterpreter.get_framecode(m)
+        fakeargs = Any[nothing for i = 1:length(framecode.scope.nargs)]
+        frame = JuliaInterpreter.prepare_frame(framecode, fakeargs, isa(m.sig, UnionAll) ? sparam_ub(m) : Core.svec())
+        src = framecode.src
+        (length(src.code) > 1 && is_self_call(src.code[end-1], src.slotnames)) || break
+        # Build the optional arg, so we can get its type
+        pc = frame.pc
+        while pc < length(src.code) - 1
+            pc = step_expr!(frame)
+        end
+        val = pc > 1 ? frame.framedata.ssavalues[pc-1] : src.code[1].args[end]
+        sig = Tuple{Base.unwrap_unionall(m.sig).parameters..., typeof(val)}
+        m = whichtt(sig)
+    end
+    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 = any(i->is_self_call(stmt, src.slotnames, i), 1:length(stmt.args))
+        hasself || return m
+        f = f.value
+        mths = methods(f)
+        if length(mths) == 1
+            return first(mths)
+        end
+    end
+    return m
+end
+
+function sparam_ub(meth::Method)
+    typs = []
+    sig = meth.sig
+    while sig isa UnionAll
+        push!(typs, Symbol(sig.var.ub))
+        sig = sig.body
+    end
+    return Core.svec(typs...)
+end
+
 # precompilation
 
 if ccall(:jl_generating_output, Cint, ()) == 1

test/runtests.jl

@@ -11,6 +11,18 @@ struct Caller end
 struct Gen{T} end
 end
 
+bodymethtest0(x) = 0
+function bodymethtest0(x)
+    y = 2x
+    y + x
+end
+bodymethtest1(x, y=1; z="hello") = 1
+bodymethtest2(x, y=Dict(1=>2); z="hello") = 2
+bodymethtest3(x::T, y=Dict(1=>2); z="hello") where T<:AbstractFloat = 3
+# No kw but has optional args
+bodymethtest4(x, y=1) = 4
+bodymethtest5(x, y=Dict(1=>2)) = 5
+
 @testset "LoweredCodeUtils.jl" begin
     signatures = Set{Any}()
     newcode = CodeInfo[]
@@ -198,4 +210,13 @@ end
     empty!(signatures)
     methoddefs!(signatures, frame; define=false)
     @test !isempty(signatures)
+
+    for m in methods(bodymethtest0)
+        @test bodymethod(m) === m
+    end
+    @test startswith(String(bodymethod(first(methods(bodymethtest1))).name), "#")
+    @test startswith(String(bodymethod(first(methods(bodymethtest2))).name), "#")
+    @test startswith(String(bodymethod(first(methods(bodymethtest3))).name), "#")
+    @test bodymethod(first(methods(bodymethtest4))).nargs == 3  # one extra for #self#
+    @test bodymethod(first(methods(bodymethtest5))).nargs == 3
 end