also cache calls that end up getting compiled in local method table

Author: KristofferC

src/construct.jl

@@ -248,7 +248,7 @@ function prepare_call(@nospecialize(f), allargs; enter_generated = false)
         # The generator threw an error. Let's generate the same error by calling it.
         f(allargs[2:end]...)
     end
-    isa(ret, Compiled) && return ret
+    isa(ret, Compiled) && return ret, argtypes
     # Typical return
     framecode, lenv = ret
     if is_generated(method) && enter_generated
@@ -543,7 +543,7 @@ See [`enter_call`](@ref) for a similar approach not based on expressions.
 function enter_call_expr(expr; enter_generated = false)
     clear_caches()
     r = determine_method_for_expr(expr; enter_generated = enter_generated)
-    if isa(r, Tuple)
+    if r !== nothing && !isa(r[1], Compiled)
         return prepare_frame(r[1:end-1]...)
     end
     nothing
@@ -597,7 +597,7 @@ function enter_call(@nospecialize(finfo), @nospecialize(args...); kwargs...)
         error(f, " is a builtin or intrinsic")
     end
     r = prepare_call(f, allargs; enter_generated=enter_generated)
-    if isa(r, Tuple)
+    if r !== nothing && !isa(r[1], Compiled)
         return prepare_frame(r[1:end-1]...)
     end
     return nothing

src/localmethtable.jl

@@ -20,8 +20,12 @@ function get_call_framecode(fargs::Vector{Any}, parentframe::FrameCode, idx::Int
             sig = d_meth.sig.parameters::SimpleVector
             if length(sig) == nargs
                 # If this is generated, match only if `enter_generated` also matches
-                fi = d_meth.frameinstance::FrameInstance
-                matches = !is_generated(scopeof(fi.framecode)) || enter_generated == fi.enter_generated
+                fi = d_meth.frameinstance
+                if fi isa FrameInstance
+                    matches = !is_generated(scopeof(fi.framecode)) || enter_generated == fi.enter_generated
+                else
+                    matches = !enter_generated
+                end
                 if matches
                     for i = 1:nargs
                         if !isa(fargs[i], sig[i])
@@ -38,7 +42,11 @@ function get_call_framecode(fargs::Vector{Any}, parentframe::FrameCode, idx::Int
                         d_methprev.next = d_meth.next
                         d_meth.next = d_meth1
                     end
-                    return fi.framecode, fi.sparam_vals
+                    if fi isa Compiled
+                        return Compiled(), nothing
+                    else
+                        return fi.framecode, fi.sparam_vals
+                    end
                 end
             end
             depth += 1
@@ -52,11 +60,16 @@ function get_call_framecode(fargs::Vector{Any}, parentframe::FrameCode, idx::Int
     fargs[1] = f = to_function(fargs[1])
     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
-    fi = FrameInstance(framecode, env, is_generated(scopeof(framecode)) && enter_generated)
-    d_meth = DispatchableMethod(nothing, fi, argtypes)
+    is_compiled = isa(ret[1], Compiled)
+    local framecode
+    if is_compiled
+        d_meth = DispatchableMethod(nothing, Compiled(), ret[2])
+    else
+        framecode, args, env, argtypes = ret
+        # Store the results of the method lookup in the local method table
+        fi = FrameInstance(framecode, env, is_generated(scopeof(framecode)) && enter_generated)
+        d_meth = DispatchableMethod(nothing, fi, argtypes)
+    end
     if isassigned(parentframe.methodtables, idx)
         d_meth.next = parentframe.methodtables[idx]
         # Drop the oldest d_meth, if necessary
@@ -74,5 +87,9 @@ function get_call_framecode(fargs::Vector{Any}, parentframe::FrameCode, idx::Int
         d_meth.next = nothing
     end
     parentframe.methodtables[idx] = d_meth
-    return framecode, env
+    if is_compiled
+        return Compiled(), nothing
+    else
+        return framecode, env
+    end
 end

src/types.jl

@@ -53,9 +53,10 @@ function breakpointchar(bps::BreakpointState)
     return bps.condition === falsecondition ? ' ' : 'd'     # no breakpoint : disabled
 end
 
+abstract type AbstractFrameInstance end
 mutable struct DispatchableMethod
     next::Union{Nothing,DispatchableMethod}  # linked-list representation
-    frameinstance::Any # really a FrameInstance but we have a cyclic dependency
+    frameinstance::Union{Compiled, AbstractFrameInstance} # really a Union{Compiled, FrameInstance} but we have a cyclic dependency
     sig::Type # for speed of matching, this is a *concrete* signature. `sig <: frameinstance.framecode.scope.sig`
 end
 
@@ -124,7 +125,7 @@ Fields:
 - `framecode`: the [`FrameCode`](@ref) for the method.
 - `sparam_vals`: the static parameter values for the method.
 """
-struct FrameInstance
+struct FrameInstance <: AbstractFrameInstance
     framecode::FrameCode
     sparam_vals::SimpleVector
     enter_generated::Bool