From f7428926682ba53add64154749ccbca7c245301b Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Tue, 5 Jan 2021 17:31:14 -0500
Subject: [PATCH 01/16] Add support for method-wise cache invalidation.

---
 README.md        |  13 +++---
 src/Memoize.jl   | 100 +++++++++++++++++++++++++++++++++++++++--------
 test/runtests.jl |  37 ++++++++++++++++--
 3 files changed, 124 insertions(+), 26 deletions(-)

diff --git a/README.md b/README.md
index 09201cd..8e6e448 100644
--- a/README.md
+++ b/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
diff --git a/src/Memoize.jl b/src/Memoize.jl
index 5fe1858..5b3afe1 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -2,12 +2,18 @@ module Memoize
 using MacroTools: isexpr, combinedef, namify, splitarg, splitdef
 export @memoize, memoize_cache
 
-cache_name(f) = Symbol("##", f, "_memoized_cache")
+export @memoize, memories, memory
 
-function try_empty_cache(f)
-    try
-        empty!(memoize_cache(f))
-    catch
+# I would call which($sig) but it's only on 1.6 I think
+function _which(tt, world = typemax(UInt))
+    meth = ccall(:jl_gf_invoke_lookup, Any, (Any, UInt), tt, world)
+    if meth !== nothing
+        if meth isa Method
+            return meth::Method
+        else
+            meth = meth.func
+            return meth::Method
+        end
     end
 end
 
@@ -40,6 +46,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 +66,11 @@ macro memoize(args...)
         end
     end
 
-    @gensym fcache
+    @gensym cache
     mod = __module__
 
     body = quote
-        get!($fcache, ($(tup...),)) do
+        get!($cache, ($(tup...),)) do
             $u($(identargs...); $(identkws...))
         end
     end
@@ -75,23 +83,81 @@ macro memoize(args...)
         def_dict[:body] = body
     end
 
+    @gensym world
+    @gensym old_meth
+    @gensym meth
+    @gensym brain
+    @gensym old_brain
+
+    sig = :(Tuple{typeof($(def_dict[:name])), $((splitarg(arg)[2] for arg in def_dict[:args])...)} where {$(def_dict[:whereparams]...)})
+
     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 = $cache_dict
+        local $world = Base.get_world_counter()
+
         $(combinedef(def_dict_unmemoized))
-        Base.@__doc__ $(combinedef(def_dict))
+        local $result = Base.@__doc__($(combinedef(def_dict)))
+
+        local $brain = if isdefined($__module__, :__Memoize_brain__)
+            brain = getfield($__module__, :__Memoize_brain__)
+        else
+            global __Memoize_brain__ = Dict()
+        end
+        
+        # If overwriting a method, empty the old cache.
+        # Notice that methods are hashed by their stored signature
+        local $old_meth = $_which($sig, $world)
+        if $old_meth !== nothing && $old_meth.sig == $sig
+            if isdefined($old_meth.module, :__Memoize_brain__)
+                $old_brain = getfield($old_meth.module, :__Memoize_brain__)
+                empty!(pop!($old_brain, $old_meth.sig, []))
+            end
+        end
+
+        # Store the cache so that it can be emptied later
+        local $meth = $_which($sig)
+        @assert $meth !== nothing
+        $brain[$meth.sig] = $cache
+
+        $result
     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.
-    getproperty(parentmodule(f), cache_name(f))
+"""
+    memories(f, [types], [module])
+    
+    Return an array of memoized method caches for the function f.
+    
+    This function takes the same arguments as the method methods.
+"""
+memories(f, args...) = _memories(methods(f, args...))
+
+function _memories(ms::Base.MethodList)
+    memories = []
+    for m in ms
+        if isdefined(m.module, :__Memoize_brain__)
+            brain = getfield(m.module, :__Memoize_brain__)
+            memory = get(brain, m.sig, nothing)
+            if memory !== nothing
+                push!(memories, memory)
+            end
+        end
+    end
+    return memories
+end
+
+"""
+    memory(m)
+    
+    Return the memoized cache for the method m, or nothing if no such method exists
+"""
+function memory(m::Method)
+    if isdefined(m.module, :__Memoize_brain__)
+        brain = getfield(m.module, :__Memoize_brain__)
+        return get(brain, m.sig, nothing)
+    end
 end
 
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 08ce032..9955804 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -29,7 +29,7 @@ end
 @test simple(6) == 6
 @test run == 2
 
