Closed tshort closed 3 years ago
Here is some code from @MasonProtter posted on Slack that unrolls a dynamic dispatch.
using Base: rewrap_unionall, unwrap_unionall, uncompressed_ast, CodeInfo
function expr_to_codeinfo(m, argnames, spnames, sp, e)
lam = Expr(:lambda, argnames,
Expr(Symbol("scope-block"),
Expr(:block,
Expr(:return,
Expr(:block,
e,
)))))
ex = if spnames === nothing
lam
else
Expr(Symbol("with-static-parameters"), lam, spnames...)
end
# Get the code-info for the generatorbody in order to use it for generating a dummy
# code info object.
ci = ccall(:jl_expand, Any, (Any, Any), ex, m)
end
@generated function static_methods_tuple(@nospecialize(m::Module), @nospecialize(f) , @nospecialize(_T::Type{T})) where {T <: Tuple}
world = typemax(UInt)
methods(f.instance)
ms = Tuple(methods(f.instance, T))
ci = expr_to_codeinfo(m, [Symbol("#self#"), :m, :f, :_T], [:T], (:T,), :($ms))
method_insts = Core.Compiler.method_instances(f.instance, T, world)
method_doesnot_exist = isempty(method_insts)
mt = f.name.mt
# Now we add the edges so if a method is defined this recompiles
if method_doesnot_exist
# No method so attach to method table
mt = f.name.mt
ci.edges = Core.Compiler.vect(mt, (mt, Tuple{Vararg{Any}}))
else # method exists, attach edges to all instances
ci.edges = method_insts
end
return ci
end
@generated function _unroll_dispatch(f, @nospecialize(args::Tuple), ::Val{N}) where {N}
T = Tuple{(Any for _ in 1:N)...}
ms = static_methods_tuple(Main, f.instance, T)
ex = Expr(:block)
for i in 1:N
sig = Tuple{(ms[i].sig.parameters[2:end])...}
_ex = quote
if args isa $sig
return invoke($(f.instance), $sig, args...)
end
end
push!(ex.args, _ex)
end
ex
end
macro unroll_dispatch(fcall)
@assert fcall.head == :call
f = fcall.args[1]
args = Tuple(fcall.args[2:end])
N = length(args)
quote
_unroll_dispatch($f, $(Expr(:tuple, args...)), $(Val(N)))
end |> esc
end
Some strategies to support this are:
jl_apply_generic
. It's hard, and Valentin says we have an ABI problem.