inference: make Core.Compiler.return_type respect max_methods set…

…ting

Previously we didn't limit the number of matching methods for
`Core.Compiler.return_type`, which could lead to a potential latency
problem.
This commit makes `Core.Compiler.return_type` and the corresponding
tfunc `return_type_tfunc` respect `max_methods` setting as like the
ordinary inference.
One caveat here is that the current `Core.Compiler.return_type` interface
unfortunately doesn't allow it to see the caller module context, so it
can't respect module-wide `max_methods` setting (I guess the change to
the interface should be made as of version 2.0), that I think can lead
to a potential confusion.

base/compiler/tfuncs.jl

@@ -2640,16 +2640,21 @@ function return_type_tfunc(interp::AbstractInterpreter, argtypes::Vector{Any}, s
         return CallMeta(Const(Union{}), EFFECTS_TOTAL, NoCallInfo())
     end
 
-    # Run the abstract_call without restricting abstract call
-    # sites. Otherwise, our behavior model of abstract_call
-    # below will be wrong.
+    # NOTE We need to sync the implementation here with that of `Core.Compiler.return_type`.
+    # In particular, this tfunc should not use the module-wide `max_methods` setting
+    # until we refactor the `Core.Compiler.return_type` API so that it can observe
+    # the caller module context.
+    f = singleton_type(argtypes_vec[1])
+    max_methods = get_max_methods(interp, f)
     if isa(sv, InferenceState)
+        # Run the abstract_call without restricting abstract call sites.
+        # Otherwise, our behavior model of abstract_call below will be wrong.
         old_restrict = sv.restrict_abstract_call_sites
         sv.restrict_abstract_call_sites = false
-        call = abstract_call(interp, ArgInfo(nothing, argtypes_vec), si, sv, -1)
+        call = abstract_call(interp, ArgInfo(nothing, argtypes_vec), si, sv, max_methods)
         sv.restrict_abstract_call_sites = old_restrict
     else
-        call = abstract_call(interp, ArgInfo(nothing, argtypes_vec), si, sv, -1)
+        call = abstract_call(interp, ArgInfo(nothing, argtypes_vec), si, sv, max_methods)
     end
     info = verbose_stmt_info(interp) ? MethodResultPure(ReturnTypeCallInfo(call.info)) : MethodResultPure()
     rt = widenslotwrapper(call.rt)

base/compiler/typeinfer.jl

@@ -1100,6 +1100,9 @@ function typeinf_ext_toplevel(interp::AbstractInterpreter, linfo::MethodInstance
     return src
 end
 
+# TODO change the interface of `return_type` to allow it to recognize caller module context
+# so that we can see its `get_max_methods` setting as the ordinary inference
+
 function return_type(@nospecialize(f), t::DataType) # this method has a special tfunc
     world = ccall(:jl_get_tls_world_age, UInt, ())
     args = Any[_return_type, NativeInterpreter(world), Tuple{Core.Typeof(f), t.parameters...}]
@@ -1129,7 +1132,10 @@ function _return_type(interp::AbstractInterpreter, t::DataType)
         rt = builtin_tfunction(interp, f, args, nothing)
         rt = widenconst(rt)
     else
-        for match in _methods_by_ftype(t, -1, get_world_counter(interp))::Vector
+        lim = get_max_methods(interp, f)
+        ms = _methods_by_ftype(t, lim, get_world_counter(interp))
+        ms isa Vector || return Any
+        for match in ms
             match = match::MethodMatch
             ty = typeinf_type(interp, match.method, match.spec_types, match.sparams)
             ty === nothing && return Any

test/compiler/inference.jl

@@ -3216,15 +3216,27 @@ j30385(T, y) = k30385(f30385(T, y))
 @test Base.return_types(Tuple, (NamedTuple{<:Any,Tuple{Any,Int}},)) == Any[Tuple{Any,Int}]
 @test Base.return_types(Base.splat(tuple), (typeof((a=1,)),)) == Any[Tuple{Int}]
 
-# test that return_type_tfunc isn't affected by max_methods differently than return_type
-_rttf_test(::Int8) = 0
-_rttf_test(::Int16) = 0
-_rttf_test(::Int32) = 0
-_rttf_test(::Int64) = 0
-_rttf_test(::Int128) = 0
-_call_rttf_test() = Core.Compiler.return_type(_rttf_test, Tuple{Any})
-@test Core.Compiler.return_type(_rttf_test, Tuple{Any}) === Int
-@test _call_rttf_test() === Int
+# test that return_type_tfunc uses the same `max_methods` configuration as
+# the dynamic implementation
+Base.Experimental.@max_methods 1 function _rttf_test1 end
+_rttf_test1(::Int8) = 0
+_rttf_test1(::Int16) = 0
+_rttf_test1(::Int32) = 0
+@test Core.Compiler.return_type(_rttf_test1, Tuple{Any}) === Any
+@test Base.return_types() do
+    Core.Compiler.return_type(_rttf_test1, Tuple{Any})
+end |> only === Type{Any}
+
+Base.Experimental.@max_methods 5 function _rttf_test5 end
+_rttf_test5(::Int8) = 0
+_rttf_test5(::Int16) = 0
+_rttf_test5(::Int32) = 0
+_rttf_test5(::Int64) = 0
+_rttf_test5(::Int128) = 0
+@test Core.Compiler.return_type(_rttf_test5, Tuple{Any}) === Int
+@test Base.return_types() do
+    Core.Compiler.return_type(_rttf_test5, Tuple{Any})
+end |> only === Type{Int}
 
 f_with_Type_arg(::Type{T}) where {T} = T
 @test Base.return_types(f_with_Type_arg, (Any,)) == Any[Type]