inlining: use already-resolved source for finalizer inlining

The source for `Core.finalizer` call is already resolved when arrived at
`sroa_mutables!`, so we can just try to inline `InliningCase` appended
to the intrinsic call there.

This commit also improves the robustness of the implementation, by
handling few edge cases e.g. when `finalizer` returns a constant value
and when the `finalizer` inliner bail out from inlining `finalizer` with
conditional block, and tests such cases explicitly.

Author: aviatesk

base/compiler/optimize.jl

@@ -578,7 +578,7 @@ function run_passes(
     @pass "Inlining"  ir = ssa_inlining_pass!(ir, ir.linetable, sv.inlining, ci.propagate_inbounds)
     # @timeit "verify 2" verify_ir(ir)
     @pass "compact 2" ir = compact!(ir)
-    @pass "SROA"      ir = sroa_pass!(ir, sv.inlining)
+    @pass "SROA"      ir = sroa_pass!(ir)
     @pass "ADCE"      ir = adce_pass!(ir)
     @pass "type lift" ir = type_lift_pass!(ir)
     @pass "compact 3" ir = compact!(ir)

base/compiler/ssair/inlining.jl

@@ -1397,6 +1397,10 @@ function handle_finalizer_call!(
         infos = MethodMatchInfo[info]
     elseif isa(info, UnionSplitInfo)
         infos = info.matches
+    # elseif isa(info, ConstCallInfo)
+    #     # NOTE currently this code path isn't active as constant propagation won't happen
+    #     # for `Core.finalizer` call because inference currently isn't able to fold a mutable
+    #     # object as a constant
     else
         return nothing
     end
@@ -1413,14 +1417,7 @@ function handle_finalizer_call!(
     cases === nothing && return nothing
     cases, all_covered = cases
     if all_covered && length(cases) == 1
-        item1 = cases[1].item
-        if isa(item1, InliningTodo)
-            push!(stmt.args, true)
-            push!(stmt.args, item1.mi)
-        elseif isa(item1, InvokeCase)
-            push!(stmt.args, false)
-            push!(stmt.args, item1.invoke)
-        end
+        push!(stmt.args, cases[1].item)
     end
     return nothing
 end

base/compiler/ssair/passes.jl

@@ -742,7 +742,7 @@ its argument).
 In a case when all usages are fully eliminated, `struct` allocation may also be erased as
 a result of succeeding dead code elimination.
 """
-function sroa_pass!(ir::IRCode, inlining::Union{Nothing, InliningState} = nothing)
+function sroa_pass!(ir::IRCode)
     compact = IncrementalCompact(ir)
     defuses = nothing # will be initialized once we encounter mutability in order to reduce dynamic allocations
     lifting_cache = IdDict{Pair{AnySSAValue, Any}, AnySSAValue}()
@@ -775,11 +775,11 @@ function sroa_pass!(ir::IRCode, inlining::Union{Nothing, InliningState} = nothin
                 widenconst(field_ordering) === Bool && (field_ordering = :unspecified)
             end
         elseif is_known_call(stmt, Core.finalizer, compact)
-            3 <= length(stmt.args) <= 5 || continue
+            3 <= length(stmt.args) <= 4 || continue
             # Inlining performs legality checks on the finalizer to determine
             # whether or not we may inline it. If so, it appends extra arguments
             # at the end of the intrinsic. Detect that here.
-            length(stmt.args) == 5 || continue
+            length(stmt.args) == 4 || continue
             is_finalizer = true
         elseif isexpr(stmt, :foreigncall)
             nccallargs = length(stmt.args[3]::SimpleVector)
@@ -940,44 +940,40 @@ function sroa_pass!(ir::IRCode, inlining::Union{Nothing, InliningState} = nothin
         used_ssas = copy(compact.used_ssas)
         simple_dce!(compact, (x::SSAValue) -> used_ssas[x.id] -= 1)
         ir = complete(compact)
-        sroa_mutables!(ir, defuses, used_ssas, lazydomtree, inlining)
+        sroa_mutables!(ir, defuses, used_ssas, lazydomtree)
         return ir
     else
         simple_dce!(compact)
         return complete(compact)
     end
 end
 
-function try_inline_finalizer!(ir::IRCode, argexprs::Vector{Any}, idx::Int, mi::MethodInstance, inlining::InliningState)
-    code = get(inlining.mi_cache, mi, nothing)
-    if code isa CodeInstance
-        if use_const_api(code)
-            # No code in the function - Nothing to do
-            inlining.et !== nothing && push!(inlining.et, mi)
-            return true
-        end
-        src = code.inferred
-    else
-        src = code
+function inline_finalizer!(ir::IRCode, argexprs::Vector{Any}, idx::Int, @nospecialize(case))
+    if isa(case, ConstantCase)
+        # No code in the function - Nothing to do
+        return true
+    elseif isa(case, InvokeCase)
+        insert_node!(ir, idx, NewInstruction(Expr(:invoke, case.invoke, argexprs...), Nothing), true)
+        return true
+    elseif case === nothing
+        @assert false "this case should never happen"
     end
 
-    src = inlining_policy(inlining.interp, src, IR_FLAG_NULL, mi, Any[])
-    src === nothing && return false
-    src = retrieve_ir_for_inlining(mi, src)
+    # now try to inline the finalizer body
+    (; mi, spec) = case::InliningTodo
+    src = (spec::ResolvedInliningSpec).ir
 
     # For now: Require finalizer to only have one basic block
     length(src.cfg.blocks) == 1 || return false
 
     # Ok, we're committed to inlining the finalizer
-    inlining.et !== nothing && push!(inlining.et, mi)
 
     linetable_offset, extra_coverage_line = ir_inline_linetable!(ir.linetable, src, mi.def, ir[SSAValue(idx)][:line])
     if extra_coverage_line != 0
         insert_node!(ir, idx, NewInstruction(Expr(:code_coverage_effect), Nothing, extra_coverage_line))
     end
 
-    # TODO: Use the actual inliner here rather than open coding this special
-    # purpose inliner.
+    # TODO: Use the actual inliner here rather than open coding this special purpose inliner.
     spvals = mi.sparam_vals
     ssa_rename = Vector{Any}(undef, length(src.stmts))
     for idx′ = 1:length(src.stmts)
@@ -1002,7 +998,7 @@ end
 
 is_nothrow(ir::IRCode, pc::Int) = ir.stmts[pc][:flag] & (IR_FLAG_EFFECT_FREE | IR_FLAG_NOTHROW) ≠ 0
 
-function try_inline_finalizer!(ir::IRCode, idx::Int, finalizer_idx::Int, defuse::SSADefUse, inlining::InliningState)
+function try_inline_finalizer!(ir::IRCode, idx::Int, finalizer_idx::Int, defuse::SSADefUse)
     # For now: Require that all uses and defs are in the same basic block,
     # so that live range calculations are easy.
     bb = ir.cfg.blocks[block_for_inst(ir.cfg, first(defuse.uses).idx)]
@@ -1025,40 +1021,30 @@ function try_inline_finalizer!(ir::IRCode, idx::Int, finalizer_idx::Int, defuse:
         return true
     end
 
-    check_in_range(idx) || return
-    _all(check_in_range, defuse.uses) || return
-    _all(check_in_range, defuse.defs) || return
+    check_in_range(idx) || return nothing
+    _all(check_in_range, defuse.uses) || return nothing
+    _all(check_in_range, defuse.defs) || return nothing
 
     # For now: Require all statements in the basic block range to be nothrow.
     _all(minval:maxval) do idx::Int
         return is_nothrow(ir, idx) || idx == finalizer_idx
-    end || return
+    end || return nothing
 
     # Ok, finalizer rewrite is legal.
     finalizer_stmt = ir[SSAValue(finalizer_idx)][:inst]
     argexprs = Any[finalizer_stmt.args[2], finalizer_stmt.args[3]]
-    mi = finalizer_stmt.args[5]::MethodInstance
-    if finalizer_stmt.args[4]::Bool # may inline
-        if try_inline_finalizer!(ir, argexprs, maxval, mi, inlining)
-            @goto done_finalizer
-        end
-        mi = compileable_specialization(inlining.et, mi, Effects()).invoke
-    end
-    if mi !== nothing
-        insert_node!(ir, maxval,
-            NewInstruction(Expr(:invoke, mi, argexprs...), Nothing),
-            true)
+    case = finalizer_stmt.args[4]
+    if inline_finalizer!(ir, argexprs, maxval, case)
+        # Erase call to finalizer
+        ir[SSAValue(finalizer_idx)][:inst] = nothing
     else
-        insert_node!(ir, maxval,
-            NewInstruction(Expr(:call, argexprs...), Nothing),
-            true)
+        pop!(finalizer_stmt.args)
+        @assert length(finalizer_stmt.args) == 3
     end
-    @label done_finalizer
-    # Erase call to finalizer
-    ir[SSAValue(finalizer_idx)][:inst] = nothing
+    return nothing
 end
 
-function sroa_mutables!(ir::IRCode, defuses::IdDict{Int, Tuple{SPCSet, SSADefUse}}, used_ssas::Vector{Int}, lazydomtree::LazyDomtree, inlining::Union{Nothing, InliningState})
+function sroa_mutables!(ir::IRCode, defuses::IdDict{Int, Tuple{SPCSet, SSADefUse}}, used_ssas::Vector{Int}, lazydomtree::LazyDomtree)
     for (idx, (intermediaries, defuse)) in defuses
         intermediaries = collect(intermediaries)
         # Check if there are any uses we did not account for. If so, the variable
@@ -1090,8 +1076,8 @@ function sroa_mutables!(ir::IRCode, defuses::IdDict{Int, Tuple{SPCSet, SSADefUse
                 finalizer_idx = use.idx
             end
         end
-        if finalizer_idx !== nothing && inlining !== nothing
-            try_inline_finalizer!(ir, idx, finalizer_idx, defuse, inlining)
+        if finalizer_idx !== nothing
+            try_inline_finalizer!(ir, idx, finalizer_idx, defuse)
             continue
         end
         # Partition defuses by field

test/compiler/inline.jl

@@ -1300,7 +1300,6 @@ mutable struct DoAllocNoEscape
         end
     end
 end
-
 let src = code_typed1() do
         for i = 1:1000
             DoAllocNoEscape()
@@ -1309,6 +1308,22 @@ let src = code_typed1() do
     @test count(isnew, src.code) == 0
 end
 
+# Test that we can inline a finalizer that just returns a constant value
+mutable struct DoAllocConst
+    function DoAllocConst()
+        finalizer(new()) do this
+            return nothing
+        end
+    end
+end
+let src = code_typed1() do
+        for i = 1:1000
+            DoAllocConst()
+        end
+    end
+    @test count(isnew, src.code) == 0
+end
+
 # Test that finalizer elision doesn't cause a throw to be inlined into a function
 # that shouldn't have it
 const finalizer_should_throw = Ref{Bool}(true)
@@ -1345,7 +1360,6 @@ mutable struct DoAllocNoEscapeSparam{T}
     end
 end
 DoAllocNoEscapeSparam(x::T) where {T} = DoAllocNoEscapeSparam{T}(x)
-
 let src = code_typed1(Tuple{Any}) do x
         for i = 1:1000
             DoAllocNoEscapeSparam(x)
@@ -1365,7 +1379,6 @@ mutable struct DoAllocNoEscapeNoInline
         finalizer(noinline_finalizer, new())
     end
 end
-
 let src = code_typed1() do
         for i = 1:1000
             DoAllocNoEscapeNoInline()
@@ -1375,6 +1388,30 @@ let src = code_typed1() do
     @test count(isinvoke(:noinline_finalizer), src.code) == 1
 end
 
+# Test that we don't form invalid `finalizer` call for cases we don't handle currently
+mutable struct DoAllocNoEscapeBranch
+    val::Int
+    function DoAllocNoEscapeBranch(val::Int)
+        finalizer(new(val)) do this
+            if this.val > 500
+                nothrow_side_effect(this.val)
+            else
+                nothrow_side_effect(nothing)
+            end
+        end
+    end
+end
+let src = code_typed1() do
+        for i = 1:1000
+            DoAllocNoEscapeBranch(i)
+        end
+    end
+    @test any(src.code) do @nospecialize x
+        iscall((src, Core.finalizer))(x) &&
+        length(x.args) == 3
+    end
+end
+
 # optimize `[push!|pushfirst!](::Vector{Any}, x...)`
 @testset "optimize `$f(::Vector{Any}, x...)`" for f = Any[push!, pushfirst!]
     @eval begin