From 38b09efed0ac51c09df6f31fcc65ed3e43bab77a Mon Sep 17 00:00:00 2001
From: willtebbutt <wtebbutt@turing.ac.uk>
Date: Wed, 31 Jul 2024 15:27:34 +0100
Subject: [PATCH] More work

---
 src/interpreter/ir_normalisation.jl | 23 +++++++++++++++++++++++
 src/test_utils.jl                   |  2 ++
 2 files changed, 25 insertions(+)

diff --git a/src/interpreter/ir_normalisation.jl b/src/interpreter/ir_normalisation.jl
index 33e9067e..2864cf6c 100644
--- a/src/interpreter/ir_normalisation.jl
+++ b/src/interpreter/ir_normalisation.jl
@@ -26,6 +26,7 @@ function normalise!(ir::IRCode, spnames::Vector{Symbol})
         inst = splatnew_to_call(inst)
         inst = intrinsic_to_function(inst)
         inst = lift_getfield_and_others(inst)
+        inst = invoke_to_invoke_wrapper(inst)
         ir.stmts.inst[n] = inst
     end
     return ir
@@ -182,3 +183,25 @@ end
 __get_arg(x::GlobalRef) = getglobal(x.mod, x.name)
 __get_arg(x::QuoteNode) = x.value
 __get_arg(x) = x
+
+"""
+    invoke_wrapper(f, TT, x::Vararg{Any, N}) where {N}
+
+Equivalent to `invoke(f, TT, x...)`, but is a primitive that won't get lowered / inlined
+away.
+"""
+invoke_wrapper(f, TT, x::Vararg{Any, N}) where {N} = invoke(f, TT, x...)
+
+@is_primitive MinimalCtx Tuple{typeof(invoke_wrapper), Vararg}
+
+"""
+    invoke_to_invoke_wrapper(inst)
+
+Replaces `:call`s to `Core.invoke` with calls to `Tapir.invoke_wrapper`. This done to
+prevent such calls being inlined / lowered away during optimisation, as access to them is
+required later on.
+"""
+function invoke_to_invoke_wrapper(inst)
+    (Meta.isexpr(inst, :call) && inst.args[1] == invoke) || return inst
+    return Expr(:call, invoke_wrapper, inst.args[2:end]...)
+end
diff --git a/src/test_utils.jl b/src/test_utils.jl
index f3a0b417..bcd29eed 100644
--- a/src/test_utils.jl
+++ b/src/test_utils.jl
@@ -1469,6 +1469,8 @@ end
 
 test_for_invoke(x::Float64, y::Float64, z::Float64...) = x + sum(y)
 
+inlinable_invoke_call(x::Float64) = invoke(test_for_invoke, Tuple{Float64}, x)
+
 function generate_test_functions()
     return Any[
         (false, :allocs, nothing, const_tester),