@memoized macro

memoized.jl

macro expect(pred)
    :( ($esc(pred)) ? nothing : error($"error: expected $pred == true") )
end

is_expr(ex, head::Symbol) = (isa(ex, Expr) && (ex.head == head))

function split_fdef(fdef::Expr)
    @expect (fdef.head == :function) || (fdef.head == :(=))
    @expect length(fdef.args) == 2
    signature, body = tuple(fdef.args...)
    @expect is_expr(signature, :call)
    @expect length(signature.args) >= 1
    (signature, body)
end
split_fdef(f::Any) = error("split_fdef: expected function definition, got\n$f")


macro memoized(fdef)
    code_memoized(fdef)
end
function code_memoized(fdef)
    signature, body = split_fdef(fdef)
    args = signature.args[2:end]

    memo = expr(:quote, ObjectIdDict())
    expr(:function, esc(signature), quote
        key = ($esc(expr(:tuple, args)))
        has(($memo), key) ? ($memo)[key] :
                           (($memo)[key] = let; ($esc(body)); end)
    end)
end

test.jl

module Test
import Base.*
load("memoized.jl")

ex = code_memoized(:(f(x) = randi()))
println(ex)

@memoized f(x) = return randi(10)
println()
println({f(x) for x=1:10})
println({f(x) for x=1:10})

@memoized fib(n) = n < 2 ? n : fib(n-1) + fib(n-2)
println()
@time fib(20)
@time fib(20)

end

Output from test.jl:

function escape(f(x))   #  line 27:
    key=escape((x,))    #  line 28:
    if has(:(ObjectIdDict()), key)  #  line 28:
        begin 
            :(ObjectIdDict())[key]
        end
    else    #  line 29:
        :(ObjectIdDict())[key]=escape(randi())
    end
end

{5, 1, 1, 10, 6, 1, 9, 9, 9, 6}
{5, 1, 1, 10, 6, 1, 9, 9, 9, 6}

elapsed time: 0.005666017532348633 seconds
elapsed time: 8.821487426757812e-6 seconds

I don't think this handles return in the body, since that will cause the whole function to return without modifying the dictionary.

Your're right. Thanks!

Ok, added a let block to fix it.