implement a system for interpreting overlay methods (#682)

* change the `recurse` interface to `AbstractInterpreter`-like interface

Align the `recurse` argument to something like the base Compiler's
`AbstractInterpreter` and make JuliaInterpreter routines overloadable
properly.

This change is quite breaking (thus bumping the minor version of this
package), but necessary to enhance the customizability of JI.
For example, it will make it easier to add changes like
#682 in a nicer way, but also should
enable better designs in packages such as Revise and JET.

* 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

docs/src/dev_reference.md

@@ -13,6 +13,7 @@ JuliaInterpreter.Interpreter
 JuliaInterpreter.RecursiveInterpreter
 JuliaInterpreter.NonRecursiveInterpreter
 JuliaInterpreter.Compiled
+JuliaInterpreter.method_table
 ```
 
 ## Frame creation

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