diff --git a/src/Memoize.jl b/src/Memoize.jl index 6406c58..5fe1858 100644 --- a/src/Memoize.jl +++ b/src/Memoize.jl @@ -4,6 +4,13 @@ export @memoize, memoize_cache cache_name(f) = Symbol("##", f, "_memoized_cache") +function try_empty_cache(f) + try + empty!(memoize_cache(f)) + catch + end +end + macro memoize(args...) if length(args) == 1 dicttype = :(IdDict) @@ -51,11 +58,8 @@ macro memoize(args...) end end - fcachename = cache_name(f) + @gensym fcache mod = __module__ - fcache = isdefined(mod, fcachename) ? - getfield(mod, fcachename) : - Core.eval(mod, :(const $fcachename = $cache_dict)) body = quote get!($fcache, ($(tup...),)) do @@ -72,17 +76,21 @@ macro memoize(args...) end esc(quote + $Memoize.try_empty_cache($f) # So that redefining a function doesn't leak memory through + # the previous cache. + # The `local` qualifier will make this performant even in the global scope. + local $fcache = $cache_dict + $(cache_name(f)) = $fcache # for `memoize_cache(f)` $(combinedef(def_dict_unmemoized)) - empty!($fcache) Base.@__doc__ $(combinedef(def_dict)) end) end function memoize_cache(f::Function) - # This will fail in certain circumstances (eg. @memoize Base.sin(::MyNumberType) = ...) but I don't think there's - # a clean answer here, because we can already have multiple caches for certain functions, if the methods are - # defined in different modules. + # This will fail in certain circumstances (eg. @memoize Base.sin(::MyNumberType) = ...) but I + # don't think there's a clean answer here, because we can already have multiple caches for + # certain functions, if the methods are defined in different modules. getproperty(parentmodule(f), cache_name(f)) end diff --git a/test/runtests.jl b/test/runtests.jl index d9639ed..08ce032 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -352,3 +352,5 @@ end @test dict_call("bb") == 2 @test run == 2 +@memoize non_allocating(x) = x+1 +@test @allocated(non_allocating(10)) == 0