effects: add reflection utility for the new effect analysis (#44785)

This commit adds new reflection utility named `Base.infer_effects` that
works in the same way as `Base.return_types` but returns inferred effects
instead. It would be helpful to test that certain method call has an
expected effects.

For example, we can now remove `Base.@pure` annotation from the definition of
`BroadcastStyle(a::A, b::B) where {A<:AbstractArrayStyle{M},B<:AbstractArrayStyle{N}} where {M,N}`
and checks it's still eligible for concrete evaluation like this
(see <#44776 (comment)> for the context):
```julia
julia> import Base.Broadcast: AbstractArrayStyle, DefaultArrayStyle, Unknown

julia> function BroadcastStyle(a::A, b::B) where {A<:AbstractArrayStyle{M},B<:AbstractArrayStyle{N}} where {M,N}
           if Base.typename(A) === Base.typename(B)
               return A(Val(max(M, N)))
           end
           return Unknown()
       end
BroadcastStyle (generic function with 1 method)

julia> # test that the above definition is eligible for concrete evaluation
       @test Base.infer_effects(BroadcastStyle, (DefaultArrayStyle{1},DefaultArrayStyle{2},)) |> Core.Compiler.is_total_or_error
Test Passed
```

Co-authored-by: Takafumi Arakaki <[email protected]>

base/compiler/abstractinterpretation.jl

@@ -48,7 +48,7 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
         # aren't any in the throw block either to enable other optimizations.
         add_remark!(interp, sv, "Skipped call in throw block")
         nonoverlayed = false
-        if isoverlayed(method_table(interp)) && sv.ipo_effects.nonoverlayed
+        if isoverlayed(method_table(interp)) && is_nonoverlayed(sv.ipo_effects)
             # as we may want to concrete-evaluate this frame in cases when there are
             # no overlayed calls, try an additional effort now to check if this call
             # isn't overlayed rather than just handling it conservatively
@@ -712,7 +712,7 @@ function concrete_eval_eligible(interp::AbstractInterpreter,
     @nospecialize(f), result::MethodCallResult, arginfo::ArgInfo, sv::InferenceState)
     # disable concrete-evaluation since this function call is tainted by some overlayed
     # method and currently there is no direct way to execute overlayed methods
-    isoverlayed(method_table(interp)) && !result.edge_effects.nonoverlayed && return false
+    isoverlayed(method_table(interp)) && !is_nonoverlayed(result.edge_effects) && return false
     return f !== nothing &&
            result.edge !== nothing &&
            is_total_or_error(result.edge_effects) &&

base/compiler/typeinfer.jl

@@ -904,15 +904,24 @@ end
 
 # compute an inferred AST and return type
 function typeinf_code(interp::AbstractInterpreter, method::Method, @nospecialize(atype), sparams::SimpleVector, run_optimizer::Bool)
+    frame = typeinf_frame(interp, method, atype, sparams, run_optimizer)
+    frame === nothing && return nothing, Any
+    frame.inferred || return nothing, Any
+    code = frame.src
+    rt = widenconst(ignorelimited(frame.result.result))
+    return code, rt
+end
+
+# compute an inferred frame
+function typeinf_frame(interp::AbstractInterpreter, method::Method, @nospecialize(atype), sparams::SimpleVector, run_optimizer::Bool)
     mi = specialize_method(method, atype, sparams)::MethodInstance
     ccall(:jl_typeinf_begin, Cvoid, ())
     result = InferenceResult(mi)
     frame = InferenceState(result, run_optimizer ? :global : :no, interp)
-    frame === nothing && return (nothing, Any)
+    frame === nothing && return nothing
     typeinf(interp, frame)
     ccall(:jl_typeinf_end, Cvoid, ())
-    frame.inferred || return (nothing, Any)
-    return (frame.src, widenconst(ignorelimited(result.result)))
+    return frame
 end
 
 # compute (and cache) an inferred AST and return type

base/compiler/types.jl

@@ -79,19 +79,25 @@ function Effects(e::Effects = EFFECTS_UNKNOWN′;
         inbounds_taints_consistency)
 end
 
+is_consistent(effects::Effects)   = effects.consistent === ALWAYS_TRUE
+is_effect_free(effects::Effects)  = effects.effect_free === ALWAYS_TRUE
+is_nothrow(effects::Effects)      = effects.nothrow === ALWAYS_TRUE
+is_terminates(effects::Effects)   = effects.terminates === ALWAYS_TRUE
+is_nonoverlayed(effects::Effects) = effects.nonoverlayed
+
 is_total_or_error(effects::Effects) =
-    effects.consistent === ALWAYS_TRUE &&
-    effects.effect_free === ALWAYS_TRUE &&
-    effects.terminates === ALWAYS_TRUE
+    is_consistent(effects) &&
+    is_effect_free(effects) &&
+    is_terminates(effects)
 
 is_total(effects::Effects) =
     is_total_or_error(effects) &&
-    effects.nothrow === ALWAYS_TRUE
+    is_nothrow(effects)
 
 is_removable_if_unused(effects::Effects) =
-    effects.effect_free === ALWAYS_TRUE &&
-    effects.terminates === ALWAYS_TRUE &&
-    effects.nothrow === ALWAYS_TRUE
+    is_effect_free(effects) &&
+    is_terminates(effects) &&
+    is_nothrow(effects)
 
 function encode_effects(e::Effects)
     return (e.consistent.state << 0) |

base/reflection.jl

@@ -1310,6 +1310,35 @@ function return_types(@nospecialize(f), @nospecialize(types=default_tt(f));
     return rt
 end
 
+function infer_effects(@nospecialize(f), @nospecialize(types=default_tt(f));
+                       world = get_world_counter(),
+                       interp = Core.Compiler.NativeInterpreter(world))
+    ccall(:jl_is_in_pure_context, Bool, ()) && error("code reflection cannot be used from generated functions")
+    types = to_tuple_type(types)
+    if isa(f, Core.Builtin)
+        args = Any[types.parameters...]
+        rt = Core.Compiler.builtin_tfunction(interp, f, args, nothing)
+        return Core.Compiler.builtin_effects(f, args, rt)
+    else
+        effects = Core.Compiler.EFFECTS_TOTAL
+        matches = _methods(f, types, -1, world)::Vector
+        if isempty(matches)
+            # although this call is known to throw MethodError (thus `nothrow=ALWAYS_FALSE`),
+            # still mark it `TRISTATE_UNKNOWN` just in order to be consistent with a result
+            # derived by the effect analysis, which can't prove guaranteed throwness at this moment
+            return Core.Compiler.Effects(effects; nothrow=Core.Compiler.TRISTATE_UNKNOWN)
+        end
+        for match in matches
+            match = match::Core.MethodMatch
+            frame = Core.Compiler.typeinf_frame(interp,
+                match.method, match.spec_types, match.sparams, #=run_optimizer=#false)
+            frame === nothing && return Core.Compiler.Effects()
+            effects = Core.Compiler.tristate_merge(effects, frame.ipo_effects)
+        end
+        return effects
+    end
+end
+
 """
     print_statement_costs(io::IO, f, types)
 

test/compiler/irpasses.jl

@@ -1036,8 +1036,9 @@ let ci = code_typed(foo_cfg_empty, Tuple{Bool}, optimize=true)[1][1]
     @test isa(ir.stmts[length(ir.stmts)][:inst], ReturnNode)
 end
 
-@test Core.Compiler.builtin_effects(getfield, Any[Complex{Int}, Symbol], Any).effect_free.state == 0x01
-@test Core.Compiler.builtin_effects(getglobal, Any[Module, Symbol], Any).effect_free.state == 0x01
+@test Core.Compiler.is_effect_free(Base.infer_effects(getfield, (Complex{Int}, Symbol)))
+@test Core.Compiler.is_effect_free(Base.infer_effects(getglobal, (Module, Symbol)))
+
 # Test that UseRefIterator gets SROA'd inside of new_to_regular (#44557)
 # expression and new_to_regular offset are arbitrary here, we just want to see the UseRefIterator erased
 let e = Expr(:call, Core.GlobalRef(Base, :arrayset), false, Core.SSAValue(4), Core.SSAValue(9), Core.SSAValue(8))

test/reflection.jl

@@ -964,3 +964,35 @@ end
     @eval m f4(a) = return
     @test Base.default_tt(m.f4) == Tuple
 end
+
+Base.@assume_effects :terminates_locally function issue41694(x::Int)
+    res = 1
+    1 < x < 20 || throw("bad")
+    while x > 1
+        res *= x
+        x -= 1
+    end
+    return res
+end
+maybe_effectful(x::Int) = 42
+maybe_effectful(x::Any) = unknown_operation()
+function f_no_methods end
+
+@testset "infer_effects" begin
+    @test Base.infer_effects(issue41694, (Int,)) |> Core.Compiler.is_terminates
+    @test Base.infer_effects((Int,)) do x
+        issue41694(x)
+    end |> Core.Compiler.is_terminates
+    @test Base.infer_effects(issue41694) |> Core.Compiler.is_terminates # use `default_tt`
+    let effects = Base.infer_effects(maybe_effectful, (Any,)) # union split
+        @test !Core.Compiler.is_consistent(effects)
+        @test !Core.Compiler.is_effect_free(effects)
+        @test !Core.Compiler.is_nothrow(effects)
+        @test !Core.Compiler.is_terminates(effects)
+        @test !Core.Compiler.is_nonoverlayed(effects)
+    end
+    @test Base.infer_effects(f_no_methods) |> !Core.Compiler.is_nothrow
+    # builtins
+    @test Base.infer_effects(typeof, (Any,)) |> Core.Compiler.is_total
+    @test Base.infer_effects(===, (Any,Any)) |> Core.Compiler.is_total
+end