-empty!(memoize_cache(simple))
+map(empty!, memories(simple))
 @test simple(6) == 6
 @test run == 3
 @test simple(6) == 6
@@ -254,6 +254,37 @@ end
 outer()
 @test !@isdefined inner
 
+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(empty!, memories(genspec, Tuple{Int}))
+@test genspec(5) == 7
+@test genrun == 2
+@test specrun == 2
+
 @memoize function typeinf(x)
     x + 1
 end
@@ -328,13 +359,13 @@ end # module
 using .MemoizeTest
 using .MemoizeTest: custom_dict
 
-empty!(memoize_cache(custom_dict))
+map(empty!, memories(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(empty!, memories(MemoizeTest.custom_dict))
 @test custom_dict(1) == 1
 @test MemoizeTest.run == 4
 

From 5b96b7f4f0819bcbe2176e24cdada5571a4ac31a Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Tue, 5 Jan 2021 17:52:47 -0500
Subject: [PATCH 02/16] typos

---
 src/Memoize.jl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 5b3afe1..3a5a096 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -1,7 +1,5 @@
 module Memoize
 using MacroTools: isexpr, combinedef, namify, splitarg, splitdef
-export @memoize, memoize_cache
-
 export @memoize, memories, memory
 
 # I would call which($sig) but it's only on 1.6 I think

From 4ec35f52a8877748928f5428356a364bfc3e3074 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Tue, 5 Jan 2021 18:12:45 -0500
Subject: [PATCH 03/16] factor out cache lookup

---
 src/Memoize.jl | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 3a5a096..7ff7d96 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -120,7 +120,6 @@ macro memoize(args...)
 
         $result
     end)
-
 end
 
 """
@@ -135,13 +134,8 @@ memories(f, args...) = _memories(methods(f, args...))
 function _memories(ms::Base.MethodList)
     memories = []
     for m in ms
-        if isdefined(m.module, :__Memoize_brain__)
-            brain = getfield(m.module, :__Memoize_brain__)
-            memory = get(brain, m.sig, nothing)
-            if memory !== nothing
-                push!(memories, memory)
-            end
-        end
+        cache = memory(m)
+        cache !== nothing && push!(memories, cache)
     end
     return memories
 end
@@ -156,6 +150,7 @@ function memory(m::Method)
         brain = getfield(m.module, :__Memoize_brain__)
         return get(brain, m.sig, nothing)
     end
+    return nothing
 end
 
 end

From 28652adcbd308233fa5be67c8ed5ee594b0ab29a Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Fri, 8 Jan 2021 16:28:10 -0500
Subject: [PATCH 04/16] style change

---
 src/Memoize.jl | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 7ff7d96..42b43de 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -97,10 +97,12 @@ macro memoize(args...)
         $(combinedef(def_dict_unmemoized))
         local $result = Base.@__doc__($(combinedef(def_dict)))
 
-        local $brain = if isdefined($__module__, :__Memoize_brain__)
-            brain = getfield($__module__, :__Memoize_brain__)
+        if isdefined($__module__, :__Memoize_brain__)
+            local $brain = $__module__.__Memoize_brain__
         else
             global __Memoize_brain__ = Dict()
+            local $brain = __Memoize_brain__
+            $__module__
         end
         
         # If overwriting a method, empty the old cache.
@@ -108,7 +110,7 @@ macro memoize(args...)
         local $old_meth = $_which($sig, $world)
         if $old_meth !== nothing && $old_meth.sig == $sig
             if isdefined($old_meth.module, :__Memoize_brain__)
-                $old_brain = getfield($old_meth.module, :__Memoize_brain__)
+                $old_brain = $old_meth.module.__Memoize_brain__
                 empty!(pop!($old_brain, $old_meth.sig, []))
             end
         end
@@ -147,7 +149,7 @@ end
 """
 function memory(m::Method)
     if isdefined(m.module, :__Memoize_brain__)
-        brain = getfield(m.module, :__Memoize_brain__)
+        brain = m.module.__Memoize_brain__
         return get(brain, m.sig, nothing)
     end
     return nothing

From 4230149b576a3493be62da5aa698cb5283d270d3 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Fri, 8 Jan 2021 18:39:49 -0500
Subject: [PATCH 05/16] passes precompile test

---
 Project.toml                              | 3 ++-
 test/TestPrecompile/Project.toml          | 7 +++++++
 test/TestPrecompile/src/TestPrecompile.jl | 5 +++++
 test/runtests.jl                          | 8 +++++++-
 4 files changed, 21 insertions(+), 2 deletions(-)
 create mode 100644 test/TestPrecompile/Project.toml
 create mode 100644 test/TestPrecompile/src/TestPrecompile.jl

diff --git a/Project.toml b/Project.toml
index 3e26e12..9531f4b 100644
--- a/Project.toml
+++ b/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"]
diff --git a/test/TestPrecompile/Project.toml b/test/TestPrecompile/Project.toml
new file mode 100644
index 0000000..e522571
--- /dev/null
+++ b/test/TestPrecompile/Project.toml
@@ -0,0 +1,7 @@
+name = "TestPrecompile"
+uuid = "ff0854c6-fbc5-40f0-bd2d-de277d7b8f28"
+authors = ["Peter Ahrens <ptrahrens@gmail.com>"]
+version = "0.1.0"
+
+[deps]
+Memoize = "c03570c3-d221-55d1-a50c-7939bbd78826"
\ No newline at end of file
diff --git a/test/TestPrecompile/src/TestPrecompile.jl b/test/TestPrecompile/src/TestPrecompile.jl
new file mode 100644
index 0000000..09482e0
--- /dev/null
+++ b/test/TestPrecompile/src/TestPrecompile.jl
@@ -0,0 +1,5 @@
+module TestPrecompile
+    using Memoize
+    @memoize forgetful(x) = true
+    forgetful(true)
+end # module
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 9955804..0eb1340 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,4 +1,4 @@
-using Memoize, Test
+using Memoize, Test, Pkg
 
 @test_throws LoadError eval(:(@memoize))
 @test_throws LoadError eval(:(@memoize () = ()))
@@ -369,6 +369,12 @@ map(empty!, memories(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 length(memories(TestPrecompile.forgetful)) >= 1
+
 run = 0
 @memoize Dict{Tuple{String},Int}() function dict_call(a::String)::Int
     global run += 1

From 3661204c96503b1a83b9032574e7ae3299273674 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Sat, 9 Jan 2021 14:53:48 -0500
Subject: [PATCH 06/16] Test precompile limitations

---
 test/TestPrecompile/src/TestPrecompile.jl   | 11 ++++++++---
 test/TestPrecompile2/Project.toml           |  7 +++++++
 test/TestPrecompile2/src/TestPrecompile2.jl |  6 ++++++
 test/runtests.jl                            | 11 ++++++++++-
 4 files changed, 31 insertions(+), 4 deletions(-)
 create mode 100644 test/TestPrecompile2/Project.toml
 create mode 100644 test/TestPrecompile2/src/TestPrecompile2.jl

diff --git a/test/TestPrecompile/src/TestPrecompile.jl b/test/TestPrecompile/src/TestPrecompile.jl
index 09482e0..6e21849 100644
--- a/test/TestPrecompile/src/TestPrecompile.jl
+++ b/test/TestPrecompile/src/TestPrecompile.jl
@@ -1,5 +1,10 @@
 module TestPrecompile
     using Memoize
-    @memoize forgetful(x) = true
-    forgetful(true)
-end # module
\ No newline at end of file
+    run = 0
+    @memoize function forgetful(x)
+        global run += 1
+        return true
+    end
+        
+    forgetful(1)
+end # module
diff --git a/test/TestPrecompile2/Project.toml b/test/TestPrecompile2/Project.toml
new file mode 100644
index 0000000..abb66cd
--- /dev/null
+++ b/test/TestPrecompile2/Project.toml
@@ -0,0 +1,7 @@
+name = "TestPrecompile2"
+uuid = "7fd9c7c1-bae8-496a-aa66-4a9878cd045a"
+authors = ["Peter Ahrens <ptrahrens@gmail.com>"]
+version = "0.1.0"
+
+[deps]
+TestPrecompile = "ff0854c6-fbc5-40f0-bd2d-de277d7b8f28"
diff --git a/test/TestPrecompile2/src/TestPrecompile2.jl b/test/TestPrecompile2/src/TestPrecompile2.jl
new file mode 100644
index 0000000..503193f
--- /dev/null
+++ b/test/TestPrecompile2/src/TestPrecompile2.jl
@@ -0,0 +1,6 @@
+module TestPrecompile2
+
+using TestPrecompile
+TestPrecompile.forgetful(2)
+
+end # module
diff --git a/test/runtests.jl b/test/runtests.jl
index 0eb1340..9e59a40 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -373,7 +373,16 @@ Pkg.activate(temp=true)
 Pkg.develop(path=joinpath(@__DIR__, "TestPrecompile"))
 using TestPrecompile
 
-@test length(memories(TestPrecompile.forgetful)) >= 1
+@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

From b5e6dc381282ffcd4093e2c3e561f765554597bd Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Sat, 9 Jan 2021 15:03:51 -0500
Subject: [PATCH 07/16] one more quick test

---
 test/runtests.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/runtests.jl b/test/runtests.jl
index 9e59a40..f8deb3c 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -373,6 +373,7 @@ Pkg.activate(temp=true)
 Pkg.develop(path=joinpath(@__DIR__, "TestPrecompile"))
 using TestPrecompile
 
+@test length(memories(TestPrecompile.forgetful)) == 1
 @test TestPrecompile.run == 1
 @test TestPrecompile.forgetful(1)
 @test TestPrecompile.run == 1

From 24ee4d33bc24e697661b198127ec99efe8641dc0 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Sat, 9 Jan 2021 15:26:47 -0500
Subject: [PATCH 08/16] make the description of memories more correct.

---
 src/Memoize.jl | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 42b43de..10f1c80 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -1,6 +1,6 @@
 module Memoize
 using MacroTools: isexpr, combinedef, namify, splitarg, splitdef
-export @memoize, memories, memory
+export @memoize, memories
 
 # I would call which($sig) but it's only on 1.6 I think
 function _which(tt, world = typemax(UInt))
@@ -127,27 +127,29 @@ end
 """
     memories(f, [types], [module])
     
-    Return an array of memoized method caches for the function f.
+    Return an array containing all the memoized method caches for the function f.
+    May also contain caches of overwritten methods.
     
     This function takes the same arguments as the method methods.
 """
 memories(f, args...) = _memories(methods(f, args...))
 
 function _memories(ms::Base.MethodList)
-    memories = []
+    caches = []
     for m in ms
-        cache = memory(m)
-        cache !== nothing && push!(memories, cache)
+        cache = memories(m)
+        cache !== nothing && push!(caches, cache)
     end
-    return memories
+    return caches
 end
 
 """
-    memory(m)
+    memories(m::Method)
     
-    Return the memoized cache for the method m, or nothing if no such method exists
+    If m has not been overwritten, return it's memoized cache. Otherwise,
+    return nothing or the cache of an overwritten method.
 """
-function memory(m::Method)
+function memories(m::Method)
     if isdefined(m.module, :__Memoize_brain__)
         brain = m.module.__Memoize_brain__
         return get(brain, m.sig, nothing)

From ddf9b4d6a3df24553c5f6d2c10131d6bd2d04c06 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Wed, 13 Jan 2021 18:36:39 -0500
Subject: [PATCH 09/16] simplified a bit

---
 src/Memoize.jl | 30 +++++++++++-------------------
 1 file changed, 11 insertions(+), 19 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 10f1c80..a283e3d 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -2,7 +2,8 @@ module Memoize
 using MacroTools: isexpr, combinedef, namify, splitarg, splitdef
 export @memoize, memories
 
-# I would call which($sig) but it's only on 1.6 I think
+# which(signature::Tuple) is only on 1.6, but previous julia versions
+# use the following code under the hood anyway.
 function _which(tt, world = typemax(UInt))
     meth = ccall(:jl_gf_invoke_lookup, Any, (Any, UInt), tt, world)
     if meth !== nothing
@@ -82,10 +83,7 @@ macro memoize(args...)
     end
 
     @gensym world
-    @gensym old_meth
     @gensym meth
-    @gensym brain
-    @gensym old_brain
 
     sig = :(Tuple{typeof($(def_dict[:name])), $((splitarg(arg)[2] for arg in def_dict[:args])...)} where {$(def_dict[:whereparams]...)})
 
@@ -97,28 +95,23 @@ macro memoize(args...)
         $(combinedef(def_dict_unmemoized))
         local $result = Base.@__doc__($(combinedef(def_dict)))
 
-        if isdefined($__module__, :__Memoize_brain__)
-            local $brain = $__module__.__Memoize_brain__
-        else
-            global __Memoize_brain__ = Dict()
-            local $brain = __Memoize_brain__
-            $__module__
+        if !isdefined($__module__, :__memories__)
+            global __memories__ = Dict()
         end
         
         # If overwriting a method, empty the old cache.
         # Notice that methods are hashed by their stored signature
-        local $old_meth = $_which($sig, $world)
-        if $old_meth !== nothing && $old_meth.sig == $sig
-            if isdefined($old_meth.module, :__Memoize_brain__)
-                $old_brain = $old_meth.module.__Memoize_brain__
-                empty!(pop!($old_brain, $old_meth.sig, []))
+        local $meth = $_which($sig, $world)
+        if $meth !== nothing && $meth.sig == $sig
+            if isdefined($meth.module, :__memories__)
+                empty!(pop!($meth.module.__memories__, $meth.sig, []))
             end
         end
 
         # Store the cache so that it can be emptied later
         local $meth = $_which($sig)
         @assert $meth !== nothing
-        $brain[$meth.sig] = $cache
+        $__module__.__memories__[$meth.sig] = $cache
 
         $result
     end)
@@ -150,9 +143,8 @@ end
     return nothing or the cache of an overwritten method.
 """
 function memories(m::Method)
-    if isdefined(m.module, :__Memoize_brain__)
-        brain = m.module.__Memoize_brain__
-        return get(brain, m.sig, nothing)
+    if isdefined(m.module, :__memories__)
+        return get(m.module.__memories__, m.sig, nothing)
     end
     return nothing
 end

From fe871aebd40bf0b011e8ee6731b4acce92112f49 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Sun, 24 Jan 2021 11:26:07 -0500
Subject: [PATCH 10/16] use local scope to store closure records

---
 src/Memoize.jl | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index a283e3d..8a6f4f7 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -95,15 +95,17 @@ macro memoize(args...)
         $(combinedef(def_dict_unmemoized))
         local $result = Base.@__doc__($(combinedef(def_dict)))
 
-        if !isdefined($__module__, :__memories__)
-            global __memories__ = Dict()
+        if !@isdefined(__memories__)
+            __memories__ = Dict()
         end
         
         # If overwriting a method, empty the old cache.
         # Notice that methods are hashed by their stored signature
         local $meth = $_which($sig, $world)
         if $meth !== nothing && $meth.sig == $sig
-            if isdefined($meth.module, :__memories__)
+            if $meth.module == $__module__ && @isdefined(__memories__)
+                empty!(pop!(__memories__, $meth.sig, []))
+            elseif isdefined($meth.module, :__memories__)
                 empty!(pop!($meth.module.__memories__, $meth.sig, []))
             end
         end
@@ -111,7 +113,7 @@ macro memoize(args...)
         # Store the cache so that it can be emptied later
         local $meth = $_which($sig)
         @assert $meth !== nothing
-        $__module__.__memories__[$meth.sig] = $cache
+        __memories__[$meth.sig] = $cache
 
         $result
     end)

From cb11d0c3ffe5a583cfe52be1a9859cd2d053d181 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Sun, 24 Jan 2021 13:20:51 -0500
Subject: [PATCH 11/16] interesting.

---
 src/Memoize.jl   |  9 +++++----
 test/runtests.jl | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 8a6f4f7..aa12430 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -122,8 +122,8 @@ end
 """
     memories(f, [types], [module])
     
-    Return an array containing all the memoized method caches for the function f.
-    May also contain caches of overwritten methods.
+    Return an array containing all the memoized method caches for the function f
+    defined at global scope. May also contain caches of overwritten methods.
     
     This function takes the same arguments as the method methods.
 """
@@ -141,8 +141,9 @@ end
 """
     memories(m::Method)
     
-    If m has not been overwritten, return it's memoized cache. Otherwise,
-    return nothing or the cache of an overwritten method.
+    If m, defined at global scope, has not been overwritten, return it's
+    memoized cache. Otherwise, return nothing or the cache of an overwritten
+    method.
 """
 function memories(m::Method)
     if isdefined(m.module, :__memories__)
diff --git a/test/runtests.jl b/test/runtests.jl
index f8deb3c..e32c22e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -254,6 +254,38 @@ 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
+    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_2(7) == (8, 42, 3)
+@test inner_2(5) == (6, 42, 1)
+@test inner_2(6) == (7, 42, 2)
+
 genrun = 0
 @memoize function genspec(a)
     global genrun += 1

From a6b3a1a817ffa1aa723de7863121fdf06b8278ce Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Sun, 24 Jan 2021 13:24:12 -0500
Subject: [PATCH 12/16] remove spurious isdefined check

---
 src/Memoize.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index aa12430..2f0900c 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -103,7 +103,7 @@ macro memoize(args...)
         # Notice that methods are hashed by their stored signature
         local $meth = $_which($sig, $world)
         if $meth !== nothing && $meth.sig == $sig
-            if $meth.module == $__module__ && @isdefined(__memories__)
+            if $meth.module == $__module__
                 empty!(pop!(__memories__, $meth.sig, []))
             elseif isdefined($meth.module, :__memories__)
                 empty!(pop!($meth.module.__memories__, $meth.sig, []))

From 4e9ab008f96dafa1dacec937fb5202160ca69f92 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Sun, 24 Jan 2021 14:41:16 -0500
Subject: [PATCH 13/16] scope-speficity to avoid querying world age and calling
 which in local functions for no good reason.

---
 src/Memoize.jl | 50 +++++++++++++++++++++++++++-----------------------
 1 file changed, 27 insertions(+), 23 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 2f0900c..5d3e2e1 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -4,8 +4,8 @@ export @memoize, memories
 
 # which(signature::Tuple) is only on 1.6, but previous julia versions
 # use the following code under the hood anyway.
-function _which(tt, world = typemax(UInt))
-    meth = ccall(:jl_gf_invoke_lookup, Any, (Any, UInt), tt, world)
+function _which(tt)
+    meth = ccall(:jl_gf_invoke_lookup, Any, (Any, UInt), tt, typemax(UInt))
     if meth !== nothing
         if meth isa Method
             return meth::Method
@@ -82,38 +82,42 @@ macro memoize(args...)
         def_dict[:body] = body
     end
 
-    @gensym world
-    @gensym meth
-
     sig = :(Tuple{typeof($(def_dict[:name])), $((splitarg(arg)[2] for arg in def_dict[:args])...)} where {$(def_dict[:whereparams]...)})
 
+    scope = gensym()
+    meth = gensym("meth")
+
     esc(quote
         # The `local` qualifier will make this performant even in the global scope.
         local $cache = $cache_dict
-        local $world = Base.get_world_counter()
 
-        $(combinedef(def_dict_unmemoized))
-        local $result = Base.@__doc__($(combinedef(def_dict)))
+        $scope = nothing
 
-        if !@isdefined(__memories__)
-            __memories__ = Dict()
-        end
-        
-        # If overwriting a method, empty the old cache.
-        # Notice that methods are hashed by their stored signature
-        local $meth = $_which($sig, $world)
-        if $meth !== nothing && $meth.sig == $sig
-            if $meth.module == $__module__
-                empty!(pop!(__memories__, $meth.sig, []))
-            elseif isdefined($meth.module, :__memories__)
+        if isdefined($__module__, $(QuoteNode(scope)))
+            if !@isdefined($(def_dict[:name]))
+                function $(def_dict[:name]) end
+            end
+
+            # If overwriting a method, empty the old cache.
+            # Notice that methods are hashed by their stored signature
+            local $meth = $_which($sig)
+            if $meth !== nothing && $meth.sig == $sig && isdefined($meth.module, :__memories__)
                 empty!(pop!($meth.module.__memories__, $meth.sig, []))
             end
         end
 
-        # Store the cache so that it can be emptied later
-        local $meth = $_which($sig)
-        @assert $meth !== nothing
-        __memories__[$meth.sig] = $cache
+        $(combinedef(def_dict_unmemoized))
+        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($sig)
+            @assert $meth !== nothing
+            __memories__[$meth.sig] = $cache
+        end
 
         $result
     end)

From 04af86b87363eb0ccc95ba4ff219ea6564b920b0 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Sun, 24 Jan 2021 15:50:54 -0500
Subject: [PATCH 14/16] sneaky way to invalidate caches in closure scope.

---
 src/Memoize.jl   | 49 ++++++++++++++++++++++++------------------------
 test/runtests.jl | 20 +++++++++++++++-----
 2 files changed, 39 insertions(+), 30 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 5d3e2e1..7f029d9 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -1,6 +1,6 @@
 module Memoize
 using MacroTools: isexpr, combinedef, namify, splitarg, splitdef
-export @memoize, memories
+export @memoize, forget!
 
 # which(signature::Tuple) is only on 1.6, but previous julia versions
 # use the following code under the hood anyway.
@@ -65,11 +65,11 @@ macro memoize(args...)
         end
     end
 
-    @gensym cache
+    cache = gensym(:__cache__)
     mod = __module__
 
     body = quote
-        get!($cache, ($(tup...),)) do
+        get!($cache[2], ($(tup...),)) do
             $u($(identargs...); $(identkws...))
         end
     end
@@ -83,13 +83,14 @@ macro memoize(args...)
     end
 
     sig = :(Tuple{typeof($(def_dict[:name])), $((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
         # The `local` qualifier will make this performant even in the global scope.
-        local $cache = $cache_dict
+        local $cache = ($tail, $cache_dict)
 
         $scope = nothing
 
@@ -102,7 +103,7 @@ macro memoize(args...)
             # Notice that methods are hashed by their stored signature
             local $meth = $_which($sig)
             if $meth !== nothing && $meth.sig == $sig && isdefined($meth.module, :__memories__)
-                empty!(pop!($meth.module.__memories__, $meth.sig, []))
+                empty!(pop!($meth.module.__memories__, $meth.sig, (nothing, []))[2])
             end
         end
 
@@ -124,36 +125,34 @@ macro memoize(args...)
 end
 
 """
-    memories(f, [types], [module])
+    forget!(f, types)
     
-    Return an array containing all the memoized method caches for the function f
-    defined at global scope. May also contain caches of overwritten methods.
-    
-    This function takes the same arguments as the method methods.
+    If the method `which(f, types)`, is memoized, `empty!` its cache in the
+    scope of `f`.
 """
-memories(f, args...) = _memories(methods(f, args...))
-
-function _memories(ms::Base.MethodList)
-    caches = []
-    for m in ms
-        cache = memories(m)
-        cache !== nothing && push!(caches, cache)
+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
-    return caches
+    forget!(which(f, types)) #otherwise, a method would suffice
 end
 
 """
-    memories(m::Method)
+    forget!(m::Method)
     
-    If m, defined at global scope, has not been overwritten, return it's
-    memoized cache. Otherwise, return nothing or the cache of an overwritten
-    method.
+    If m, defined at global scope, is a memoized function, `empty!` its
+    cache.
 """
-function memories(m::Method)
+function forget!(m::Method)
     if isdefined(m.module, :__memories__)
-        return get(m.module.__memories__, m.sig, nothing)
+        empty!(get(m.module.__memories__, m.sig, (nothing, []))[2])
     end
-    return nothing
 end
 
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index e32c22e..5131379 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -29,7 +29,7 @@ end
 @test simple(6) == 6
 @test run == 2
 
-map(empty!, memories(simple))
+map(forget!, methods(simple))
 @test simple(6) == 6
 @test run == 3
 @test simple(6) == 6
@@ -274,6 +274,10 @@ function outer_overwrite(y)
     @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
 
@@ -282,9 +286,16 @@ 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)
@@ -312,7 +323,7 @@ end
 @test genrun == 2
 @test specrun == 1
 
-map(empty!, memories(genspec, Tuple{Int}))
+map(forget!, methods(genspec, Tuple{Int}))
 @test genspec(5) == 7
 @test genrun == 2
 @test specrun == 2
@@ -391,13 +402,13 @@ end # module
 using .MemoizeTest
 using .MemoizeTest: custom_dict
 
-map(empty!, memories(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
 
-map(empty!, memories(MemoizeTest.custom_dict))
+map(forget!, methods(MemoizeTest.custom_dict))
 @test custom_dict(1) == 1
 @test MemoizeTest.run == 4
 
@@ -405,7 +416,6 @@ Pkg.activate(temp=true)
 Pkg.develop(path=joinpath(@__DIR__, "TestPrecompile"))
 using TestPrecompile
 
-@test length(memories(TestPrecompile.forgetful)) == 1
 @test TestPrecompile.run == 1
 @test TestPrecompile.forgetful(1)
 @test TestPrecompile.run == 1

From 6de8680197626b36f7b1bcb4010dee8e4ce1d537 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Sun, 24 Jan 2021 15:52:05 -0500
Subject: [PATCH 15/16] in the event we rename our function, we must define it
 first.

---
 src/Memoize.jl | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 7f029d9..58dcfa4 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -95,9 +95,7 @@ macro memoize(args...)
         $scope = nothing
 
         if isdefined($__module__, $(QuoteNode(scope)))
-            if !@isdefined($(def_dict[:name]))
-                function $(def_dict[:name]) end
-            end
+            function $(def_dict[:name]) end
 
             # If overwriting a method, empty the old cache.
             # Notice that methods are hashed by their stored signature

From 6057d381a68778bd02292bc371a7955edcddb71c Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Sun, 24 Jan 2021 17:20:14 -0500
Subject: [PATCH 16/16] simplify by removing custom which

---
 src/Memoize.jl | 31 ++++++++++---------------------
 1 file changed, 10 insertions(+), 21 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 58dcfa4..345e217 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -2,20 +2,6 @@ module Memoize
 using MacroTools: isexpr, combinedef, namify, splitarg, splitdef
 export @memoize, forget!
 
-# which(signature::Tuple) is only on 1.6, but previous julia versions
-# use the following code under the hood anyway.
-function _which(tt)
-    meth = ccall(:jl_gf_invoke_lookup, Any, (Any, UInt), tt, typemax(UInt))
-    if meth !== nothing
-        if meth isa Method
-            return meth::Method
-        else
-            meth = meth.func
-            return meth::Method
-        end
-    end
-end
-
 macro memoize(args...)
     if length(args) == 1
         dicttype = :(IdDict)
@@ -82,7 +68,8 @@ macro memoize(args...)
         def_dict[:body] = body
     end
 
-    sig = :(Tuple{typeof($(def_dict[:name])), $((splitarg(arg)[2] for arg in def_dict[:args])...)} where {$(def_dict[:whereparams]...)})
+    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()
@@ -95,13 +82,16 @@ macro memoize(args...)
         $scope = nothing
 
         if isdefined($__module__, $(QuoteNode(scope)))
-            function $(def_dict[:name]) end
+            function $f end
 
             # If overwriting a method, empty the old cache.
             # Notice that methods are hashed by their stored signature
-            local $meth = $_which($sig)
-            if $meth !== nothing && $meth.sig == $sig && isdefined($meth.module, :__memories__)
-                empty!(pop!($meth.module.__memories__, $meth.sig, (nothing, []))[2])
+            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
 
@@ -113,8 +103,7 @@ macro memoize(args...)
                 __memories__ = Dict()
             end
             # Store the cache so that it can be emptied later
-            local $meth = $_which($sig)
-            @assert $meth !== nothing
+            local $meth = $which($f, $tail)
             __memories__[$meth.sig] = $cache
         end