dynamic_overlay.jl

# adapted from https://github.com/JuliaCompilerPlugins/Mixtape.jl/blob/75fc4aa25e918f92bfd0ba706236f61a466efc11/examples/dynamic_overlay.jl

module DynamicOverlay

using Mixtape
import Mixtape: CompilationContext, transform, allow
using MacroTools
using CodeInfoTools
using BenchmarkTools

foo(x) = x^5
bar(x) = x^n
n = 10
apply(f, x1, x2::Val{T}) where {T} = f(x1, T)

function f(x)
    g = x < 5 ? foo : bar
    return g(2)
end

struct MyMix <: CompilationContext end
allow(ctx::MyMix, m::Module) = m == DynamicOverlay

swap(e) = e
function swap(e::Expr)
    new = MacroTools.postwalk(e) do s
        isexpr(s, :call) || return s
        s.args[1] == Base.literal_pow || return s
        return Expr(:call, apply, Base.:(*), s.args[3:end]...)
    end
    return new
end

function transform(::MyMix, src)
    b = CodeInfoTools.Builder(src)
    for (v, st) in b
        b[v] = swap(st)
    end
    return CodeInfoTools.finish(b)
end

# JIT compile an entry and call.
fn = Mixtape.jit(f, Tuple{Int64}; ctx = MyMix())
println(fn(3))
println(fn(6))
@btime fn(6)

# Mixtape cached call.
Mixtape.@load_abi()
println(call(f, 3; ctx = MyMix()))
println(call(f, 6; ctx = MyMix()))
@btime call(f, 6; ctx = MyMix())

# Native.
println(f(3))
println(f(6))
@btime f(5)

end # module

#  jl examples/dynamic_overlay.jl
# 10
# 1024
#   55.046 ns (2 allocations: 112 bytes)
# 10
# 1024
#   246.617 ns (2 allocations: 112 bytes)
# 32
# 1024
#   26.067 ns (1 allocation: 16 bytes)