Fix problem with union spliting during inlining, add regression test (#42347)

The crash in #42264 showed up when attempting to record the new value for
an SSA rename during unionsplit inlining. Such a crash would only happen when
`compact.idx == 1`, which is an unusual condition, because it implies that no
statement of the original function had yet been processed (which is odd, because
the call being inlined must have been processed by this point). As it turns out, there
is currently exactly one case where this happens. If the inliner sees an `_apply_iterate`
(e.g. from splatting) of something that is not a Tuple or SimpleVector (thus
requiring calls to `iterate` to expand the value during `apply`), and if inference
was able to determine the total length of the resulting iteration, a special case early
inliner, will expand out the splat into explicit iterate calls. E.g. a call like:

     %r = tuple((x::Pair)...)

will be expanded out to

     %a = iterate(x::Pair)
     %b = getfield(%a, 1)
     %c = getfield(%a, 2)
     %d = iterate(x, %c)
     %e = getfield(%d, 1)
     %f = getfield(%d, 2)
     iterate(x, %f)
     %r = tuple(%b, %e)

where the inserted `iterate` calls may themselves be inlined. These newly inserted calls are
"pending nodes" during the actual inlining. Thus, if the original apply call was the first statement
of the function, these nodes would be processed before processing the statements in the function themselves.

In particular, this investigation also shows that `compact.idx`, which is the current location
in the function being inlined into is the wrong value to use for SSA renaming. Rather, we
need to use the SSA value of the just-inserted statement. In the absence of pending nodes,
these are equivalent concepts, but not for pending nodes. Fortunately the IncrementalCompact
iterator provides the old SSA value for just this reason and in fact, non-UnionSplit inlining was
already correct here. Thus, to fix the issue, simply adjust union splitting to work the same way
as a non-UnionSplit inline.

In coming up with the test case, an additional complication is that we currently do not perform this optimization for any calls where the apply call itself was unionsplit. It
is somewhat unusual (which explains why we haven't seen this much) for the apply
to not be union-split, but for the subsequent `iterate` to require union-splitting. In the
problem case, what happened is that for two out of the three union cases, the
`iterate` calls themselves would error, causing the resulting info object to look like
a non-unionsplit apply call, triggering this issue.

Co-authored-by: Keno Fischer <[email protected]>
(cherry picked from commit b5f3a99)

Author: 2 people

base/compiler/ssair/inlining.jl

@@ -492,8 +492,7 @@ function ir_inline_unionsplit!(compact::IncrementalCompact, idx::Int,
     end
 
     # We're now in the join block.
-    compact.ssa_rename[compact.idx-1] = insert_node_here!(compact, pn, typ, line)
-    nothing
+    return insert_node_here!(compact, NewInstruction(pn, typ, line))
 end
 
 function batch_inline!(todo::Vector{Pair{Int, Any}}, ir::IRCode, linetable::Vector{LineInfoNode}, propagate_inbounds::Bool)
@@ -548,7 +547,7 @@ function batch_inline!(todo::Vector{Pair{Int, Any}}, ir::IRCode, linetable::Vect
                 if isa(item, InliningTodo)
                     compact.ssa_rename[old_idx] = ir_inline_item!(compact, idx, argexprs, linetable, item, boundscheck, state.todo_bbs)
                 elseif isa(item, UnionSplit)
-                    ir_inline_unionsplit!(compact, idx, argexprs, linetable, item, boundscheck, state.todo_bbs)
+                    compact.ssa_rename[old_idx] = ir_inline_unionsplit!(compact, idx, argexprs, linetable, item, boundscheck, state.todo_bbs)
                 end
                 compact[idx] = nothing
                 refinish && finish_current_bb!(compact, 0)

test/compiler/inline.jl

@@ -331,3 +331,11 @@ struct NonIsBitsDimsUndef
 end
 @test Core.Compiler.is_inlineable_constant(NonIsBitsDimsUndef())
 @test !Core.Compiler.is_inlineable_constant((("a"^1000, "b"^1000), nothing))
+
+# Issue #42264 - crash on certain union splits
+let f(x) = (x...,)
+    # Test splatting with a Union of non-{Tuple, SimpleVector} types that require creating new `iterate` calls
+    # in inlining. For this particular case, we're relying on `iterate(::CaretesianIndex)` throwing an error, such
+    # the the original apply call is not union-split, but the inserted `iterate` call is.
+    @test code_typed(f, Tuple{Union{Int64, CartesianIndex{1}, CartesianIndex{3}}})[1][2] == Tuple{Int64}
+end