From 8cbe183f63cd05c8fbb5bf0f85b8e67f2cf9214c Mon Sep 17 00:00:00 2001
From: Keno Fischer <keno@juliacomputing.com>
Date: Thu, 11 Jan 2024 22:11:48 +0000
Subject: [PATCH] ir: Don't accidentally consider pending statements from
 previous block

When IncremetnalCompact deletes and edge between blocks, it needs
to go to the target and delete any reference to the edge from
the PhiNodes therein. What happened in #52858, is that we had
a situtation like:

```
# Block n-1
%a Expr(...)
Expr(...) # <- This node is pending at the start of compaction
# Block n
%b PhiNode
```

To do the deletion, we use `CompactPeekIterator`, which gets
a list of original statements and iterates over all the statements
inside the range, including pending ones. However, there is a
semantic confusion about whether to include all all statements
since the previous statement (%a above), or only those statements
that are actually attached to the first instruction (%b above).
This of course doesn't really matter usually unless you're at
a BasicBlock boundary.

For the present case, we really do not want the instructions in
the previous basic block - the iteration of PhiNodes terminates
if it seems a stmt that cannot be in the phi block, so mistakenly
seeing one of those instructions causes it to fail to fix a PhiNode
that it should have, causing #52858.

We fix #52858 by giving the CompactPeekIterator a flag to only
return stmts in the correct basic block. While we're at it, also
fix the condition for what kinds of statements are allowed in the
phi block, as that was clarified more recently than this code
was written.
---
 base/compiler/ssair/ir.jl | 24 +++++++++++++++++++-----
 test/compiler/irpasses.jl | 24 ++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/base/compiler/ssair/ir.jl b/base/compiler/ssair/ir.jl
index 8029f6c57fb83..759997feb04da 100644
--- a/base/compiler/ssair/ir.jl
+++ b/base/compiler/ssair/ir.jl
@@ -1340,8 +1340,8 @@ function kill_edge!(compact::IncrementalCompact, active_bb::Int, from::Int, to::
         else
             stmts = compact.ir.cfg.blocks[to].stmts
             for stmt in CompactPeekIterator(compact, first(stmts), last(stmts))
-                stmt === nothing && continue
-                isa(stmt, PhiNode) || break
+                is_valid_phiblock_stmt(stmt) || break
+                isa(stmt, PhiNode) || continue
                 i = findfirst(x::Int32->x==from, stmt.edges)
                 if i !== nothing
                     deleteat!(stmt.edges, i)
@@ -1684,13 +1684,27 @@ struct CompactPeekIterator
     compact::IncrementalCompact
     start_idx::Int
     end_idx::Int
+    include_stmts_before_start::Bool
 end
+CompactPeekIterator(compact::IncrementalCompact, start_idx::Int, end_idx::Int) =
+    CompactPeekIterator(compact, start_idx, end_idx, false)
+
 
 function CompactPeekIterator(compact::IncrementalCompact, start_idx::Int)
     return CompactPeekIterator(compact, start_idx, 0)
 end
 
-entry_at_idx(entry::NewNodeInfo, idx::Int) = entry.attach_after ? entry.pos == idx - 1 : entry.pos == idx
+function entry_at_idx(entry::NewNodeInfo, idx::Int, start_idx::Int, include_stmts_before_start::Bool)
+    if entry.attach_after
+        if !include_stmts_before_start
+            entry.pos >= start_idx || return false
+        end
+        return entry.pos == idx - 1
+    else
+        return entry.pos == idx
+    end
+end
+
 function iterate(it::CompactPeekIterator, (idx, aidx, bidx)::NTuple{3, Int}=(it.start_idx, it.compact.new_nodes_idx, 1))
     if it.end_idx > 0 && idx > it.end_idx
         return nothing
@@ -1702,7 +1716,7 @@ function iterate(it::CompactPeekIterator, (idx, aidx, bidx)::NTuple{3, Int}=(it.
     if compact.new_nodes_idx <= length(compact.perm)
         new_nodes = compact.ir.new_nodes
         for eidx in aidx:length(compact.perm)
-            if entry_at_idx(new_nodes.info[compact.perm[eidx]], idx)
+            if entry_at_idx(new_nodes.info[compact.perm[eidx]], idx, it.start_idx, it.include_stmts_before_start)
                 entry = new_nodes.stmts[compact.perm[eidx]]
                 return (entry[:stmt], (idx, eidx+1, bidx))
             end
@@ -1710,7 +1724,7 @@ function iterate(it::CompactPeekIterator, (idx, aidx, bidx)::NTuple{3, Int}=(it.
     end
     if !isempty(compact.pending_perm)
         for eidx in bidx:length(compact.pending_perm)
-            if entry_at_idx(compact.pending_nodes.info[compact.pending_perm[eidx]], idx)
+            if entry_at_idx(compact.pending_nodes.info[compact.pending_perm[eidx]], idx, it.start_idx, it.include_stmts_before_start)
                 entry = compact.pending_nodes.stmts[compact.pending_perm[eidx]]
                 return (entry[:stmt], (idx, aidx, eidx+1))
             end
diff --git a/test/compiler/irpasses.jl b/test/compiler/irpasses.jl
index 43027ff740e43..4f4faf3ae6617 100644
--- a/test/compiler/irpasses.jl
+++ b/test/compiler/irpasses.jl
@@ -1740,3 +1740,27 @@ end
     return 0
 end
 @test code_typed(f52703)[1][2] === Int
+
+# Issue #52858 - compaction gets confused by pending node
+let code = Any[
+    # Block 1
+    GotoIfNot(true, 6),
+    # Block 2
+    Expr(:call, println, 1),
+    Expr(:call, Base.inferencebarrier, true),
+    GotoIfNot(SSAValue(3), 6),
+    # Block 3
+    nothing,
+    # Block 4
+    PhiNode(Int32[1, 4, 5], Any[1, 2, 3]),
+    ReturnNode(SSAValue(6))
+]
+    ir = make_ircode(code)
+    Core.Compiler.insert_node!(ir, SSAValue(5),
+        Core.Compiler.NewInstruction(
+            Expr(:call, println, 2), Nothing, Int32(1)),
+            #= attach_after = =# true)
+    ir = Core.Compiler.compact!(ir, true)
+    @test Core.Compiler.verify_ir(ir) === nothing
+    @test count(x->isa(x, GotoIfNot), ir.stmts.stmt) == 1
+end