From 17791bbde6189d982df5f2fed13c7d0697fe6293 Mon Sep 17 00:00:00 2001
From: Julian P Samaroo <jpsamaroo@jpsamaroo.me>
Date: Sun, 6 Feb 2022 16:11:33 -0600
Subject: [PATCH] optimizer: Add early finalize calls

---
 base/compiler/optimize.jl | 53 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 52 insertions(+), 1 deletion(-)

diff --git a/base/compiler/optimize.jl b/base/compiler/optimize.jl
index a6816fb077298d..63aebbced9a9e3 100644
--- a/base/compiler/optimize.jl
+++ b/base/compiler/optimize.jl
@@ -97,7 +97,7 @@ function stmt_effect_free end # imported by EscapeAnalysis
 function alloc_array_ndims end # imported by EscapeAnalysis
 include("compiler/ssair/driver.jl")
 using .EscapeAnalysis
-import .EscapeAnalysis: EscapeState, ArgEscapeInfo, is_ipo_profitable
+import .EscapeAnalysis: EscapeState, EscapeInfo, ArgEscapeInfo, is_ipo_profitable
 
 """
     cache_escapes!(caller::InferenceResult, estate::EscapeState)
@@ -121,6 +121,56 @@ function getargescapes(mi_cache::MICache) where MICache
     end
 end
 
+"""
+    early_finalize!(ir::IRCode, estate::EscapeState) -> IRCode
+
+Analyzes `ir` for heap allocations which escape only via `FinalizerEscape`
+(thus having a `finalizer` call associated), and inserts a call to
+`finalize(obj)` just before any return where the allocation doesn't escape.
+"""
+function early_finalize!(ir::IRCode, estate::EscapeState)
+    # Find all allocations that escape only through a finalizer
+    to_finalize = Pair{Int,EscapeInfo}[]
+    for idx in (estate.nargs+1):length(estate.escapes)
+        info = estate.escapes[idx]
+        if has_finalizer_escape(info)# && !has_return_escape(info)
+            push!(to_finalize, (idx-estate.nargs)=>info)
+        end
+    end
+
+    # Find all locations to insert calls to `finalize`
+    # currently this is just all returns
+    return_locs = Int[]
+    for bb in ir.cfg.blocks
+        inst = ir.stmts.inst[bb.stmts.stop]
+        if inst isa ReturnNode
+            push!(return_locs, bb.stmts.stop)
+        end
+    end
+
+    # Insert `finalize` calls in order
+    for idx in return_locs
+        non_escaping = Pair{Int,EscapeInfo}[]
+        for (alloc_idx, alloc_info) in to_finalize
+            if !has_return_escape(alloc_info, idx) && !has_arg_escape(alloc_info)
+                push!(non_escaping, alloc_idx=>alloc_info)
+            end
+        end
+        if length(non_escaping) > 0
+            ct_expr = Expr(:foreigncall, :(:jl_get_current_task), Ref{Task}, svec(), 0, :(:ccall))
+            ct_insn = NewInstruction(ct_expr, Ref{Task})
+            ct = insert_node!(ir, idx, ct_insn)
+            for (alloc_idx, alloc_info) in non_escaping
+                fin_expr = Expr(:foreigncall, :(:jl_finalize_th), Nothing, svec(Any, Any), 0, :(:ccall), ct, SSAValue(alloc_idx))
+                fin_insn = NewInstruction(fin_expr, Nothing)
+                fin = insert_node!(ir, idx, fin_insn)
+            end
+        end
+    end
+
+    compact!(ir)
+end
+
 mutable struct OptimizationState
     linfo::MethodInstance
     src::CodeInfo
@@ -548,6 +598,7 @@ function run_passes(ci::CodeInfo, sv::OptimizationState, caller::InferenceResult
             state = analyze_escapes(ir, nargs, false, getter)
             cache_escapes!(caller, state)
         end
+        @timeit "Early Finalize" ir = early_finalize!(ir, state)
     end
     @timeit "Inlining"  ir = ssa_inlining_pass!(ir, ir.linetable, sv.inlining, ci.propagate_inbounds)
     # @timeit "verify 2" verify_ir(ir)