As @simeonschaub has recommended, I have rewritten my code to return CodeInfo object, thought the documentation is scarce. My current solution looks like
# Generated functions
using Dictionaries
function retrieve_code_info(sigtypes, world = Base.get_world_counter())
S = Tuple{map(s -> Core.Compiler.has_free_typevars(s) ? typeof(s.parameters[1]) : s, sigtypes)...}
_methods = Base._methods_by_ftype(S, -1, world)
isempty(_methods) && @error("method $(sigtypes) does not exist, may-be run it once")
type_signature, raw_static_params, method = _methods[1] # method is the same as we would get by invoking methods(+, (Int, Int)).ms[1]
# this provides us with the CodeInfo
method_instance = Core.Compiler.specialize_method(method, type_signature, raw_static_params, false)
code_info = Core.Compiler.retrieve_code_info(method_instance)
end
struct Calls
stamps::Vector{Float64} # contains the time stamps
event::Vector{Symbol} # name of the function that is being recorded
startstop::Vector{Symbol} # if the time stamp corresponds to start or to stop
i::Ref{Int}
end
function Calls(n::Int)
Calls(Vector{Float64}(undef, n+1), Vector{Symbol}(undef, n+1), Vector{Symbol}(undef, n+1), Ref{Int}(0))
end
function Base.show(io::IO, calls::Calls)
for i in 1:calls.i[]
println(io, calls.stamps[i] - calls.stamps[1]," ", calls.startstop[i]," ",calls.event[i])
end
end
function Base.push!(calls::Calls, s::Symbol, ev::Symbol)
n = calls.i[] = calls.i[] + 1
n > length(calls.stamps) && return
calls.event[n] = ev
calls.startstop[n] = s
calls.stamps[n] = time()
end
reset!(calls::Calls) = calls.i[] = 0
function overdubbable(ex::Expr)
ex.head != :call && return(false)
length(ex.args) < 2 && return(false)
(ex.args[1] isa Core.IntrinsicFunction) && return(false)
return(true)
end
overdubbable(ex) = false
timable(ex::Expr) = ex.head == :call
timable(ex) = false
rename_args(ex, ssamap) = ex
rename_args(ex::Expr, ssamap) = Expr(ex.head, rename_args(ex.args, ssamap)...)
rename_args(args::AbstractArray, ssamap) = map(a -> rename_args(a, ssamap), args)
rename_args(r::Core.ReturnNode, ssamap) = Core.ReturnNode(rename_args(r.val, ssamap))
rename_args(a::Core.SSAValue, ssamap) = Core.SSAValue(ssamap[a.id])
exportname(ex::GlobalRef) = ex.name
exportname(ex::Expr) = ex.args[1]
using Base: invokelatest
dummy() = return
overdub(f::Core.IntrinsicFunction, args...) = f(args...)
@generated function overdub(f::F, args...) where {F}
F = typeof(foo)
args = (Float64, Float64)
# ci = code_lowered((F, args))[1]
ci = retrieve_code_info((F, args...))
# this is to initialize a new CodeInfo and fill it with values from the
# overdubbed function
new_ci = code_lowered(dummy, Tuple{})[1]
empty!(new_ci.code)
empty!(new_ci.slotnames)
foreach(s -> push!(new_ci.slotnames, s), ci.slotnames)
new_ci.slotnames = vcat([Symbol("#self#"), :f, :args], ci.slotnames[length(args)+2:end])
empty!(new_ci.linetable)
foreach(s -> push!(new_ci.linetable, s), ci.linetable)
empty!(new_ci.codelocs)
ssamap = Dict{Int, Int}()
slotvar = Dict{Int, Any}()
for i in 1:length(args)
push!(new_ci.code, :(Base.getindex(args, $(i))))
slotvar[i+1] = Core.SSAValue(i)
push!(new_ci.codelocs, ci.codelocs[1])
end
slotvar[1] = Core.SlotNumber(1)
foreach(i -> slotvar[i[2]] = Core.SlotNumber(i[1]+3), enumerate(length(args)+2:length(ci.slotnames)))
j = length(args)
for (i, ex) in enumerate(ci.code)
ex = rename_args(ex, slotvar, ssamap)
if timable(ex)
fname = exportname(ex)
fname = :(Symbol($(fname)))
push!(new_ci.code, Expr(:call, :push!, :to, :(:start), fname))
j += 1
push!(new_ci.codelocs, ci.codelocs[i])
ex = overdubbable(ex) ? Expr(:call, :overdub, ex.args...) : ex
# ex = i ∈ used ? Expr(:(=) , ssa_vars[i], ex) : ex
push!(new_ci.code, ex)
push!(new_ci.codelocs, ci.codelocs[i])
j += 1
ssamap[i] = j
push!(new_ci.code, Expr(:call, :push!, :to, :(:stop), fname))
push!(new_ci.codelocs, ci.codelocs[i])
j += 1
else
# ex = i ∈ used ? Expr(:(=) , ssa_vars[i], ex) : ex
push!(new_ci.code, ex)
push!(new_ci.codelocs, ci.codelocs[i])
j += 1
ssamap[i] = j
end
end
new_ci
# Core.Compiler.replace_code_newstyle!(ci, ir, length(ir.argtypes)-1)
new_ci.inferred = false
new_ci.ssavaluetypes = length(new_ci.code)
new_ci
end
But when I call it as follows
function foo(x, y)
z = x * y
z + sin(y)
end
global const to = Calls(100)
reset!(to)
overdub(foo, 1.0, 1.0)
, the overdub does not executed foo but the Core.CodeInfo. This quite puzzling, as based on Simon’s post I would expect it is done automatically.
Thanks for any help.
Best wishes,
Tomas