From d7d8c2a00345360dfd697b3a9bf3ea9f136313d0 Mon Sep 17 00:00:00 2001
From: Paul Berg <paul@plutojl.org>
Date: Wed, 29 Nov 2023 15:04:55 +0100
Subject: [PATCH 1/2] add support for EnterNode

the EnterNode struct was added in https://github.com/JuliaLang/julia/pull/52300
This commit adds support for ingesting and exporting CodeInfo with EnterNode
but still uses `Expr(:enter)` inside IRTools IR. We may want to switch
the representation to `IRTools.EnterNode` similarly to `IRTools.Slot` if
more fields are added to the `Core.EnterNode` struct.
---
 src/ir/print.jl  |  3 +++
 src/ir/wrap.jl   | 16 +++++++++++++---
 test/compiler.jl | 18 ++++++++++++++++++
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/ir/print.jl b/src/ir/print.jl
index 0969f80..14fe5ff 100644
--- a/src/ir/print.jl
+++ b/src/ir/print.jl
@@ -79,6 +79,9 @@ end
 
 function print_stmt(io::IO, ::Val{:enter}, ex)
   print(io, "try #$(ex.args[1])")
+  if length(ex.args) >= 2
+      print(io, " with scope ", ex.args[2])
+  end
 end
 
 function print_stmt(io::IO, ::Val{:leave}, ex)
diff --git a/src/ir/wrap.jl b/src/ir/wrap.jl
index 6df715e..e9d7f24 100644
--- a/src/ir/wrap.jl
+++ b/src/ir/wrap.jl
@@ -29,6 +29,9 @@ function IRCode(ir::IR)
     for (v, st) in b
       defs[v] = Variable(length(stmts)+1)
       ex = varmap(x -> get(defs, x, x), st.expr) |> unvars
+      @static if isdefined(Core.IR, :EnterNode)
+        ex = isexpr(ex, :enter) ? Core.EnterNode(ex.args...) : ex
+      end
       push!(stmts, ex)
       push!(types, st.type)
       push!(lines, st.line)
@@ -43,7 +46,7 @@ function IRCode(ir::IR)
         else
           push!(stmts, Expr(:call, GlobalRef(Core, :throw), "unreachable"))
         end
-      elseif br.condition == nothing
+      elseif br.condition === nothing
         push!(stmts, GotoNode(br.block))
       else
         cond = get(defs, br.condition, br.condition) |> unvars
@@ -130,8 +133,15 @@ function IR(ci::CodeInfo, nargs::Integer; meta = nothing)
     ex = ci.code[i]
     if ex isa Core.NewvarNode
       continue
-    elseif isexpr(ex, :enter)
-      _rename[Core.SSAValue(i)] = push!(ir, Expr(:enter, findfirst(==(ex.args[1]), bs)+1))
+    elseif @static if isdefined(Core.IR, :EnterNode); ex isa Core.IR.EnterNode else isexpr(ex, :enter) end
+      catch_dest = @static if isdefined(Core.IR, :EnterNode)
+          ex.catch_dest
+      else
+          ex.args[1]
+      end
+    enter_expr = Expr(:enter, findfirst(==(catch_dest), bs)+1)
+      isdefined(ex, :scope) && push!(enter_expr.args, ex.scope)
+      _rename[Core.SSAValue(i)] = push!(ir, enter_expr)
     elseif ex isa GotoNode
       branch!(ir, findfirst(==(ex.label), bs)+1)
     elseif isgotoifnot(ex)
diff --git a/test/compiler.jl b/test/compiler.jl
index 634ccec..6d887e8 100644
--- a/test/compiler.jl
+++ b/test/compiler.jl
@@ -211,3 +211,21 @@ end
     @test (code_typed(func_ir, Tuple{typeof(func_ir)}) |> only
            isa Pair{Core.CodeInfo,DataType})
 end
+
+function f_try_catch(x)
+    y = 0.
+    try
+        y = sqrt(x)
+    catch
+
+    end
+    y
+end
+
+@testset "try/catch" begin
+    ir = @code_ir f_try_catch(1.)
+    @test true
+    fir = func(ir)
+    @test fir(nothing,1.) == 1.
+    @test_broken fir(nothing,-1.) == 1.
+end

From 5cdd66233725b6035767a9f537e5932e4946b222 Mon Sep 17 00:00:00 2001
From: Paul Berg <paul@plutojl.org>
Date: Mon, 11 Dec 2023 11:44:02 +0100
Subject: [PATCH 2/2] make inlining tests more robusts to lowering changes

In particular, https://github.com/JuliaLang/julia/pull/51970/ which
adds more nodes in lowered ir.
---
 test/ir.jl | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/test/ir.jl b/test/ir.jl
index f715125..f436b13 100644
--- a/test/ir.jl
+++ b/test/ir.jl
@@ -42,7 +42,12 @@ end
 let
   ir1 = @code_ir foo(1, 2)
   ir2 = @code_ir bar(1, 2)
-  ir = IRTools.inline(ir1, IRTools.var(4), ir2)
+  inline_at = findfirst(ir1) do (_, stmt)
+    IRTools.isexpr(stmt.expr, :call) &&
+      stmt.expr.args[1] == GlobalRef(@__MODULE__, :bar)
+  end
+  @test inline_at !== nothing
+  ir = IRTools.inline(ir1, inline_at, ir2)
   f = IRTools.func(ir)
   @test f(nothing, 2, 3) == 3
   @test f(nothing, 3, 2) == 3
@@ -66,7 +71,12 @@ end
 let
   ir = @code_ir foo2(1)
   ir2 = @code_ir foo1(1)
-  ir3 = IRTools.inline(ir, IRTools.var(4), ir2)
+  inline_at = findfirst(ir) do (_, stmt)
+    IRTools.isexpr(stmt.expr, :call) &&
+      stmt.expr.args[1] == GlobalRef(@__MODULE__, :foo1)
+  end
+  @test inline_at !== nothing
+  ir3 = IRTools.inline(ir, inline_at, ir2)
   @test IRTools.func(ir3)(nothing, 2) == 12
   @test IRTools.func(ir3)(nothing, 101) == 101
 end
@@ -87,7 +97,12 @@ end
 let
   ir = @code_ir foo2(1)
   ir2 = @code_ir foo1(1)
-  ir3 = IRTools.inline(ir, IRTools.var(3), ir2)
+  inline_at = findfirst(ir) do (_, stmt)
+    IRTools.isexpr(stmt.expr, :call) &&
+      stmt.expr.args[1] == GlobalRef(@__MODULE__, :foo1)
+  end
+  @test inline_at !== nothing
+  ir3 = IRTools.inline(ir, inline_at, ir2)
   @test IRTools.func(ir3)(nothing, 2) == foo2(2)
   @test IRTools.func(ir3)(nothing, -2) == foo2(-2)
 end