How to write @memoize that works in both local and global scope

General Usage
1

Let’s write a @memoize f(x) = ... macro in Julia. Glossing over gensyms etc. it can just macroexpand into

memo = Dict()
f(x) = get(memo, x) do ... end

At the global scope memo will be a slow global variable, so let’s instead macroexpand to:

const memo = Dict()
f(x) = get(memo, x) do ... end

But oops, Julia doesn’t like const in the local scope. This won’t work to memoize a local function. Alright, let’s try

let memo = Dict()
    f(x) = get(memo, x) do ... end
end

to make our Lisp fathers proud. But the funny rule here is that at the global scope, f’s name will be mangled to make it a local function unless it is prefixed with a module, like Main.f(x) = get(memo, x) do ... end. I could use @__MODULE__ to get that but, of course, it won’t work at the local scope.

It’s quite a frustrating situation. Memoize.jl uses eval in the macro code to get around it, which is evil in its own way. And Memoization.jl uses generated functions (!)

This shouldn’t be so hard! Did we all miss something obvious?

2

Can’t you just assign assign the function returned by the let into the global scope

fib = let memo=Dict()
    innerfib(x) = get!(memo, x) do
        println(x)
        x<=2 && return BigInt(1)
        innerfib(x-1)+innerfib(x-2)
    end
end

As a macro that works on recursive, or mutually recursive functions, you might need

fib = let memo=Dict()
    innerfib(x) = get!(memo, x) do
        x<=2 && return BigInt(1)
        fib(x-1)+fib(x-2)
    end
end

so that the name of the recursive calls is unaltered.

3

Yes, but then A) it only works if fib is the only method for that function and B) it will have poor performance because fib is a global.

4

Ok, how about constructing the cache within the macro body itself and then closing over it in the generated code?

julia> macro memoize(f, body)
         memo = Dict()
         quote
           function $(esc(f))(x)
             get!($memo, x) do
               println("Cache miss")
               $(esc(body))(x)
             end
           end
         end
       end
@memoize (macro with 1 method)

julia> @memoize fib x -> begin
         x <= 2 && return BigInt(1)
         fib(x - 1) + fib(x - 2)
       end
fib (generic function with 2 methods)

julia> fib(1)
Cache miss
1

julia> fib(1)
1

julia> fib(5)
Cache miss
Cache miss
Cache miss
Cache miss
5

This macro results in a normal function definition, so you can extend it at global scope as desired:

julia> fib(x::Float64) = 1.0
fib (generic function with 2 methods)

julia> fib(5)  # uses the memoized `fib`
5

julia> fib(5.0)  # uses `fib(x::Float64)` method
1.0

and it also doesn’t leak anything into global scope when run from a local scope:

julia> let
         @memoize foo x -> begin
           x + 1
         end
         
         foo(2)
       end
Cache miss
3

julia> foo   # not defined because it was local to the `let` block
ERROR: UndefVarError: foo not defined

I’ve been a little lazy in requiring that the user split up the fib symbol from the x -> begin...end body, but you could parse that data out of a normal-looking function definition so that the macro works like:

@memoize fib(x)
   ... # body here
end
1 Like
5

Yes, that’s basically what Memoize.jl does, but it supports arbitrary AbstractDict, like @memoize LRUDict(...) foo(x) = ..., which is why it needs to eval in the containing module.

It fails for local functions because they all share the same cache, since the macro body is evaluated only once:

julia> g = []
       for x in 1:10
       push!(g, 
             @memoize f(y) = x+y)
       end

julia> g[1](1)
2

julia> g[2](1)
2

The let solution would work fine here.

6

Recently I have been avoiding the “magic” of macros where I think the explicit code is simple enough, and can benefit from customization: e.g. for memoizing:

  • local vs global scope
  • new function versus adding method to existing function (perhaps in another module)
  • typed Dict for caching
  • custom caching where cache value is function of argument

Is the fib example faster with typed Dict cache?

let
    memo = Dict{Int,BigInt}()
    global function fib2(x)
        get!(memo, x) do
            x <= 2 && return BigInt(1)
            fib2(x-1) + fib2(x-2)
        end
    end
end
1 Like
7

I suppose one could get around this problem by putting g inside of the key. Mostly, I’m annoyed by the eval. Lisp forefathers would disapprove, but maybe that’s one instance where it’s legitimate and necessary?