Add support for method-wise cache invalidation.

Project.toml

@@ -15,6 +15,7 @@ julia = "1.2"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 
 [targets]
-test = ["Test"]
+test = ["Test", "Pkg"]

README.md

@@ -2,7 +2,7 @@
 
 [![Build Status](https://travis-ci.org/JuliaCollections/Memoize.jl.png?branch=master)](https://travis-ci.org/JuliaCollections/Memoize.jl) [![Coverage Status](https://coveralls.io/repos/github/JuliaCollections/Memoize.jl/badge.svg?branch=master)](https://coveralls.io/github/JuliaCollections/Memoize.jl?branch=master)
 
-Easy memoization for Julia.
+Easy method memoization for Julia.
 
 ## Usage
 
@@ -19,15 +19,16 @@ julia> x(1)
 Running
 2
 
-julia> memoize_cache(x)
-IdDict{Any,Any} with 1 entry:
-  (1,) => 2
+julia> memories(x)
+1-element Array{Any,1}:
+ IdDict{Any,Any}((1,) => 2)
 
 julia> x(1)
 2
 
-julia> empty!(memoize_cache(x))
-IdDict{Any,Any}()
+julia> map(empty!, memories(x))
+1-element Array{IdDict{Tuple{Any},Any},1}:
+ IdDict()
 
 julia> x(1)
 Running

src/Memoize.jl

@@ -1,15 +1,6 @@
 module Memoize
 using MacroTools: isexpr, combinedef, namify, splitarg, splitdef
-export @memoize, memoize_cache
-
-cache_name(f) = Symbol("##", f, "_memoized_cache")
-
-function try_empty_cache(f)
-    try
-        empty!(memoize_cache(f))
-    catch
-    end
-end
+export @memoize, forget!
 
 macro memoize(args...)
     if length(args) == 1
@@ -40,6 +31,8 @@ macro memoize(args...)
     # Set up arguments for tuple
     tup = [splitarg(arg)[1] for arg in vcat(args, kws)]
 
+    @gensym result
+
     # Set up identity arguments to pass to unmemoized function
     identargs = map(args) do arg
         arg_name, typ, slurp, default = splitarg(arg)
@@ -58,11 +51,11 @@ macro memoize(args...)
         end
     end
 
-    @gensym fcache
+    cache = gensym(:__cache__)
     mod = __module__
 
     body = quote
-        get!($fcache, ($(tup...),)) do
+        get!($cache[2], ($(tup...),)) do
             $u($(identargs...); $(identkws...))
         end
     end
@@ -75,23 +68,78 @@ macro memoize(args...)
         def_dict[:body] = body
     end
 
+    f = def_dict[:name]
+    sig = :(Tuple{typeof($f), $((splitarg(arg)[2] for arg in def_dict[:args])...)} where {$(def_dict[:whereparams]...)})
+    tail = :(Tuple{$((splitarg(arg)[2] for arg in def_dict[:args])...)} where {$(def_dict[:whereparams]...)})
+
+    scope = gensym()
+    meth = gensym("meth")
+
     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)`
+        local $cache = ($tail, $cache_dict)
+
+        $scope = nothing
+
+        if isdefined($__module__, $(QuoteNode(scope)))
+            function $f end
+
+            # If overwriting a method, empty the old cache.
+            # Notice that methods are hashed by their stored signature
+            try
+                local $meth = which($f, $tail)
+                if $meth.sig == $sig && isdefined($meth.module, :__memories__)
+                    empty!(pop!($meth.module.__memories__, $meth.sig, (nothing, []))[2])
+                end
+            catch
+            end
+        end
+
         $(combinedef(def_dict_unmemoized))
-        Base.@__doc__ $(combinedef(def_dict))
+        local $result = Base.@__doc__($(combinedef(def_dict)))
+
+        if isdefined($__module__, $(QuoteNode(scope)))
+            if !@isdefined __memories__
+                __memories__ = Dict()
+            end
+            # Store the cache so that it can be emptied later
+            local $meth = $which($f, $tail)
+            __memories__[$meth.sig] = $cache
+        end
+
+        $result
     end)
+end
 
+"""
+    forget!(f, types)
+
+    If the method `which(f, types)`, is memoized, `empty!` its cache in the
+    scope of `f`.
+"""
+function forget!(f, types)
+    for name in propertynames(f) #if f is a closure, we walk its fields
+        if first(string(name), length("##__cache__")) == "##__cache__"
+            cache = getproperty(f, name)
+            if cache isa Core.Box
+                cache = cache.contents
+            end
+            (cache[1] == types) && empty!(cache[2])
+        end
+    end
+    forget!(which(f, types)) #otherwise, a method would suffice
 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.
-    getproperty(parentmodule(f), cache_name(f))
+"""
+    forget!(m::Method)
+
+    If m, defined at global scope, is a memoized function, `empty!` its
+    cache.
+"""
+function forget!(m::Method)
+    if isdefined(m.module, :__memories__)
+        empty!(get(m.module.__memories__, m.sig, (nothing, []))[2])
+    end
 end
 
 end

test/TestPrecompile/Project.toml

@@ -0,0 +1,7 @@
+name = "TestPrecompile"
+uuid = "ff0854c6-fbc5-40f0-bd2d-de277d7b8f28"
+authors = ["Peter Ahrens <[email protected]>"]
+version = "0.1.0"
+
+[deps]
+Memoize = "c03570c3-d221-55d1-a50c-7939bbd78826"

test/TestPrecompile/src/TestPrecompile.jl

@@ -0,0 +1,10 @@
+module TestPrecompile
+    using Memoize
+    run = 0
+    @memoize function forgetful(x)
+        global run += 1
+        return true
+    end
+
+    forgetful(1)
+end # module

test/TestPrecompile2/Project.toml

@@ -0,0 +1,7 @@
+name = "TestPrecompile2"
+uuid = "7fd9c7c1-bae8-496a-aa66-4a9878cd045a"
+authors = ["Peter Ahrens <[email protected]>"]
+version = "0.1.0"
+
+[deps]
+TestPrecompile = "ff0854c6-fbc5-40f0-bd2d-de277d7b8f28"

test/TestPrecompile2/src/TestPrecompile2.jl

@@ -0,0 +1,6 @@
+module TestPrecompile2
+
+using TestPrecompile
+TestPrecompile.forgetful(2)
+
+end # module

test/runtests.jl

@@ -1,4 +1,4 @@
-using Memoize, Test
+using Memoize, Test, Pkg
 
 @test_throws LoadError eval(:(@memoize))
 @test_throws LoadError eval(:(@memoize () = ()))
@@ -29,7 +29,7 @@ end
 @test simple(6) == 6
 @test run == 2
 
-empty!(memoize_cache(simple))
+map(forget!, methods(simple))
 @test simple(6) == 6
 @test run == 3
 @test simple(6) == 6
@@ -254,6 +254,80 @@ end
 outer()
 @test !@isdefined inner
 
+function outer_overwrite(y)
+    run = 0
+    @memoize function inner(x)
+        run += 1
+        (x, y, run)
+    end
+    #note that calling inner here would result in an error,
+    #since both definitions of inner are evaluated before the
+    #body of outer_overwrite runs, and the cache for the second definition
+    #of inner has not been set up yet.
+    @memoize function inner(x)
+        run += 1
+        (x + 1, y, run)
+    end
+    @test inner(5) == (6, y, 1)
+    @test run == 1
+    @test inner(5) == (6, y, 1)
+    @test run == 1
+    @test inner(6) == (7, y, 2)
+    @test run == 2
+    @memoize function inner(x::String)
+        run += 1
+        (x, y, run)
+    end
+    return inner
+end
+
+inner_1 = outer_overwrite(7)
+inner_2 = outer_overwrite(42)
+@test inner_1(5) == (6, 7, 1)
+@test inner_1(6) == (7, 7, 2)
+@test inner_1(7) == (8, 7, 3)
+@test inner_1("hello") == ("hello", 7, 4)
+@test inner_2(7) == (8, 42, 3)
+@test inner_2(5) == (6, 42, 1)
+@test inner_2(6) == (7, 42, 2)
+@test inner_2("goodbye") == ("goodbye", 42, 4)
+@test inner_2("hello") == ("hello", 42, 5)
+forget!(inner_1, Tuple{Any})
+@test inner_1(5) == (6, 7, 5)
+@test inner_2(6) == (7, 42, 2)
+@test inner_1("hello") == ("hello", 7, 4)
+
+genrun = 0
+@memoize function genspec(a)
+    global genrun += 1
+    a + 1
+end
+specrun = 0
+@test genspec(5) == 6
+@test genrun == 1
+@test specrun == 0
+@memoize function genspec(a::Int)
+    global specrun += 1
+    a + 2
+end
+@test genspec(5) == 7
+@test genrun == 1
+@test specrun == 1
+@test genspec(5) == 7
+@test genrun == 1
+@test specrun == 1
+@test genspec(true) == 2
+@test genrun == 2
+@test specrun == 1
+@test invoke(genspec, Tuple{Any}, 5) == 6
+@test genrun == 2
+@test specrun == 1
+
+map(forget!, methods(genspec, Tuple{Int}))
+@test genspec(5) == 7
+@test genrun == 2
+@test specrun == 2
+
 @memoize function typeinf(x)
     x + 1
 end
@@ -328,16 +402,31 @@ end # module
 using .MemoizeTest
 using .MemoizeTest: custom_dict
 
-empty!(memoize_cache(custom_dict))
+map(forget!, methods(custom_dict))
 @test custom_dict(1) == 1
 @test MemoizeTest.run == 3
 @test custom_dict(1) == 1
 @test MemoizeTest.run == 3
 
-empty!(memoize_cache(MemoizeTest.custom_dict))
+map(forget!, methods(MemoizeTest.custom_dict))
 @test custom_dict(1) == 1
 @test MemoizeTest.run == 4
 
+Pkg.activate(temp=true)
+Pkg.develop(path=joinpath(@__DIR__, "TestPrecompile"))
+using TestPrecompile
+
+@test TestPrecompile.run == 1
+@test TestPrecompile.forgetful(1)
+@test TestPrecompile.run == 1
+
+Pkg.develop(path=joinpath(@__DIR__, "TestPrecompile2"))
+using TestPrecompile2
+
+@test TestPrecompile.run == 1
+@test TestPrecompile.forgetful(2)
+@test TestPrecompile.run == 2
+
 run = 0
 @memoize Dict{Tuple{String},Int}() function dict_call(a::String)::Int
     global run += 1