From 8827eb3e927ec2559458b2a7d00739b12b407c04 Mon Sep 17 00:00:00 2001
From: Cody Tapscott <84105208+topolarity@users.noreply.github.com>
Date: Thu, 13 Jul 2023 11:54:09 -0400
Subject: [PATCH] ssair: Correctly handle stmt insertion at end of basic block
 (#50528)

In the presence of `attach_after` insertions, we have to be careful to
extend the basic block to include everything up to the last insertion.
We were accounting for "new" nodes (before the compaction point), but
not "pending" nodes (after the compaction point).

Fixes https://github.com/JuliaLang/julia/issues/50379.

(cherry picked from commit cdec4c21bdfffa9b198b8ecdde83aa082ebfc4e9)
---
 base/compiler/ssair/ir.jl | 22 ++++++++++++++++------
 test/compiler/ssair.jl    | 27 +++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/base/compiler/ssair/ir.jl b/base/compiler/ssair/ir.jl
index 5c6751c1e1dda..d133b88721238 100644
--- a/base/compiler/ssair/ir.jl
+++ b/base/compiler/ssair/ir.jl
@@ -1508,10 +1508,20 @@ function finish_current_bb!(compact::IncrementalCompact, active_bb::Int,
     return skipped
 end
 
-function attach_after_stmt_after(compact::IncrementalCompact, idx::Int)
-    compact.new_nodes_idx > length(compact.perm) && return false
-    entry = compact.ir.new_nodes.info[compact.perm[compact.new_nodes_idx]]
-    return entry.pos == idx && entry.attach_after
+"""
+    stmts_awaiting_insertion(compact::IncrementalCompact, idx::Int)
+
+Returns true if there are new/pending instructions enqueued for insertion into
+`compact` on any instruction in the range `1:idx`. Otherwise, returns false.
+"""
+function stmts_awaiting_insertion(compact::IncrementalCompact, idx::Int)
+
+    new_node_waiting = compact.new_nodes_idx <= length(compact.perm) &&
+        compact.ir.new_nodes.info[compact.perm[compact.new_nodes_idx]].pos <= idx
+    pending_node_waiting = !isempty(compact.pending_perm) &&
+        compact.pending_nodes.info[compact.pending_perm[1]].pos <= idx
+
+    return new_node_waiting || pending_node_waiting
 end
 
 function process_newnode!(compact::IncrementalCompact, new_idx::Int, new_node_entry::Instruction, new_node_info::NewNodeInfo, idx::Int, active_bb::Int, do_rename_ssa::Bool)
@@ -1523,7 +1533,7 @@ function process_newnode!(compact::IncrementalCompact, new_idx::Int, new_node_en
     compact.result_idx = result_idx
     # If this instruction has reverse affinity and we were at the end of a basic block,
     # finish it now.
-    if new_node_info.attach_after && idx == last(bb.stmts)+1 && !attach_after_stmt_after(compact, idx-1)
+    if new_node_info.attach_after && idx == last(bb.stmts)+1 && !stmts_awaiting_insertion(compact, idx-1)
         active_bb += 1
         finish_current_bb!(compact, active_bb, old_result_idx)
     end
@@ -1653,7 +1663,7 @@ function iterate_compact(compact::IncrementalCompact)
     compact.result[old_result_idx] = compact.ir.stmts[idx]
     result_idx = process_node!(compact, old_result_idx, compact.ir.stmts[idx], idx, idx, active_bb, true)
     compact.result_idx = result_idx
-    if idx == last(bb.stmts) && !attach_after_stmt_after(compact, idx)
+    if idx == last(bb.stmts) && !stmts_awaiting_insertion(compact, idx)
         finish_current_bb!(compact, active_bb, old_result_idx)
         active_bb += 1
     end
diff --git a/test/compiler/ssair.jl b/test/compiler/ssair.jl
index 43f17d4ad69f2..98cab349b46ee 100644
--- a/test/compiler/ssair.jl
+++ b/test/compiler/ssair.jl
@@ -566,6 +566,33 @@ let ir = Base.code_ircode((Int,Int); optimize_until="inlining") do a, b
     @test call2.args[2] === SSAValue(2)
 end
 
+# Issue #50379 - insert_node!(::IncrementalCompact, ...) at end of basic block
+let ci = make_ci([
+        # block 1
+        #= %1: =# Core.Compiler.GotoIfNot(Expr(:boundscheck), 3),
+        # block 2
+        #= %2: =# Expr(:call, println, Argument(1)),
+        # block 3
+        #= %3: =# Core.PhiNode(),
+        #= %4: =# Core.Compiler.ReturnNode(),
+    ])
+    ir = Core.Compiler.inflate_ir(ci)
+
+    # Insert another call at end of "block 2"
+    compact = Core.Compiler.IncrementalCompact(ir)
+    new_inst = NewInstruction(Expr(:call, println, Argument(1)), Nothing)
+    insert_node!(compact, SSAValue(2), new_inst, #= attach_after =# true)
+
+    # Complete iteration
+    x = Core.Compiler.iterate(compact)
+    while x !== nothing
+        x = Core.Compiler.iterate(compact, x[2])
+    end
+    ir = Core.Compiler.complete(compact)
+
+    @test Core.Compiler.verify_ir(ir) === nothing
+end
+
 # insert_node! with new instruction with flag computed
 let ir = Base.code_ircode((Int,Int); optimize_until="inlining") do a, b
         a^b