implement a system for interpreting overlay methods

Implemented, being inspired by #680,
thinking it might be useful in the near future for JET's use case.

The overlay interpretertation is enabled by overloading
`JuliaInterpreter.method_table(recurse)`, but maybe the entire
`recurse`-overload mechanism itself needs to be aligned with the
`AbstractInterpreter` design and properly organized.

Since this interface will also be needed for JET,
I'll probably work on it soon.

Author: aviatesk

src/construct.jl

@@ -229,7 +229,9 @@ julia> argtypes
 Tuple{typeof(mymethod), Vector{Float64}}
 ```
 """
-function prepare_call(@nospecialize(f), allargs; enter_generated = false)
+function prepare_call(@nospecialize(f), allargs;
+                      enter_generated::Bool=false,
+                      method_table::Union{Nothing,MethodTable}=nothing)
     # Can happen for thunks created by generated functions
     if isa(f, Core.Builtin) || isa(f, Core.IntrinsicFunction)
         return nothing
@@ -247,13 +249,13 @@ function prepare_call(@nospecialize(f), allargs; enter_generated = false)
             return nothing
         end
     else
-        method = whichtt(argtypes)
+        method = whichtt(argtypes, method_table)
     end
     if method === nothing
         # Call it to generate the exact error
         return f(allargs[2:end]...)
     end
-    ret = prepare_framecode(method, argtypes; enter_generated=enter_generated)
+    ret = prepare_framecode(method, argtypes; enter_generated)
     # Exceptional returns
     if ret === nothing
         # The generator threw an error. Let's generate the same error by calling it.
@@ -580,7 +582,9 @@ Prepare all the information needed to execute a particular `:call` expression `e
 For example, try `JuliaInterpreter.determine_method_for_expr(:(\$sum([1,2])))`.
 See [`JuliaInterpreter.prepare_call`](@ref) for information about the outputs.
 """
-function determine_method_for_expr(expr; enter_generated = false)
+function determine_method_for_expr(expr::Expr;
+                                   enter_generated::Bool=false,
+                                   method_table::Union{Nothing,MethodTable}=nothing)
     f = to_function(expr.args[1])
     allargs = expr.args
     # Extract keyword args
@@ -589,7 +593,7 @@ function determine_method_for_expr(expr; enter_generated = false)
         kwargs = splice!(allargs, 2)::Expr
     end
     f, allargs = prepare_args(f, allargs, kwargs.args)
-    return prepare_call(f, allargs; enter_generated=enter_generated)
+    return prepare_call(f, allargs; enter_generated, method_table)
 end
 
 """
@@ -626,9 +630,11 @@ T = Float64
 
 See [`enter_call`](@ref) for a similar approach not based on expressions.
 """
-function enter_call_expr(expr; enter_generated = false)
+function enter_call_expr(expr::Expr;
+                         enter_generated::Bool=false,
+                         method_table::Union{Nothing,MethodTable}=nothing)
     clear_caches()
-    r = determine_method_for_expr(expr; enter_generated = enter_generated)
+    r = determine_method_for_expr(expr; enter_generated, method_table)
     if r !== nothing && !isa(r[1], Compiled)
         return prepare_frame(Base.front(r)...)
     end
@@ -666,7 +672,9 @@ would be created by the generator.
 
 See [`enter_call_expr`](@ref) for a similar approach based on expressions.
 """
-function enter_call(@nospecialize(finfo), @nospecialize(args...); kwargs...)
+function enter_call(@nospecialize(finfo), @nospecialize(args...);
+                    method_table::Union{Nothing,MethodTable}=nothing,
+                    kwargs...)
     clear_caches()
     if isa(finfo, Tuple)
         f = finfo[1]
@@ -680,7 +688,7 @@ function enter_call(@nospecialize(finfo), @nospecialize(args...); kwargs...)
     if isa(f, Core.Builtin) || isa(f, Core.IntrinsicFunction)
         error(f, " is a builtin or intrinsic")
     end
-    r = prepare_call(f, allargs; enter_generated=enter_generated)
+    r = prepare_call(f, allargs; enter_generated, method_table)
     if r !== nothing && !isa(r[1], Compiled)
         return prepare_frame(Base.front(r)...)
     end

src/interpret.jl

