diff --git a/REQUIRE b/REQUIRE index 94237c0..436456a 100644 --- a/REQUIRE +++ b/REQUIRE @@ -1 +1,2 @@ -julia 0.5 +MacroTools +julia 0.6 diff --git a/src/Memoize.jl b/src/Memoize.jl index 132d358..f1cc925 100644 --- a/src/Memoize.jl +++ b/src/Memoize.jl @@ -1,4 +1,5 @@ module Memoize +using MacroTools: splitdef, combinedef, splitarg export @memoize macro memoize(args...) @@ -10,85 +11,39 @@ macro memoize(args...) else error("Memoize accepts at most two arguments") end - # a return type declaration of Any is a No-op because everything is <: Any - rettype = Any - # if the return type is provided we need to strip it out and put it back later - if ex.args[1].head == :(::) - rettype = ex.args[1].args[2] - ex.args[1] = ex.args[1].args[1] - end - # error handling for expressions that are not method definitions - if !isa(ex,Expr) || (ex.head != :function && ex.head != Symbol("=")) || - isempty(ex.args) || ex.args[1].head != :call || isempty(ex.args[1].args) + + def_dict = try + splitdef(ex) + catch error("@memoize must be applied to a method definition") end - f = ex.args[1].args[1] - ex.args[1].args[1] = u = Symbol("##",f,"_unmemoized") - - args = ex.args[1].args[2:end] - - # Extract keywords from AST - kws = Any[] - vals = copy(args) - if length(vals) > 0 && isa(vals[1], Expr) && vals[1].head == :parameters - kws = shift!(vals).args - end + # a return type declaration of Any is a No-op because everything is <: Any + rettype = get(def_dict, :rtype, Any) + f = def_dict[:name] + def_dict_unmemoized = copy(def_dict) + def_dict_unmemoized[:name] = u = Symbol("##",f,"_unmemoized") - # Set up arguments for tuple to encode keywords - tup = Array{Any}(length(kws)+length(vals)) - i = 1 - for val in vals - tup[i] = if isa(val, Expr) - if val.head == :... || val.head == :kw - val.args[1] - elseif val.head == :(::) - val - else - error("@memoize did not understand method syntax $val") - end - else - val - end - i += 1 - end - - for kw in kws - if isa(kw, Expr) && (kw.head == :kw || kw.head == :...) - tup[i] = kw.args[1] - else - error("@memoize did not understand method syntax") - end - i += 1 - end + args = def_dict[:args] + kws = def_dict[:kwargs] + # Set up arguments for tuple + tup = [splitarg(arg)[1] for arg in vcat(args, kws)] # Set up identity arguments to pass to unmemoized function - identargs = Array{Any}((length(kws) > 0)+length(vals)) - i = (length(kws) > 0) + 1 - for val in vals - if isa(val, Expr) - if val.head == :kw - val = val.args[1] - end - if isa(val, Expr) && val.head == :(::) - val = val.args[1] - end + identargs = map(args) do arg + arg_name, typ, slurp, default = splitarg(arg) + if slurp + Expr(:..., arg_name) + else + arg_name end - identargs[i] = val - i += 1 end - if length(kws) > 0 - identkws = map(kws) do kw - if kw.head == :kw - key = kw.args[1] - if isa(key, Expr) && key.head == :(::) - key = key.args[1] - end - Expr(:kw, key, key) - else - kw - end + identkws = map(kws) do kw + arg_name, typ, slurp, default = splitarg(kw) + if slurp + Expr(:..., arg_name) + else + Expr(:kw, arg_name, arg_name) end - identargs[1] = Expr(:parameters, identkws...) end fcachename = Symbol("##",f,"_memoized_cache") @@ -103,13 +58,15 @@ macro memoize(args...) lookup = :($fcache[($(tup...),)]) end + def_dict[:body] = quote + haskey($fcache, ($(tup...),)) ? $lookup : + ($fcache[($(tup...),)] = $u($(identargs...),; $(identkws...))) + end esc(quote - $ex + $(combinedef(def_dict_unmemoized)) empty!($fcache) - $f($(args...),)::$rettype = - haskey($fcache, ($(tup...),)) ? $lookup : - ($fcache[($(tup...),)] = $u($(identargs...),)) + $(combinedef(def_dict)) end) - + end end diff --git a/test/runtests.jl b/test/runtests.jl index 565571a..a172034 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -194,6 +194,18 @@ end @test multiple_dispatch(1.0) == 2 @test run == 2 +if VERSION >= v"0.6.0" + run = 0 + @memoize function where_clause(a::T) where T + global run += 1 + T + end + @test where_clause(1) == Int + @test run == 1 + @test where_clause(1) == Int + @test run == 1 +end + function outer() run = 0 @memoize function inner(x)