compiler: apply more accurate effects to return_type_tfunc (#55338)

In extreme cases, the compiler could mark this function for
concrete-eval, even though that is illegal unless the compiler has first
deleted this instruction. Otherwise the attempt to concrete-eval will
re-run the function repeatedly until it hits a StackOverflow.

Workaround to fix #55147

@aviatesk You might know how to solve this even better, using
post-optimization effect refinements? Since we should actually only
apply the refinement of terminates=false => terminates=true (and thus
allowing concrete eval) if the optimization occurs, and not just in
inference thinks the optimization would be legal.

---------

Co-authored-by: Shuhei Kadowaki <[email protected]>

Author: vtjnash

base/boot.jl

@@ -284,7 +284,8 @@ macro _foldable_meta()
         #=:inaccessiblememonly=#true,
         #=:noub=#true,
         #=:noub_if_noinbounds=#false,
-        #=:consistent_overlay=#false))
+        #=:consistent_overlay=#false,
+        #=:nortcall=#true))
 end
 
 macro inline()   Expr(:meta, :inline)   end

base/cmd.jl

@@ -482,7 +482,7 @@ function cmd_gen(parsed)
     end
 end
 
-@assume_effects :effect_free :terminates_globally :noub function cmd_gen(
+@assume_effects :foldable !:consistent function cmd_gen(
     parsed::Tuple{Vararg{Tuple{Vararg{Union{String, SubString{String}}}}}}
 )
     return @invoke cmd_gen(parsed::Any)

base/compiler/abstractinterpretation.jl

@@ -980,7 +980,7 @@ function concrete_eval_eligible(interp::AbstractInterpreter,
         end
     end
     mi = result.edge
-    if mi !== nothing && is_foldable(effects)
+    if mi !== nothing && is_foldable(effects, #=check_rtcall=#true)
         if f !== nothing && is_all_const_arg(arginfo, #=start=#2)
             if (is_nonoverlayed(interp) || is_nonoverlayed(effects) ||
                 # Even if overlay methods are involved, when `:consistent_overlay` is
@@ -2910,8 +2910,9 @@ function override_effects(effects::Effects, override::EffectsOverride)
         notaskstate = override.notaskstate ? true : effects.notaskstate,
         inaccessiblememonly = override.inaccessiblememonly ? ALWAYS_TRUE : effects.inaccessiblememonly,
         noub = override.noub ? ALWAYS_TRUE :
-               (override.noub_if_noinbounds && effects.noub !== ALWAYS_TRUE) ? NOUB_IF_NOINBOUNDS :
-               effects.noub)
+            (override.noub_if_noinbounds && effects.noub !== ALWAYS_TRUE) ? NOUB_IF_NOINBOUNDS :
+            effects.noub,
+        nortcall = override.nortcall ? true : effects.nortcall)
 end
 
 isdefined_globalref(g::GlobalRef) = !iszero(ccall(:jl_globalref_boundp, Cint, (Any,), g))

base/compiler/compiler.jl

@@ -49,10 +49,11 @@ struct EffectsOverride
     noub::Bool
     noub_if_noinbounds::Bool
     consistent_overlay::Bool
+    nortcall::Bool
 end
 function EffectsOverride(
     override::EffectsOverride =
-        EffectsOverride(false, false, false, false, false, false, false, false, false, false);
+        EffectsOverride(false, false, false, false, false, false, false, false, false, false, false);
     consistent::Bool = override.consistent,
     effect_free::Bool = override.effect_free,
     nothrow::Bool = override.nothrow,
@@ -62,7 +63,8 @@ function EffectsOverride(
     inaccessiblememonly::Bool = override.inaccessiblememonly,
     noub::Bool = override.noub,
     noub_if_noinbounds::Bool = override.noub_if_noinbounds,
-    consistent_overlay::Bool = override.consistent_overlay)
+    consistent_overlay::Bool = override.consistent_overlay,
+    nortcall::Bool = override.nortcall)
     return EffectsOverride(
         consistent,
         effect_free,
@@ -73,9 +75,10 @@ function EffectsOverride(
         inaccessiblememonly,
         noub,
         noub_if_noinbounds,
-        consistent_overlay)
+        consistent_overlay,
+        nortcall)
 end
-const NUM_EFFECTS_OVERRIDES = 10 # sync with julia.h
+const NUM_EFFECTS_OVERRIDES = 11 # sync with julia.h
 
 # essential files and libraries
 include("essentials.jl")

base/compiler/effects.jl

@@ -58,6 +58,9 @@ following meanings:
     methods are `:consistent` with their non-overlayed original counterparts
     (see [`Base.@assume_effects`](@ref) for the exact definition of `:consistenct`-cy).
   * `ALWAYS_FALSE`: this method may invoke overlayed methods.
+- `nortcall::Bool`: this method does not call `Core.Compiler.return_type`,
+  and it is guaranteed that any other methods this method might call also do not call
+  `Core.Compiler.return_type`.
 
 Note that the representations above are just internal implementation details and thus likely
 to change in the future. See [`Base.@assume_effects`](@ref) for more detailed explanation
@@ -103,6 +106,9 @@ The output represents the state of different effect properties in the following
     - `+o` (green): `ALWAYS_TRUE`
     - `-o` (red): `ALWAYS_FALSE`
     - `?o` (yellow): `CONSISTENT_OVERLAY`
+9. `:nortcall` (`r`):
+    - `+r` (green): `true`
+    - `-r` (red): `false`
 """
 struct Effects
     consistent::UInt8
@@ -113,6 +119,7 @@ struct Effects
     inaccessiblememonly::UInt8
     noub::UInt8
     nonoverlayed::UInt8
+    nortcall::Bool
     function Effects(
         consistent::UInt8,
         effect_free::UInt8,
@@ -121,7 +128,8 @@ struct Effects
         notaskstate::Bool,
         inaccessiblememonly::UInt8,
         noub::UInt8,
-        nonoverlayed::UInt8)
+        nonoverlayed::UInt8,
+        nortcall::Bool)
         return new(
             consistent,
             effect_free,
@@ -130,7 +138,8 @@ struct Effects
             notaskstate,
             inaccessiblememonly,
             noub,
-            nonoverlayed)
+            nonoverlayed,
+            nortcall)
     end
 end
 
@@ -160,10 +169,10 @@ const NOUB_IF_NOINBOUNDS = 0x01 << 1
 # :nonoverlayed bits
 const CONSISTENT_OVERLAY = 0x01 << 1
 
-const EFFECTS_TOTAL    = Effects(ALWAYS_TRUE,  ALWAYS_TRUE,  true,  true,  true,  ALWAYS_TRUE,  ALWAYS_TRUE,  ALWAYS_TRUE)
-const EFFECTS_THROWS   = Effects(ALWAYS_TRUE,  ALWAYS_TRUE,  false, true,  true,  ALWAYS_TRUE,  ALWAYS_TRUE,  ALWAYS_TRUE)
-const EFFECTS_UNKNOWN  = Effects(ALWAYS_FALSE, ALWAYS_FALSE, false, false, false, ALWAYS_FALSE, ALWAYS_FALSE, ALWAYS_TRUE) # unknown mostly, but it's not overlayed at least (e.g. it's not a call)
-const _EFFECTS_UNKNOWN = Effects(ALWAYS_FALSE, ALWAYS_FALSE, false, false, false, ALWAYS_FALSE, ALWAYS_FALSE, ALWAYS_FALSE) # unknown really
+const EFFECTS_TOTAL    = Effects(ALWAYS_TRUE,  ALWAYS_TRUE,  true,  true,  true,  ALWAYS_TRUE,  ALWAYS_TRUE,  ALWAYS_TRUE,  true)
+const EFFECTS_THROWS   = Effects(ALWAYS_TRUE,  ALWAYS_TRUE,  false, true,  true,  ALWAYS_TRUE,  ALWAYS_TRUE,  ALWAYS_TRUE,  true)
+const EFFECTS_UNKNOWN  = Effects(ALWAYS_FALSE, ALWAYS_FALSE, false, false, false, ALWAYS_FALSE, ALWAYS_FALSE, ALWAYS_TRUE,  false) # unknown mostly, but it's not overlayed at least (e.g. it's not a call)
+const _EFFECTS_UNKNOWN = Effects(ALWAYS_FALSE, ALWAYS_FALSE, false, false, false, ALWAYS_FALSE, ALWAYS_FALSE, ALWAYS_FALSE, false) # unknown really
 
 function Effects(effects::Effects = _EFFECTS_UNKNOWN;
     consistent::UInt8 = effects.consistent,
@@ -173,7 +182,8 @@ function Effects(effects::Effects = _EFFECTS_UNKNOWN;
     notaskstate::Bool = effects.notaskstate,
     inaccessiblememonly::UInt8 = effects.inaccessiblememonly,
     noub::UInt8 = effects.noub,
-    nonoverlayed::UInt8 = effects.nonoverlayed)
+    nonoverlayed::UInt8 = effects.nonoverlayed,
+    nortcall::Bool = effects.nortcall)
     return Effects(
         consistent,
         effect_free,
@@ -182,7 +192,8 @@ function Effects(effects::Effects = _EFFECTS_UNKNOWN;
         notaskstate,
         inaccessiblememonly,
         noub,
-        nonoverlayed)
+        nonoverlayed,
+        nortcall)
 end
 
 function is_better_effects(new::Effects, old::Effects)
@@ -247,6 +258,11 @@ function is_better_effects(new::Effects, old::Effects)
     elseif new.nonoverlayed != old.nonoverlayed
         return false
     end
+    if new.nortcall
+        any_improved |= !old.nortcall
+    elseif new.nortcall != old.nortcall
+        return false
+    end
     return any_improved
 end
 
@@ -259,7 +275,8 @@ function merge_effects(old::Effects, new::Effects)
         merge_effectbits(old.notaskstate, new.notaskstate),
         merge_effectbits(old.inaccessiblememonly, new.inaccessiblememonly),
         merge_effectbits(old.noub, new.noub),
-        merge_effectbits(old.nonoverlayed, new.nonoverlayed))
+        merge_effectbits(old.nonoverlayed, new.nonoverlayed),
+        merge_effectbits(old.nortcall, new.nortcall))
 end
 
 function merge_effectbits(old::UInt8, new::UInt8)
@@ -279,16 +296,18 @@ is_inaccessiblememonly(effects::Effects) = effects.inaccessiblememonly === ALWAY
 is_noub(effects::Effects)                = effects.noub === ALWAYS_TRUE
 is_noub_if_noinbounds(effects::Effects)  = effects.noub === NOUB_IF_NOINBOUNDS
 is_nonoverlayed(effects::Effects)        = effects.nonoverlayed === ALWAYS_TRUE
+is_nortcall(effects::Effects)            = effects.nortcall
 
 # implies `is_notaskstate` & `is_inaccessiblememonly`, but not explicitly checked here
-is_foldable(effects::Effects) =
+is_foldable(effects::Effects, check_rtcall::Bool=false) =
     is_consistent(effects) &&
     (is_noub(effects) || is_noub_if_noinbounds(effects)) &&
     is_effect_free(effects) &&
-    is_terminates(effects)
+    is_terminates(effects) &&
+    (!check_rtcall || is_nortcall(effects))
 
-is_foldable_nothrow(effects::Effects) =
-    is_foldable(effects) &&
+is_foldable_nothrow(effects::Effects, check_rtcall::Bool=false) =
+    is_foldable(effects, check_rtcall) &&
     is_nothrow(effects)
 
 # TODO add `is_noub` here?
@@ -318,7 +337,8 @@ function encode_effects(e::Effects)
            ((e.notaskstate         % UInt32) << 7)  |
            ((e.inaccessiblememonly % UInt32) << 8)  |
            ((e.noub                % UInt32) << 10) |
-           ((e.nonoverlayed        % UInt32) << 12)
+           ((e.nonoverlayed        % UInt32) << 12) |
+           ((e.nortcall            % UInt32) << 14)
 end
 
 function decode_effects(e::UInt32)
@@ -330,7 +350,8 @@ function decode_effects(e::UInt32)
         _Bool((e >> 7) & 0x01),
         UInt8((e >> 8) & 0x03),
         UInt8((e >> 10) & 0x03),
-        UInt8((e >> 12) & 0x03))
+        UInt8((e >> 12) & 0x03),
+        _Bool((e >> 14) & 0x01))
 end
 
 function encode_effects_override(eo::EffectsOverride)
@@ -345,6 +366,7 @@ function encode_effects_override(eo::EffectsOverride)
     eo.noub                && (e |= (0x0001 << 7))
     eo.noub_if_noinbounds  && (e |= (0x0001 << 8))
     eo.consistent_overlay  && (e |= (0x0001 << 9))
+    eo.nortcall            && (e |= (0x0001 << 10))
     return e
 end
 
@@ -359,7 +381,8 @@ function decode_effects_override(e::UInt16)
         !iszero(e & (0x0001 << 6)),
         !iszero(e & (0x0001 << 7)),
         !iszero(e & (0x0001 << 8)),
-        !iszero(e & (0x0001 << 9)))
+        !iszero(e & (0x0001 << 9)),
+        !iszero(e & (0x0001 << 10)))
 end
 
 decode_statement_effects_override(ssaflag::UInt32) =

base/compiler/optimize.jl

@@ -42,13 +42,16 @@ const IR_FLAG_NOUB        = one(UInt32) << 8
 const IR_FLAG_EFIIMO      = one(UInt32) << 9
 # This statement is :inaccessiblememonly == INACCESSIBLEMEM_OR_ARGMEMONLY
 const IR_FLAG_INACCESSIBLEMEM_OR_ARGMEM = one(UInt32) << 10
+# This statement is :nortcall
+const IR_FLAG_NORTCALL    = one(UInt32) << 11
 # This statement has no users and may be deleted if flags get refined to IR_FLAGS_REMOVABLE
-const IR_FLAG_UNUSED      = one(UInt32) << 11
+const IR_FLAG_UNUSED      = one(UInt32) << 12
 
-const NUM_IR_FLAGS = 12 # sync with julia.h
+const NUM_IR_FLAGS = 13 # sync with julia.h
 
 const IR_FLAGS_EFFECTS =
-    IR_FLAG_CONSISTENT | IR_FLAG_EFFECT_FREE | IR_FLAG_NOTHROW | IR_FLAG_TERMINATES | IR_FLAG_NOUB
+    IR_FLAG_CONSISTENT | IR_FLAG_EFFECT_FREE | IR_FLAG_NOTHROW |
+    IR_FLAG_TERMINATES | IR_FLAG_NOUB | IR_FLAG_NORTCALL
 
 const IR_FLAGS_REMOVABLE = IR_FLAG_EFFECT_FREE | IR_FLAG_NOTHROW | IR_FLAG_TERMINATES
 
@@ -78,6 +81,9 @@ function flags_for_effects(effects::Effects)
     if is_noub(effects)
         flags |= IR_FLAG_NOUB
     end
+    if is_nortcall(effects)
+        flags |= IR_FLAG_NORTCALL
+    end
     return flags
 end
 
@@ -583,26 +589,28 @@ mutable struct PostOptAnalysisState
     all_nothrow::Bool
     all_noub::Bool
     any_conditional_ub::Bool
+    nortcall::Bool
     function PostOptAnalysisState(result::InferenceResult, ir::IRCode)
         inconsistent = BitSetBoundedMinPrioritySet(length(ir.stmts))
         tpdum = TwoPhaseDefUseMap(length(ir.stmts))
         lazypostdomtree = LazyPostDomtree(ir)
         lazyagdomtree = LazyAugmentedDomtree(ir)
         return new(result, ir, inconsistent, tpdum, lazypostdomtree, lazyagdomtree, Int[],
-                   true, true, nothing, true, true, false)
+                   true, true, nothing, true, true, false, true)
     end
 end
 
 give_up_refinements!(sv::PostOptAnalysisState) =
     sv.all_retpaths_consistent = sv.all_effect_free = sv.effect_free_if_argmem_only =
-    sv.all_nothrow = sv.all_noub = false
+    sv.all_nothrow = sv.all_noub = sv.nortcall = false
 
 function any_refinable(sv::PostOptAnalysisState)
     effects = sv.result.ipo_effects
     return ((!is_consistent(effects) & sv.all_retpaths_consistent) |
             (!is_effect_free(effects) & sv.all_effect_free) |
             (!is_nothrow(effects) & sv.all_nothrow) |
-            (!is_noub(effects) & sv.all_noub))
+            (!is_noub(effects) & sv.all_noub) |
+            (!is_nortcall(effects) & sv.nortcall))
 end
 
 struct GetNativeEscapeCache{CodeCache}
@@ -647,7 +655,8 @@ function refine_effects!(interp::AbstractInterpreter, sv::PostOptAnalysisState)
         effect_free = sv.all_effect_free ? ALWAYS_TRUE :
                       sv.effect_free_if_argmem_only === true ? EFFECT_FREE_IF_INACCESSIBLEMEMONLY : effects.effect_free,
         nothrow = sv.all_nothrow ? true : effects.nothrow,
-        noub = sv.all_noub ? (sv.any_conditional_ub ? NOUB_IF_NOINBOUNDS : ALWAYS_TRUE) : effects.noub)
+        noub = sv.all_noub ? (sv.any_conditional_ub ? NOUB_IF_NOINBOUNDS : ALWAYS_TRUE) : effects.noub,
+        nortcall = sv.nortcall ? true : effects.nortcall)
     return true
 end
 
@@ -772,6 +781,13 @@ function scan_non_dataflow_flags!(inst::Instruction, sv::PostOptAnalysisState)
             sv.all_noub = false
         end
     end
+    if !has_flag(flag, IR_FLAG_NORTCALL)
+        # if a function call that might invoke `Core.Compiler.return_type` has been deleted,
+        # there's no need to taint with `:nortcall`, allowing concrete evaluation
+        if isexpr(stmt, :call) || isexpr(stmt, :invoke)
+            sv.nortcall = false
+        end
+    end
 end
 
 function scan_inconsistency!(inst::Instruction, sv::PostOptAnalysisState)

base/compiler/ssair/show.jl

@@ -1050,6 +1050,8 @@ function Base.show(io::IO, e::Effects)
     printstyled(io, effectbits_letter(e, :noub, 'u'); color=effectbits_color(e, :noub))
     print(io, ',')
     printstyled(io, effectbits_letter(e, :nonoverlayed, 'o'); color=effectbits_color(e, :nonoverlayed))
+    print(io, ',')
+    printstyled(io, effectbits_letter(e, :nortcall, 'r'); color=effectbits_color(e, :nortcall))
     print(io, ')')
 end
 

base/compiler/tfuncs.jl

@@ -2871,7 +2871,7 @@ end
 # since abstract_call_gf_by_type is a very inaccurate model of _method and of typeinf_type,
 # while this assumes that it is an absolutely precise and accurate and exact model of both
 function return_type_tfunc(interp::AbstractInterpreter, argtypes::Vector{Any}, si::StmtInfo, sv::AbsIntState)
-    UNKNOWN = CallMeta(Type, Any, EFFECTS_THROWS, NoCallInfo())
+    UNKNOWN = CallMeta(Type, Any, Effects(EFFECTS_THROWS; nortcall=false), NoCallInfo())
     if !(2 <= length(argtypes) <= 3)
         return UNKNOWN
     end
@@ -2899,8 +2899,12 @@ function return_type_tfunc(interp::AbstractInterpreter, argtypes::Vector{Any}, s
         return UNKNOWN
     end
 
+    # effects are not an issue if we know this statement will get removed, but if it does not get removed,
+    # then this could be recursively re-entering inference (via concrete-eval), which will not terminate
+    RT_CALL_EFFECTS = Effects(EFFECTS_TOTAL; nortcall=false)
+
     if contains_is(argtypes_vec, Union{})
-        return CallMeta(Const(Union{}), Union{}, EFFECTS_TOTAL, NoCallInfo())
+        return CallMeta(Const(Union{}), Union{}, RT_CALL_EFFECTS, NoCallInfo())
     end
 
     # Run the abstract_call without restricting abstract call
@@ -2918,25 +2922,25 @@ function return_type_tfunc(interp::AbstractInterpreter, argtypes::Vector{Any}, s
     rt = widenslotwrapper(call.rt)
     if isa(rt, Const)
         # output was computed to be constant
-        return CallMeta(Const(typeof(rt.val)), Union{}, EFFECTS_TOTAL, info)
+        return CallMeta(Const(typeof(rt.val)), Union{}, RT_CALL_EFFECTS, info)
     end
     rt = widenconst(rt)
     if rt === Bottom || (isconcretetype(rt) && !iskindtype(rt))
         # output cannot be improved so it is known for certain
-        return CallMeta(Const(rt), Union{}, EFFECTS_TOTAL, info)
+        return CallMeta(Const(rt), Union{}, RT_CALL_EFFECTS, info)
     elseif isa(sv, InferenceState) && !isempty(sv.pclimitations)
         # conservatively express uncertainty of this result
         # in two ways: both as being a subtype of this, and
         # because of LimitedAccuracy causes
-        return CallMeta(Type{<:rt}, Union{}, EFFECTS_TOTAL, info)
+        return CallMeta(Type{<:rt}, Union{}, RT_CALL_EFFECTS, info)
     elseif isa(tt, Const) || isconstType(tt)
         # input arguments were known for certain
         # XXX: this doesn't imply we know anything about rt
-        return CallMeta(Const(rt), Union{}, EFFECTS_TOTAL, info)
+        return CallMeta(Const(rt), Union{}, RT_CALL_EFFECTS, info)
     elseif isType(rt)
-        return CallMeta(Type{rt}, Union{}, EFFECTS_TOTAL, info)
+        return CallMeta(Type{rt}, Union{}, RT_CALL_EFFECTS, info)
     else
-        return CallMeta(Type{<:rt}, Union{}, EFFECTS_TOTAL, info)
+        return CallMeta(Type{<:rt}, Union{}, RT_CALL_EFFECTS, info)
     end
 end