@@ -260,7 +260,9 @@ function evaluate_call!(interp::Interpreter, frame::Frame, call_expr::Expr, ente
         lenv === nothing && return framecode  # this was a Builtin
         fargs = fargs_pruned
     else
-        framecode, lenv = get_call_framecode(fargs, frame.framecode, frame.pc; enter_generated=enter_generated)
+        method_table = JuliaInterpreter.method_table(interp)
+        framecode, lenv = get_call_framecode(fargs, frame.framecode, frame.pc;
+                                             enter_generated, method_table)
         if lenv === nothing
             if isa(framecode, Compiled)
                 return native_call(fargs, frame)

src/localmethtable.jl

@@ -7,7 +7,9 @@ Return the framecode and environment for a call specified by `fargs = [f, args..
 `parentframecode` is the caller, and `idx` is the program-counter index.
 If possible, `framecode` will be looked up from the local method tables of `parentframe`.
 """
-function get_call_framecode(fargs::Vector{Any}, parentframe::FrameCode, idx::Int; enter_generated::Bool=false)
+function get_call_framecode(fargs::Vector{Any}, parentframe::FrameCode, idx::Int;
+                            enter_generated::Bool=false,
+                            method_table::Union{Nothing,MethodTable}=nothing)
     nargs = length(fargs)  # includes f as the first "argument"
     # Determine whether we can look up the appropriate framecode in the local method table
     if isassigned(parentframe.methodtables, idx)  # if this is the first call, this may not yet be set
@@ -60,7 +62,7 @@ function get_call_framecode(fargs::Vector{Any}, parentframe::FrameCode, idx::Int
     end
     # We haven't yet encountered this argtype combination and need to look it up by dispatch
     fargs[1] = f = to_function(fargs[1])
-    ret = prepare_call(f, fargs; enter_generated=enter_generated)
+    ret = prepare_call(f, fargs; enter_generated, method_table)
     ret === nothing && return invokelatest(f, fargs[2:end]...), nothing
     is_compiled = isa(ret[1], Compiled)
     local framecode

src/types.jl

@@ -39,6 +39,14 @@ redefined as a completely different type in v0.11 or later.
 const Compiled = NonRecursiveInterpreter # for backward compatibility
 Base.similar(::Compiled, sz) = Compiled()  # to support similar(stack, 0)
 
+"""
+    method_table(interpreter::Interpreter) -> mt::Union{Nothing,MethodTable}
+
+Configures the method table used for method lookups performed by the interpreter.
+Uses the global method table by default.
+"""
+method_table(::Interpreter) = nothing
+
 # Our own replacements for Core types. We need to do this to ensure we can tell the difference
 # between "data" (Core types) and "code" (our types) if we step into Core.Compiler
 struct SSAValue

src/utils.jl

@@ -31,16 +31,29 @@ end
 Like `which` except it operates on the complete tuple-type `tt`,
 and doesn't throw when there is no matching method.
 """
-function whichtt(@nospecialize(tt), mt::Union{Nothing, MethodTable} = nothing)
+function whichtt(@nospecialize(tt), mt::Union{Nothing,MethodTable}=nothing)
     # TODO: provide explicit control over world age? In case we ever need to call "old" methods.
-    # branch on https://github.com/JuliaLang/julia/pull/44515
-    # for now, code execution doesn't have the capability to use an overlayed method table,
-    # which is meant to be addressed in https://github.com/JuliaDebug/JuliaInterpreter.jl/pull/682.
-    match, _ = Core.Compiler._findsup(tt, mt, get_world_counter())
+    # TODO Use `CachedMethodTable` for better performance once `teh/worldage` is merged
+    match, _ = findsup_mt(tt, Base.get_world_counter(), mt)
     match === nothing && return nothing
     return match.method
 end
 
+@static if VERSION ≥ v"1.12-"
+using Base.Compiler: findsup_mt
+else
+function findsup_mt(@nospecialize(tt), world, method_table)
+    if method_table === nothing
+        table = Core.Compiler.InternalMethodTable(world)
+    elseif method_table isa Core.MethodTable
+        table = Core.Compiler.OverlayMethodTable(world, method_table)
+    else
+        table = method_table
+    end
+    return Core.Compiler.findsup(tt, table)
+end
+end
+
 instantiate_type_in_env(arg, spsig::UnionAll, spvals::Vector{Any}) =
     ccall(:jl_instantiate_type_in_env, Any, (Any, Any, Ptr{Any}), arg, spsig, spvals)
 

test/interpret.jl

@@ -975,3 +975,25 @@ end
     @test_throws "Invalid @interpret call" macroexpand(@__MODULE__, :(@interpret interp sin(42)))
     @test_throws "Invalid @interpret call" macroexpand(@__MODULE__, :(@interpret _interp_=RecursiveInterpreter() sin(42)))
 end
+
+function func_overlay end
+func_overlay(x) = sin(x)
+call_func_overlay(x) = func_overlay(x)
+Base.Experimental.@MethodTable ex_method_table
+Base.Experimental.@overlay ex_method_table func_overlay(x) = cos(x)
+struct OverlayInterpreter <: Interpreter end
+JuliaInterpreter.method_table(::OverlayInterpreter) = ex_method_table
+
+@testset "Interpret overlay method" begin
+    let frame = JuliaInterpreter.Frame(@__MODULE__, :(func_overlay(42.0)))
+        @test JuliaInterpreter.finish_and_return!(frame, true) == sin(42.0)
+    end
+    @test sin(42.0) == @interpret call_func_overlay(42.0)
+    let frame = JuliaInterpreter.Frame(@__MODULE__, :(func_overlay(42.0)))
+        @test JuliaInterpreter.finish_and_return!(OverlayInterpreter(), frame, true) == cos(42.0)
+    end
+    let frame = JuliaInterpreter.enter_call(func_overlay, 42.0; method_table=ex_method_table)
+        @test JuliaInterpreter.finish_and_return!(OverlayInterpreter(), frame) == cos(42.0)
+    end
+    @test cos(42.0) == @interpret interp=OverlayInterpreter() call_func_overlay(42.0)
+end