tshort / StaticCompiler.jl

Compiles Julia code to a standalone library (experimental)
Other
500 stars 30 forks source link

Supporting dynamic calls (jl_apply_generic / call) #18

Closed tshort closed 3 years ago

tshort commented 4 years ago

Some strategies to support this are:

tshort commented 4 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