Fix local method tables for enter_generated

Author: timholy

src/interpret.jl

@@ -173,7 +173,8 @@ function bypass_builtins(frame, call_expr, pc)
     return nothing
 end
 
-function evaluate_call_compiled!(::Compiled, frame::Frame, call_expr::Expr)
+function evaluate_call_compiled!(::Compiled, frame::Frame, call_expr::Expr; enter_generated::Bool=false)
+    # @assert !enter_generated
     pc = frame.pc
     ret = bypass_builtins(frame, call_expr, pc)
     isa(ret, Some{Any}) && return ret.value
@@ -185,7 +186,7 @@ function evaluate_call_compiled!(::Compiled, frame::Frame, call_expr::Expr)
     return f(fargs...)
 end
 
-function evaluate_call_recurse!(@nospecialize(recurse), frame::Frame, call_expr::Expr)
+function evaluate_call_recurse!(@nospecialize(recurse), frame::Frame, call_expr::Expr; enter_generated::Bool=false)
     pc = frame.pc
     ret = bypass_builtins(frame, call_expr, pc)
     isa(ret, Some{Any}) && return ret.value
@@ -198,7 +199,7 @@ function evaluate_call_recurse!(@nospecialize(recurse), frame::Frame, call_expr:
         err = length(fargs) > 1 ? fargs[2] : frame.framedata.last_exception[]
         throw(err)
     end
-    framecode, lenv = get_call_framecode(fargs, frame.framecode, frame.pc)
+    framecode, lenv = get_call_framecode(fargs, frame.framecode, frame.pc; enter_generated=enter_generated)
     if lenv === nothing
         if isa(framecode, Compiled)
             popfirst!(fargs)  # now it's really just `args`
@@ -231,9 +232,9 @@ The first causes it to be executed using Julia's normal dispatch (compiled code)
 whereas the second recurses in via the interpreter.
 `recurse` has a default value of [`JuliaInterpreter.finish_and_return!`](@ref).
 """
-evaluate_call!(::Compiled, frame::Frame, call_expr::Expr) = evaluate_call_compiled!(Compiled(), frame, call_expr)
-evaluate_call!(@nospecialize(recurse), frame::Frame, call_expr::Expr) = evaluate_call_recurse!(recurse, frame, call_expr)
-evaluate_call!(frame::Frame, call_expr::Expr) = evaluate_call!(finish_and_return!, frame, call_expr)
+evaluate_call!(::Compiled, frame::Frame, call_expr::Expr; kwargs...) = evaluate_call_compiled!(Compiled(), frame, call_expr; kwargs...)
+evaluate_call!(@nospecialize(recurse), frame::Frame, call_expr::Expr; kwargs...) = evaluate_call_recurse!(recurse, frame, call_expr; kwargs...)
+evaluate_call!(frame::Frame, call_expr::Expr; kwargs...) = evaluate_call!(finish_and_return!, frame, call_expr; kwargs...)
 
 # The following come up only when evaluating toplevel code
 function evaluate_methoddef(frame, node)

src/localmethtable.jl

@@ -7,7 +7,7 @@ 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, parentframe::FrameCode, idx::Int)
+function get_call_framecode(fargs, parentframe::FrameCode, idx::Int; enter_generated::Bool=false)
     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
@@ -19,11 +19,15 @@ function get_call_framecode(fargs, parentframe::FrameCode, idx::Int)
             # Determine whether the argument types match the signature
             sig = tme.sig.parameters::SimpleVector
             if length(sig) == nargs
-                matches = true
-                for i = 1:nargs
-                    if !isa(fargs[i], sig[i])
-                        matches = false
-                        break
+                # If this is generated, match only if `enter_generated` also matches
+                mi = tme.func::FrameInstance
+                matches = !is_generated(scopeof(mi.framecode)) || enter_generated == mi.enter_generated
+                if matches
+                    for i = 1:nargs
+                        if !isa(fargs[i], sig[i])
+                            matches = false
+                            break
+                        end
                     end
                 end
                 if matches
@@ -34,7 +38,6 @@ function get_call_framecode(fargs, parentframe::FrameCode, idx::Int)
                         tmeprev.next = tme.next
                         tme.next = tme1
                     end
-                    mi = tme.func::FrameInstance
                     return mi.framecode, mi.sparam_vals
                 end
             end
@@ -47,12 +50,12 @@ function get_call_framecode(fargs, 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)
+    ret = prepare_call(f, fargs; enter_generated=enter_generated)
     ret === nothing && return f(fargs[2:end]...), nothing
     isa(ret, Compiled) && return ret, nothing
     framecode, args, env, argtypes = ret
     # Store the results of the method lookup in the local method table
-    mi = FrameInstance(framecode, env)
+    mi = FrameInstance(framecode, env, is_generated(scopeof(framecode)) & enter_generated)
     # it's sort of odd to call this a TypeMapEntry, then set most of the fields incorrectly
     # but since we're just using it as a linked list, it's probably ok
     tme = ccall(:jl_new_struct_uninit, Any, (Any,), TypeMapEntry)::TypeMapEntry

src/types.jl

@@ -90,10 +90,11 @@ Fields:
 struct FrameInstance
     framecode::FrameCode
     sparam_vals::SimpleVector
+    enter_generated::Bool
 end
 
 Base.show(io::IO, instance::FrameInstance) =
-    print(io, "FrameInstance(", scopeof(instance.framecode), ", ", instance.sparam_vals, ')')
+    print(io, "FrameInstance(", scopeof(instance.framecode), ", ", instance.sparam_vals, ", ", instance.enter_generated, ')')
 
 """
 `FrameData` holds the arguments, local variables, and intermediate execution state