Add the specialize_vararg macro by Mason Protter

Author: YingboMa

Project.toml

@@ -8,6 +8,7 @@ Cassette = "7057c7e9-c182-5462-911a-8362d720325c"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 

src/dual_context.jl

@@ -3,6 +3,62 @@ using ChainRules
 using ChainRulesCore
 import ChainRulesCore: Zero
 
+# TODO: remove the copy pasted code and add that package
+# copyed from SpecializeVarargs.jl, written by @MasonProtter
+using MacroTools: MacroTools, splitdef, combinedef, @capture
+
+macro specialize_vararg(n::Int, fdef::Expr)
+    @assert n > 0
+
+    macros = Symbol[]
+    while fdef.head == :macrocall && length(fdef.args) == 3
+        push!(macros, fdef.args[1])
+        fdef = fdef.args[3]
+    end
+
+    d = splitdef(fdef)
+    args = d[:args][end]
+    @assert d[:args][end] isa Expr && d[:args][end].head == Symbol("...") && d[:args][end].args[] isa Symbol
+    args_symbol = d[:args][end].args[]
+
+    fdefs = Expr(:block)
+
+    for i in 1:n-1
+        di = deepcopy(d)
+        pop!(di[:args])
+        args = Tuple(gensym("arg$j") for j in 1:i)
+        Ts   = Tuple(gensym("T$j")   for j in 1:i)
+
+        args_with_Ts = ((arg, T) -> :($arg :: $T)).(args, Ts)
+
+        di[:whereparams] = (di[:whereparams]..., Ts...)
+
+        push!(di[:args], args_with_Ts...)
+        pushfirst!(di[:body].args, :($args_symbol = $(Expr(:tuple, args...))))
+        cfdef = combinedef(di)
+        mcfdef = isempty(macros) ? cfdef : foldr((m,f) -> Expr(:macrocall, m, nothing, f), macros, init=cfdef)
+        push!(fdefs.args, mcfdef)
+    end
+
+    di = deepcopy(d)
+    pop!(di[:args])
+    args = tuple((gensym() for j in 1:n)..., :($(gensym("args"))...))
+    Ts   = Tuple(gensym("T$j")   for j in 1:n)
+
+    args_with_Ts = (((arg, T) -> :($arg :: $T)).(args[1:end-1], Ts)..., args[end])
+
+    di[:whereparams] = (di[:whereparams]..., Ts...)
+
+    push!(di[:args], args_with_Ts...)
+    pushfirst!(di[:body].args, :($args_symbol = $(Expr(:tuple, args...))))
+
+    cfdef = combinedef(di)
+    mcfdef = isempty(macros) ? cfdef : foldr((m,f) -> Expr(:macrocall, m, nothing, f), macros, init=cfdef)
+    push!(fdefs.args, mcfdef)
+
+    esc(fdefs)
+end
+
 using Cassette: overdub, Context, nametype, similarcontext
 
 Cassette.@context DualContext