How to write a memoizer that is type inference friendly

Question

function memoize(f, T)
    cache = Dict{Any, T}()
    function g(args...)::T
        key = make_key(args...)
        get!(cache, key) do
            f(args...)
        end
    end
    g
end

fib = memoize(x::Int -> begin
    if x == 2
        return 2 
    end
    if x == 1
        return 1
    end
    fib(x - 1) + fib(x - 2)
end, Int)

This is what I get, sadly it doesn't recognise the return type though I annotated.

Also, is there a way to annotate the return type of an anonymous function?

@code_warntype fib(3)

Variables:
  #self#::#g#40{##44#45,DataType,Dict{Any,Int64}}
  args::Tuple{Int64}
  key::Tuple{Int64}
  #39::##39#41{Tuple{Int64},##44#45}

Body:
  begin 
      SSAValue(0) = (Core.getfield)(#self#::#g#40{##44#45,DataType,Dict{Any,Int64}},:T)::DataType
      key::Tuple{Int64} = (Core.tuple)((Core.getfield)(args::Tuple{Int64},1)::Int64)::Tuple{Int64} # line 20:
      #39::##39#41{Tuple{Int64},##44#45} = $(Expr(:new, ##39#41{Tuple{Int64},##44#45}, :(args), :((Core.getfield)(#self#,:f)::##44#45)))
      SSAValue(1) = #39::##39#41{Tuple{Int64},##44#45}
      SSAValue(2) = (Core.getfield)(#self#::#g#40{##44#45,DataType,Dict{Any,Int64}},:cache)::Dict{Any,Int64}
      return (Core.typeassert)((Base.convert)(SSAValue(0),$(Expr(:invoke, LambdaInfo for get!(::##39#41{Tuple{Int64},##44#45}, ::Dict{Any,Int64}, ::Tuple{Int64}), :(Main.get!), SSAValue(1), SSAValue(2), :(key))))::ANY,SSAValue(0))::ANY
  end::ANY

Update

I made a package that provides basic support for type inference friendly generic function memoization via macro. It also allows to customize cache key from function arguments.

https://github.com/colinfang/Memoize.jl

Answer 1

In order to get Julia to specialize its implementation for a specific DataType, you must use the ::Type{T} parametric type:

function memoize{T}(f, ::Type{T})
    …

That simple change means that Julia will specialize methods for each and every type memoize encounters, instead of just making one specialization for all DataTypes.