From 5f8f33b7141957ba0b8b89574018f99ad3eaca15 Mon Sep 17 00:00:00 2001
From: Paul Berg <paul@plutojl.org>
Date: Mon, 16 Oct 2023 02:41:46 +0200
Subject: [PATCH] Allow using `ReturnNode()` in `@generated` code (#51715)

IRTools.jl currently tries to use `ReturnNode()` to model unreachable
block terminators but it fails in `find_ssavalue_uses`. This PR adds a
check to enable using `ReturnNode()` in untyped code.

One other alternative for frontends is to use an instruction known to
terminate (`Core.throw`) instead. See
https://github.com/FluxML/IRTools.jl/pull/115 for more context.

(cherry picked from commit ff03e5116d452d0e6e5970a50c187d468febf594)
---
 base/compiler/utilities.jl |  1 +
 test/staged.jl             | 29 +++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/base/compiler/utilities.jl b/base/compiler/utilities.jl
index f3c5694535ce6..7c3b19ead54cc 100644
--- a/base/compiler/utilities.jl
+++ b/base/compiler/utilities.jl
@@ -390,6 +390,7 @@ function find_ssavalue_uses(body::Vector{Any}, nvals::Int)
     for line in 1:length(body)
         e = body[line]
         if isa(e, ReturnNode)
+            isdefined(e, :val) || continue
             e = e.val
         elseif isa(e, GotoIfNot)
             e = e.cond
diff --git a/test/staged.jl b/test/staged.jl
index 5204f5f6ca777..76d02c1938e4d 100644
--- a/test/staged.jl
+++ b/test/staged.jl
@@ -346,3 +346,32 @@ let world = Base.get_world_counter()
     @test all(lin->lin.method === :sin, src.linetable)
     @test sin_generated(42) == sin(42)
 end
+
+# Allow passing unreachable insts in generated codeinfo
+let
+    dummy() = return
+    dummy_m = which(dummy, Tuple{})
+
+    src = Base.uncompressed_ir(dummy_m)
+    src.code = Any[
+        # block 1
+        Core.ReturnNode(nothing),
+        # block 2
+        Core.ReturnNode(),
+    ]
+    nstmts = length(src.code)
+    nslots = 1
+    src.ssavaluetypes = nstmts
+    src.codelocs = fill(Int32(1), nstmts)
+    src.ssaflags = fill(Int32(0), nstmts)
+    src.slotflags = fill(0, nslots)
+    src.slottypes = Any[Any]
+
+    @eval function f_unreachable()
+        $(Expr(:meta, :generated, Returns(src)))
+        $(Expr(:meta, :generated_only))
+    end
+
+    ir, _ = Base.code_ircode(f_unreachable, ()) |> only
+    @test length(ir.cfg.blocks) == 1
+end