EA: minor updates on EA (#51365)

- update outdated test code
- add more comments on TODO items
- use `MethodInstance` for better frame-identity check
- minor cleanups on the inlining code

Author: aviatesk

base/compiler/ssair/EscapeAnalysis/EscapeAnalysis.jl

@@ -1096,6 +1096,7 @@ function escape_exception!(astate::AnalysisState, tryregions::Vector{UnitRange{I
     # NOTE if `:the_exception` is the only way to access the exception, we can do:
     # exc = SSAValue(pc)
     # excinfo = estate[exc]
+    # TODO? set up a special effect bit that checks the existence of `rethrow` and `current_exceptions` and use it here
     excinfo = ⊤
     escapes = estate.escapes
     for i in 1:length(escapes)

base/compiler/ssair/inlining.jl

@@ -211,10 +211,9 @@ function cfg_inline_item!(ir::IRCode, idx::Int, todo::InliningTodo, state::CFGIn
     return nothing
 end
 
-function cfg_inline_unionsplit!(ir::IRCode, idx::Int,
-                                (; fully_covered, #=atype,=# cases, bbs)::UnionSplit,
-                                state::CFGInliningState,
-                                params::OptimizationParams)
+function cfg_inline_unionsplit!(ir::IRCode, idx::Int, union_split::UnionSplit,
+                                state::CFGInliningState, params::OptimizationParams)
+    (; fully_covered, #=atype,=# cases, bbs) = union_split
     inline_into_block!(state, block_for_inst(ir, idx))
     from_bbs = Int[]
     delete!(state.split_targets, length(state.new_cfg_blocks))
@@ -234,9 +233,7 @@ function cfg_inline_unionsplit!(ir::IRCode, idx::Int,
             end
         end
         push!(from_bbs, length(state.new_cfg_blocks))
-        # TODO: Right now we unconditionally generate a fallback block
-        # in case of subtyping errors - This is probably unnecessary.
-        if i != length(cases) || (!fully_covered)
+        if !(i == length(cases) && fully_covered)
             # This block will have the next condition or the final else case
             push!(state.new_cfg_blocks, BasicBlock(StmtRange(idx, idx)))
             push!(state.new_cfg_blocks[cond_bb].succs, length(state.new_cfg_blocks))
@@ -582,7 +579,7 @@ function ir_inline_unionsplit!(compact::IncrementalCompact, idx::Int, argexprs::
         cond = true
         nparams = fieldcount(atype)
         @assert nparams == fieldcount(mtype)
-        if i != ncases || !fully_covered
+        if !(i == ncases && fully_covered)
             for i = 1:nparams
                 a, m = fieldtype(atype, i), fieldtype(mtype, i)
                 # If this is always true, we don't need to check for it
@@ -1475,7 +1472,7 @@ function handle_match!(cases::Vector{InliningCase},
     # We may see duplicated dispatch signatures here when a signature gets widened
     # during abstract interpretation: for the purpose of inlining, we can just skip
     # processing this dispatch candidate (unless unmatched type parameters are present)
-    !allow_typevars && _any(case->case.sig === spec_types, cases) && return true
+    !allow_typevars && any(case::InliningCase->case.sig === spec_types, cases) && return true
     item = analyze_method!(match, argtypes, info, flag, state; allow_typevars)
     item === nothing && return false
     push!(cases, InliningCase(spec_types, item))

test/compiler/EscapeAnalysis/EAUtils.jl

@@ -41,9 +41,10 @@ function code_escapes(@nospecialize(f), @nospecialize(types=Base.default_tt(f));
                       world::UInt = get_world_counter(),
                       debuginfo::Symbol = :none)
     tt = Base.signature_type(f, types)
-    interp = EscapeAnalyzer(world, tt)
-    results = Base.code_typed_by_type(tt; optimize=true, world, interp)
-    isone(length(results)) || throw(ArgumentError("`code_escapes` only supports single analysis result"))
+    match = Base._which(tt; world, raise=true)
+    mi = Core.Compiler.specialize_method(match)::MethodInstance
+    interp = EscapeAnalyzer(world, mi)
+    Core.Compiler.typeinf_ext(interp, mi)
     return EscapeResult(interp.ir, interp.estate, interp.mi, debuginfo === :source, interp)
 end
 
@@ -93,17 +94,17 @@ mutable struct EscapeAnalyzer <: AbstractInterpreter
     const inf_cache::Vector{InferenceResult}
     const code_cache::CodeCache
     const escape_cache::EscapeCache
-    const entry_tt
+    const entry_mi::MethodInstance
     ir::IRCode
     estate::EscapeState
     mi::MethodInstance
-    function EscapeAnalyzer(world::UInt, @nospecialize(tt),
+    function EscapeAnalyzer(world::UInt, entry_mi::MethodInstance,
                             code_cache::CodeCache=GLOBAL_CODE_CACHE,
                             escape_cache::EscapeCache=GLOBAL_ESCAPE_CACHE)
         inf_params = InferenceParams()
         opt_params = OptimizationParams()
         inf_cache = InferenceResult[]
-        return new(world, inf_params, opt_params, inf_cache, code_cache, escape_cache, tt)
+        return new(world, inf_params, opt_params, inf_cache, code_cache, escape_cache, entry_mi)
     end
 end
 
@@ -210,10 +211,10 @@ function run_passes_ipo_safe_with_ea(interp::EscapeAnalyzer,
         @timeit "EA" estate = analyze_escapes(ir, nargs, get_escape_cache)
     catch err
         @error "error happened within EA, inspect `Main.failed_escapeanalysis`"
-        @eval Main failed_escapeanalysis = $(FailedAnalysis(ir, nargs, get_escape_cache))
+        Main.failed_escapeanalysis = FailedAnalysis(ir, nargs, get_escape_cache)
         rethrow(err)
     end
-    if caller.linfo.specTypes === interp.entry_tt
+    if caller.linfo === interp.entry_mi
         # return back the result
         interp.ir = cccopy(ir)
         interp.estate = estate

test/compiler/EscapeAnalysis/EscapeAnalysis.jl

@@ -74,12 +74,13 @@ end
 @testset "basics" begin
     let # arg return
         result = code_escapes((Any,)) do a # return to caller
-            Base.donotdelete(1) # TODO #51143
+            println("prevent ConstABI")
             return nothing
         end
         @test has_arg_escape(result.state[Argument(2)])
         # return
         result = code_escapes((Any,)) do a
+            println("prevent ConstABI")
             return a
         end
         i = only(findall(isreturn, result.ir.stmts.stmt))
@@ -114,9 +115,9 @@ end
     end
     let # :gc_preserve_begin / :gc_preserve_end
         result = code_escapes((String,)) do s
-            Base.donotdelete(1) # TODO #51143
             m = SafeRef(s)
             GC.@preserve m begin
+                println(s)
                 return nothing
             end
         end
@@ -125,7 +126,7 @@ end
         @test has_no_escape(result.state[SSAValue(i)])
     end
     let # :isdefined
-        result = code_escapes((String, Bool, )) do a, b
+        result = code_escapes((String, Bool,)) do a, b
             if b
                 s = Ref(a)
             end
@@ -510,7 +511,7 @@ end
         end
         i = only(findall(isnew, result.ir.stmts.stmt))
         r = only(findall(isreturn, result.ir.stmts.stmt))
-        @test_broken !has_return_escape(result.state[SSAValue(i)], r)
+        @test_broken !has_return_escape(result.state[SSAValue(i)], r) # TODO? see `escape_exception!`
     end
     let # sequential: escape information imposed on `err1` and `err2 should propagate separately
         result = @eval M $code_escapes() do
@@ -537,7 +538,7 @@ end
         r = only(findall(isreturn, result.ir.stmts.stmt))
         @test has_all_escape(result.state[SSAValue(i1)])
         @test has_return_escape(result.state[SSAValue(i2)], r)
-        @test_broken !has_all_escape(result.state[SSAValue(i2)])
+        @test_broken !has_all_escape(result.state[SSAValue(i2)]) # TODO? see `escape_exception!`
     end
     let # nested: escape information imposed on `inner` shouldn't propagate to `s`
         result = @eval M $code_escapes() do
@@ -2350,13 +2351,11 @@ end
 # no method error
 @noinline identity_if_string(x::SafeRef) = (println("preventing inlining"); nothing)
 let result = code_escapes((SafeRef{String},)) do x
-        Base.donotdelete(1) # TODO #51143
         identity_if_string(x)
     end
     @test has_no_escape(ignore_argescape(result.state[Argument(2)]))
 end
 let result = code_escapes((Union{SafeRef{String},Nothing},)) do x
-        Base.donotdelete(1) # TODO #51143
         identity_if_string(x)
     end
     i = only(findall(iscall((result.ir, identity_if_string)), result.ir.stmts.stmt))