Fix return_type model when run in toplevel context (#44670)

Fixes the test failure observed in #44652.

base/compiler/abstractinterpretation.jl

@@ -644,6 +644,7 @@ function abstract_call_method(interp::AbstractInterpreter, method::Method, @nosp
     if edge === nothing
         edgecycle = edgelimited = true
     end
+
     # we look for the termination effect override here as well, since the :terminates effect
     # may have been tainted due to recursion at this point even if it's overridden
     if is_effect_overridden(sv, :terminates_globally)

base/compiler/inferencestate.jl

@@ -59,6 +59,10 @@ mutable struct InferenceState
     inferred::Bool
     dont_work_on_me::Bool
 
+    # Whether to restrict inference of abstract call sites to avoid excessive work
+    # Set by default for toplevel frame.
+    restrict_abstract_call_sites::Bool
+
     # Inferred purity flags
     ipo_effects::Effects
 
@@ -133,7 +137,7 @@ mutable struct InferenceState
             #=callers_in_cycle=#Vector{InferenceState}(),
             #=parent=#nothing,
             #=cached=#cache === :global,
-            #=inferred=#false, #=dont_work_on_me=#false,
+            #=inferred=#false, #=dont_work_on_me=#false, #=restrict_abstract_call_sites=# isa(linfo.def, Module),
             #=ipo_effects=#Effects(EFFECTS_TOTAL; consistent, inbounds_taints_consistency),
             interp)
         result.result = frame

base/compiler/tfuncs.jl

@@ -1993,7 +1993,13 @@ function return_type_tfunc(interp::AbstractInterpreter, argtypes::Vector{Any}, s
                     if contains_is(argtypes_vec, Union{})
                         return CallMeta(Const(Union{}), false)
                     end
+                    # 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), sv, -1)
+                    sv.restrict_abstract_call_sites = old_restrict
                     info = verbose_stmt_info(interp) ? ReturnTypeCallInfo(call.info) : false
                     rt = widenconditional(call.rt)
                     if isa(rt, Const)

base/compiler/types.jl

@@ -407,7 +407,7 @@ It also bails out from local statement/frame inference when any lattice element
 but `AbstractInterpreter` doesn't provide a specific interface for configuring it.
 """
 bail_out_toplevel_call(::AbstractInterpreter, @nospecialize(callsig), sv#=::InferenceState=#) =
-    return isa(sv.linfo.def, Module) && !isdispatchtuple(callsig)
+    return sv.restrict_abstract_call_sites && !isdispatchtuple(callsig)
 bail_out_call(::AbstractInterpreter, @nospecialize(rt), sv#=::InferenceState=#) =
     return rt === Any
 bail_out_apply(::AbstractInterpreter, @nospecialize(rt), sv#=::InferenceState=#) =

test/compiler/inference.jl

@@ -4108,3 +4108,10 @@ end |> !Core.Compiler.is_foldable
 @test !fully_eliminated() do
     entry_to_be_invalidated('a')
 end
+
+# Make sure return_type_tfunc doesn't accidentally cause bad inference if used
+# at top level.
+@test let
+    Base.Experimental.@force_compile
+    Core.Compiler.return_type(+, NTuple{2, Rational})
+end == Rational