Performant Recursive Anonymous Functions

General Usage
, , ,
17

Yeah, this is a funny, and very aggravating property of how we handle anonymous functions. One other solution I’ve seen, first shown to me by @tkf is this:

julia> using MacroTools, ExprTools

julia> fibby = let
           function _fib(f, n)
               n ≤ 1 ? n : f(f, n-1) + f(f, n-2)
           end
           fib(n) = _fib(_fib, n)
       end

julia> @btime $fibby($10)
  221.356 ns (0 allocations: 0 bytes)
55

The idea here is to basically just play some sleight of hand to avoid julia thinking we’re doing recursion, and to stop it from doing something dumb.

This can be automated into a macro with a little elbow grease:

using ExprTools: splitdef, combinedef
using MacroTools: postwalk, @capture

"""
This is only worth using for locally scoped recursive functions.
"""
macro fast_recustion(fdef)
    d = splitdef(fdef)
    name = get!(d, :name, gensym(:f))
    fargs = copy(get!(d, :args, []))
    fkwargs = get!(d, :kwargs, [])
    _name = gensym(name)
    __name = gensym(Symbol(:_, name))
    d[:body] = postwalk(d[:body]) do ex
        if @capture(ex, f_(args__))
            if f == name
                return :($__name($__name, $(args...),))
            end
        elseif @capture(ex, f(args__; kwargs__))
            if f == name
                return :($__name($__name, $(args...); $(kwargs...),))
            end
        end
        ex
    end
    d[:name] = _name
    d[:args] = pushfirst!(d[:args], __name)
    quote
        $(combinedef(d))
        $name($(fargs...); $(fkwargs...),) = $_name($_name, $(fargs...); $(fkwargs...),)
    end |> esc
end

so that we can write

julia> fibby = let
           @fast_recustion function fib(n)
               n ≤ 1 ? n : fib(n-1) + fib(n-2)
           end
       end
(::var"#fib#56"{var"##fib#456#55"}) (generic function with 1 method)

julia> @btime $fibby($10)
  218.555 ns (0 allocations: 0 bytes)
55
1 Like