inference: continue const-prop' when concrete-eval returns non-inlineable

base/compiler/abstractinterpretation.jl

@@ -767,11 +767,13 @@ struct ConstCallResults
     const_result::ConstResult
     effects::Effects
     edge::MethodInstance
-    ConstCallResults(@nospecialize(rt),
-                     const_result::ConstResult,
-                     effects::Effects,
-                     edge::MethodInstance) =
-        new(rt, const_result, effects, edge)
+    function ConstCallResults(
+        @nospecialize(rt),
+        const_result::ConstResult,
+        effects::Effects,
+        edge::MethodInstance)
+        return new(rt, const_result, effects, edge)
+    end
 end
 
 function abstract_call_method_with_const_args(interp::AbstractInterpreter,
@@ -785,24 +787,33 @@ function abstract_call_method_with_const_args(interp::AbstractInterpreter,
         return nothing
     end
     eligibility = concrete_eval_eligible(interp, f, result, arginfo, sv)
+    concrete_eval_result = nothing
     if eligibility === :concrete_eval
-        return concrete_eval_call(interp, f, result, arginfo, sv, invokecall)
+        concrete_eval_result = concrete_eval_call(interp, f, result, arginfo, sv, invokecall)
+        # if we don't inline the result of this concrete evaluation,
+        # give const-prop' a chance to inline a better method body
+        if !may_optimize(interp) || (
+            may_inline_concrete_result(concrete_eval_result.const_result::ConcreteResult) ||
+            concrete_eval_result.rt === Bottom) # unless this call deterministically throws and thus is non-inlineable
+            return concrete_eval_result
+        end
+        # TODO allow semi-concrete interp for this call?
     end
     mi = maybe_get_const_prop_profitable(interp, result, f, arginfo, si, match, sv)
-    mi === nothing && return nothing
+    mi === nothing && return concrete_eval_result
     if is_constprop_recursed(result, mi, sv)
         add_remark!(interp, sv, "[constprop] Edge cycle encountered")
         return nothing
     end
     # try semi-concrete evaluation
     if eligibility === :semi_concrete_eval
-        res = semi_concrete_eval_call(interp, mi, result, arginfo, sv)
-        if res !== nothing
-            return res
+        irinterp_result = semi_concrete_eval_call(interp, mi, result, arginfo, sv)
+        if irinterp_result !== nothing
+            return irinterp_result
         end
     end
     # try constant prop'
-    return const_prop_call(interp, mi, result, arginfo, sv)
+    return const_prop_call(interp, mi, result, arginfo, sv, concrete_eval_result)
 end
 
 function const_prop_enabled(interp::AbstractInterpreter, sv::AbsIntState, match::MethodMatch)
@@ -895,7 +906,7 @@ function concrete_eval_call(interp::AbstractInterpreter,
         Core._call_in_world_total(world, f, args...)
     catch
         # The evaluation threw. By :consistent-cy, we're guaranteed this would have happened at runtime
-        return ConstCallResults(Union{}, ConcreteResult(edge, result.effects), result.effects, edge)
+        return ConstCallResults(Bottom, ConcreteResult(edge, result.effects), result.effects, edge)
     end
     return ConstCallResults(Const(value), ConcreteResult(edge, EFFECTS_TOTAL, value), EFFECTS_TOTAL, edge)
 end
@@ -1165,7 +1176,8 @@ function semi_concrete_eval_call(interp::AbstractInterpreter,
 end
 
 function const_prop_call(interp::AbstractInterpreter,
-    mi::MethodInstance, result::MethodCallResult, arginfo::ArgInfo, sv::AbsIntState)
+    mi::MethodInstance, result::MethodCallResult, arginfo::ArgInfo, sv::AbsIntState,
+    concrete_eval_result::Union{Nothing,ConstCallResults}=nothing)
     inf_cache = get_inference_cache(interp)
     𝕃ᵢ = typeinf_lattice(interp)
     inf_result = cache_lookup(𝕃ᵢ, mi, arginfo.argtypes, inf_cache)
@@ -1188,6 +1200,11 @@ function const_prop_call(interp::AbstractInterpreter,
             return nothing
         end
         @assert inf_result.result !== nothing
+        if concrete_eval_result !== nothing
+            # override return type and effects with concrete evaluation result if available
+            inf_result.result = concrete_eval_result.rt
+            inf_result.ipo_effects = concrete_eval_result.effects
+        end
     else
         # found the cache for this constant prop'
         if inf_result.result === nothing

test/compiler/inference.jl

@@ -5099,3 +5099,16 @@ end
     refine_partial_struct2(42, s)
 end |> only === String
 # JET.test_call(s::AbstractString->Base._string(s, 'c'))
+
+# override const-prop' return type with the concrete-eval result
+# if concrete-eval returns non-inlineable constant
+Base.@assume_effects :foldable function continue_const_prop(i, j)
+    chars = map(Char, i:j)
+    String(chars)
+end
+@test Base.return_types() do
+    Val(length(continue_const_prop(1, 5)))
+end |> only === Val{5}
+@test fully_eliminated() do
+    length(continue_const_prop(1, 5))
+end

test/compiler/inline.jl

@@ -2082,3 +2082,17 @@ end
     z = erase_before_inlining(true, y)
     return length(z), length(z)
 end
+
+# continue const-prop' when concrete-eval result is too big
+const THE_BIG_TUPLE_2 = ntuple(identity, 1024)
+return_the_big_tuple2(a) = (a, THE_BIG_TUPLE_2)
+let src = code_typed1() do
+        return return_the_big_tuple2(42)[2]
+    end
+    @test count(isinvoke(:return_the_big_tuple2), src.code) == 0
+end
+let src = code_typed1() do
+        return iterate(("1", '2'), 1)
+    end
+    @test count(isinvoke(:iterate), src.code) == 0
+end