From db41394c11362a4a9a8bc1e7d2eb114732298335 Mon Sep 17 00:00:00 2001
From: cstjean <cedric.stjean@gmail.com>
Date: Wed, 28 Oct 2020 13:25:42 -0400
Subject: [PATCH 01/26] Get rid of eval

Fix #48, at the cost of putting the variable in a poorly-performing global.

Not sure if this is acceptable. It's frustrating that Julia seemingly lacks the tools to deal with
this elegantly.

- If `const` worked in a local function, we'd just put `const` and be done with it.
- If `typeconst` existed for global variables, that would work too.

Memoization.jl uses generated functions, which causes other problems. And it feels like the
wrong solution too.
---
 src/Memoize.jl | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 6406c58..94f2e85 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -53,9 +53,6 @@ macro memoize(args...)
 
     fcachename = cache_name(f)
     mod = __module__
-    fcache = isdefined(mod, fcachename) ?
-             getfield(mod, fcachename) :
-             Core.eval(mod, :(const $fcachename = $cache_dict))
 
     body = quote
         get!($fcache, ($(tup...),)) do
@@ -72,8 +69,10 @@ macro memoize(args...)
     end
 
     esc(quote
+        $fcachename = $cache_dict  # this should be `const` for performance, but then this
+                                   # fails the local-function cache test.
         $(combinedef(def_dict_unmemoized))
-        empty!($fcache)
+        empty!($fcachename)
         Base.@__doc__ $(combinedef(def_dict))
     end)
 

From 3dc0431a7dbd004ddf1cd9a42ff2ff4d245d7050 Mon Sep 17 00:00:00 2001
From: Cedric St-Jean <cedric.stjean@gmail.com>
Date: Wed, 30 Dec 2020 03:51:07 -0500
Subject: [PATCH 02/26] Use `local` directive so that it's also performant

at the global scope
---
 src/Memoize.jl   | 4 +---
 test/runtests.jl | 4 +++-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 94f2e85..353a113 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -69,10 +69,8 @@ macro memoize(args...)
     end
 
     esc(quote
-        $fcachename = $cache_dict  # this should be `const` for performance, but then this
-                                   # fails the local-function cache test.
+        local $fcachename = $cache_dict  # see #48 comment for `local` explanation
         $(combinedef(def_dict_unmemoized))
-        empty!($fcachename)
         Base.@__doc__ $(combinedef(def_dict))
     end)
 
diff --git a/test/runtests.jl b/test/runtests.jl
index d9639ed..38242ae 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -269,7 +269,7 @@ end
 method_rewrite()
 @memoize function method_rewrite() end
 GC.gc()
-@test finalized
+@test_broken finalized
 
 run = 0
 """ documented function """
@@ -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

From 81ddd465459a7a8d320dfc31b8aec16b255c412d Mon Sep 17 00:00:00 2001
From: Cedric St-Jean <cedric.stjean@gmail.com>
Date: Wed, 30 Dec 2020 04:25:17 -0500
Subject: [PATCH 03/26] Support memoize_cache too

---
 src/Memoize.jl | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 353a113..b9a7c11 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -51,7 +51,7 @@ macro memoize(args...)
         end
     end
 
-    fcachename = cache_name(f)
+    @gensym fcache
     mod = __module__
 
     body = quote
@@ -69,7 +69,9 @@ macro memoize(args...)
     end
 
     esc(quote
-        local $fcachename = $cache_dict  # see #48 comment for `local` explanation
+        # 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))
         Base.@__doc__ $(combinedef(def_dict))
     end)
@@ -77,9 +79,9 @@ macro memoize(args...)
 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
 

From d7bb206a87e823b3a8b5e0c8e3210ae099ed6bb0 Mon Sep 17 00:00:00 2001
From: Cedric St-Jean <cedric.stjean@gmail.com>
Date: Wed, 30 Dec 2020 04:30:51 -0500
Subject: [PATCH 04/26] Fix the broken finalized test

---
 src/Memoize.jl   | 5 +++++
 test/runtests.jl | 2 +-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index b9a7c11..8c06eb1 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -69,6 +69,11 @@ macro memoize(args...)
     end
 
     esc(quote
+        try
+            empty!(memoize_cache($f))
+        catch
+        end
+
         # The `local` qualifier will make this performant even in the global scope.
         local $fcache = $cache_dict
         $(cache_name(f)) = $fcache   # for `memoize_cache(f)`
diff --git a/test/runtests.jl b/test/runtests.jl
index 38242ae..08ce032 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -269,7 +269,7 @@ end
 method_rewrite()
 @memoize function method_rewrite() end
 GC.gc()
-@test_broken finalized
+@test finalized
 
 run = 0
 """ documented function """

From 84b61edc033459db35c117697ef793b2be9bb61f Mon Sep 17 00:00:00 2001
From: Cedric St-Jean <cedric.stjean@gmail.com>
Date: Wed, 30 Dec 2020 04:32:47 -0500
Subject: [PATCH 05/26] Comment

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

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 8c06eb1..285208f 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -70,6 +70,7 @@ macro memoize(args...)
 
     esc(quote
         try
+            # So that redefining a function doesn't leak memory through the previous cache.
             empty!(memoize_cache($f))
         catch
         end

From a99ab82051d7efaaa70975f88bdd4a4ebe418013 Mon Sep 17 00:00:00 2001
From: Cedric St-Jean <cedric.stjean@gmail.com>
Date: Fri, 1 Jan 2021 06:12:14 -0500
Subject: [PATCH 06/26] Factor out try_empty_cache.

It makes the macro expansion much more palatable
---
 src/Memoize.jl | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 285208f..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)
@@ -69,12 +76,8 @@ macro memoize(args...)
     end
 
     esc(quote
-        try
-            # So that redefining a function doesn't leak memory through the previous cache.
-            empty!(memoize_cache($f))
-        catch
-        end
-
+        $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)`

From 7de9def3f2404e2f2ac479bafdd8f882d2a2d99e Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 19:46:35 -0500
Subject: [PATCH 07/26] passes tests

---
 src/Memoize.jl   | 188 ++++++++++++++++++++++++++++++++++-------------
 test/runtests.jl |  10 +--
 2 files changed, 141 insertions(+), 57 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 5fe1858..bdf4459 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -1,97 +1,181 @@
 module Memoize
-using MacroTools: isexpr, combinedef, namify, splitarg, splitdef
-export @memoize, memoize_cache
+using MacroTools: isexpr, combinearg, combinedef, namify, splitarg, splitdef, @capture
+export @memoize, function_memories, method_memories
 
-cache_name(f) = Symbol("##", f, "_memoized_cache")
-
-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
 
+const _memories = Dict()
+
 macro memoize(args...)
     if length(args) == 1
-        dicttype = :(IdDict)
+        cache_constructor = :(IdDict())
         ex = args[1]
     elseif length(args) == 2
-        (dicttype, ex) = args
+        (cache_constructor, ex) = args
     else
         error("Memoize accepts at most two arguments")
     end
 
-    cache_dict = isexpr(dicttype, :call) ? dicttype : :(($dicttype)())
-
-    def_dict = try
+    def = try
         splitdef(ex)
     catch
         error("@memoize must be applied to a method definition")
     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")
+    # Ensure that all args have names that can be passed to the inner function
+    function tag_arg(arg)
+        arg_name, arg_type, is_splat, default = splitarg(arg)
+        arg_name === nothing && (arg_name = gensym())
+        return combinearg(arg_name, arg_type, is_splat, default)
+    end
+    args = def[:args] = map(tag_arg, def[:args])
+    kwargs = def[:kwargs] = map(tag_arg, def[:kwargs])
 
-    args = def_dict[:args]
-    kws = def_dict[:kwargs]
-    # Set up arguments for tuple
-    tup = [splitarg(arg)[1] for arg in vcat(args, kws)]
+    # Get argument types for function signature
+    arg_sigs = map(def[:args]) do arg
+        arg_name, arg_type, is_splat, default = splitarg(arg)
+        if is_splat
+            return :(Vararg{$arg_type})
+        else
+            return arg_type
+        end
+    end
+    kwarg_sigs = map(def[:args]) do arg
+        arg_name, arg_type, is_splat, default = splitarg(arg)
+        if is_splat
+            return :(Vararg{$arg_type})
+        else
+            return arg_type
+        end
+    end
 
     # Set up identity arguments to pass to unmemoized function
-    identargs = map(args) do arg
-        arg_name, typ, slurp, default = splitarg(arg)
-        if slurp || namify(typ) === :Vararg
+    pass_args = map(args) do arg
+        arg_name, arg_type, is_splat, default = splitarg(arg)
+        if is_splat || namify(arg_type) === :Vararg
             Expr(:..., arg_name)
         else
             arg_name
         end
     end
-    identkws = map(kws) do kw
-        arg_name, typ, slurp, default = splitarg(kw)
-        if slurp
-            Expr(:..., arg_name)
+    pass_kwargs = map(kwargs) do kwarg
+        kwarg_name, kwarg_type, is_splat, default = splitarg(kwarg)
+        if is_splat
+            Expr(:..., kwarg_name)
+        else
+            Expr(:kw, kwarg_name, kwarg_name)
+        end
+    end
+
+    # A return type declaration of Any is a No-op because everything is <: Any
+    return_type = get(def, :rtype, Any)
+
+    # Set up arguments for memo key
+    key_args = [splitarg(arg)[1] for arg in vcat(args, kwargs)]
+    key_arg_types = [arg_sigs; kwarg_sigs]
+
+    @gensym inner
+    inner_def = copy(def)
+    inner_def[:name] = inner
+    pop!(inner_def, :params, nothing)
+
+    @gensym result
+
+    # If this is a method of a callable object, the definition returns nothing.
+    # Thus, we must construct the type of the method on our own.
+    if haskey(def, :name)
+        if haskey(def, :params)
+            cstr_type = :($(def[:name]){$(def[:params]...)})
+            sig = :(Tuple{$cstr_type, $(arg_sigs...)} where {$(def[:whereparams]...)})
+            pushfirst!(inner_def[:args], gensym())
+            pushfirst!(pass_args, cstr_type)
+            pushfirst!(key_args, cstr_type)
+            pushfirst!(key_arg_types, :(Type{cstr_type}))
+        elseif @capture(def[:name], obj_::obj_type_)
+            obj === nothing && (obj = gensym())
+            obj_type === nothing && (obj_type = Any)
+            def[:name] = :($obj::$obj_type)
+            sig = :(Tuple{$obj_type, $(arg_sigs...)} where {$(def[:whereparams]...)})
+            pushfirst!(inner_def[:args], :($obj::$obj_type))
+            pushfirst!(pass_args, obj)
+            pushfirst!(key_args, obj)
+            pushfirst!(key_arg_types, obj_type)
         else
-            Expr(:kw, arg_name, arg_name)
+            sig = :(Tuple{typeof($(def[:name])), $(arg_sigs...)} where {$(def[:whereparams]...)})
         end
+    else
+        sig = :(Tuple{typeof($result), $(arg_sigs...)} where {$(def[:whereparams]...)})
     end
 
-    @gensym fcache
-    mod = __module__
+    @gensym cache
 
-    body = quote
-        get!($fcache, ($(tup...),)) do
-            $u($(identargs...); $(identkws...))
+    def[:body] = quote
+        $(combinedef(inner_def))
+        get!($cache, ($(key_args...),)) do
+            $inner($(pass_args...); $(pass_kwargs...))
         end
     end
 
-    if length(kws) == 0
-        def_dict[:body] = quote
-            $(body)::Core.Compiler.return_type($u, typeof(($(identargs...),)))
+    if length(kwargs) == 0
+        def[:body] = quote
+            $(def[:body])::Core.Compiler.return_type($inner, typeof(($(pass_args...),)))
         end
-    else
-        def_dict[:body] = body
     end
 
+    @gensym world
+    @gensym old_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)`
-        $(combinedef(def_dict_unmemoized))
-        Base.@__doc__ $(combinedef(def_dict))
-    end)
+        local $cache = $cache_constructor
 
+        $world = Base.get_world_counter()
+
+        $result = Base.@__doc__($(combinedef(def)))
+        
+        # If overwriting a method, empty the old cache.
+        $old_meth = $_which($sig, $world)
+        if $old_meth !== nothing
+            empty!(pop!($_memories, $old_meth, []))
+        end
+
+        # Store the cache so that it can be emptied later
+        $meth = $_which($sig)
+        @assert $meth !== nothing
+        $_memories[$meth] = $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))
+function_memories(f) = _function_memories(methods(f))
+function_memories(f, types) = _function_memories(methods(f, types))
+function_memories(f, types, mod) = _function_memories(methods(f, types, mod))
+
+function _function_memories(ms)
+    memories = []
+    for m in ms
+        memory = method_memory(m)
+        if memory !== nothing
+            push!(memories, memory)
+        end
+    end
+    return memories
 end
 
+function method_memory(m::Method)
+    return get(_memories, m, nothing)
 end
+
+end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 08ce032..93f68b2 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!, function_memories(simple))
 @test simple(6) == 6
 @test run == 3
 @test simple(6) == 6
@@ -183,7 +183,7 @@ end
 @test run == 2
 
 run = 0
-@memoize Dict function kw_ellipsis(;a...)
+@memoize Dict() function kw_ellipsis(;a...)
     global run += 1
     a
 end
@@ -308,7 +308,7 @@ using Memoize
 const MyDict = Dict
 
 run = 0
-@memoize MyDict function custom_dict(a)
+@memoize MyDict() function custom_dict(a)
     global run += 1
     a
 end
@@ -328,13 +328,13 @@ end # module
 using .MemoizeTest
 using .MemoizeTest: custom_dict
 
-empty!(memoize_cache(custom_dict))
+map(empty!, function_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!, function_memories(MemoizeTest.custom_dict))
 @test custom_dict(1) == 1
 @test MemoizeTest.run == 4
 

From d72397154e8b99db3cb63a57576991da962ef2f1 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 20:21:31 -0500
Subject: [PATCH 08/26] callable object tests

---
 src/Memoize.jl   |  2 +-
 test/runtests.jl | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index bdf4459..2734eed 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -86,7 +86,7 @@ macro memoize(args...)
     key_arg_types = [arg_sigs; kwarg_sigs]
 
     @gensym inner
-    inner_def = copy(def)
+    inner_def = deepcopy(def)
     inner_def[:name] = inner
     pop!(inner_def, :params, nothing)
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 93f68b2..6a57c78 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -254,6 +254,28 @@ end
 outer()
 @test !@isdefined inner
 
+
+run = 0
+struct callable_object
+    a
+end
+@memoize function (o::callable_object)(b)
+    global run += 1
+    (o.a, b)
+end
+@test callable_object(1)(2) == (1, 2)
+@test run == 1
+@test callable_object(1)(2) == (1, 2)
+@test run == 1
+@test callable_object(1)(3) == (1, 3)
+@test run == 2
+@test callable_object(1)(3) == (1, 3)
+@test run == 2
+@test callable_object(2)(3) == (2, 3)
+@test run == 3
+@test callable_object(2)(3) == (2, 3)
+@test run == 3
+
 @memoize function typeinf(x)
     x + 1
 end

From 7c238b4b77619025047da47ae175ac57a5cdce97 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 20:50:35 -0500
Subject: [PATCH 09/26] add trait tests

---
 src/Memoize.jl   | 34 ++++++++++++++++++++---------
 test/runtests.jl | 57 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 10 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 2734eed..24f6856 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -33,10 +33,21 @@ macro memoize(args...)
         error("@memoize must be applied to a method definition")
     end
 
+    # Set up arguments for memo key
+    key_args = []
+    key_arg_types = []
+
     # Ensure that all args have names that can be passed to the inner function
     function tag_arg(arg)
         arg_name, arg_type, is_splat, default = splitarg(arg)
-        arg_name === nothing && (arg_name = gensym())
+        if arg_name === nothing
+            arg_name = gensym()
+            push!(key_args, arg_type)
+            push!(key_arg_types, :(Type{$arg_type}))
+        else
+            push!(key_args, arg_name)
+            push!(key_arg_types, arg_type)
+        end
         return combinearg(arg_name, arg_type, is_splat, default)
     end
     args = def[:args] = map(tag_arg, def[:args])
@@ -81,10 +92,6 @@ macro memoize(args...)
     # A return type declaration of Any is a No-op because everything is <: Any
     return_type = get(def, :rtype, Any)
 
-    # Set up arguments for memo key
-    key_args = [splitarg(arg)[1] for arg in vcat(args, kwargs)]
-    key_arg_types = [arg_sigs; kwarg_sigs]
-
     @gensym inner
     inner_def = deepcopy(def)
     inner_def[:name] = inner
@@ -102,15 +109,20 @@ macro memoize(args...)
             pushfirst!(pass_args, cstr_type)
             pushfirst!(key_args, cstr_type)
             pushfirst!(key_arg_types, :(Type{cstr_type}))
-        elseif @capture(def[:name], obj_::obj_type_)
-            obj === nothing && (obj = gensym())
+        elseif @capture(def[:name], obj_::obj_type_ | ::obj_type_)
             obj_type === nothing && (obj_type = Any)
+            if obj === nothing
+                obj = gensym()
+                pushfirst!(key_args, obj_type)
+                pushfirst!(key_arg_types, :(Type{$obj_type}))
+            else
+                pushfirst!(key_args, obj)
+                pushfirst!(key_arg_types, obj_type)
+            end
             def[:name] = :($obj::$obj_type)
             sig = :(Tuple{$obj_type, $(arg_sigs...)} where {$(def[:whereparams]...)})
             pushfirst!(inner_def[:args], :($obj::$obj_type))
             pushfirst!(pass_args, obj)
-            pushfirst!(key_args, obj)
-            pushfirst!(key_arg_types, obj_type)
         else
             sig = :(Tuple{typeof($(def[:name])), $(arg_sigs...)} where {$(def[:whereparams]...)})
         end
@@ -137,7 +149,7 @@ macro memoize(args...)
     @gensym old_meth
     @gensym meth
 
-    esc(quote
+    res = esc(quote
         # The `local` qualifier will make this performant even in the global scope.
         local $cache = $cache_constructor
 
@@ -157,6 +169,8 @@ macro memoize(args...)
         $_memories[$meth] = $cache
         $result
     end)
+    #println(res)
+    res
 end
 
 function_memories(f) = _function_memories(methods(f))
diff --git a/test/runtests.jl b/test/runtests.jl
index 6a57c78..270c198 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -254,6 +254,42 @@ end
 outer()
 @test !@isdefined inner
 
+trait_function(a, ::Bool) = (-a,)
+run = 0
+@memoize function trait_function(a, ::Int)
+    global run += 1
+    (a,)
+end
+@test trait_function(1, true) == (-1,)
+@test run == 0
+@test trait_function(2, true) == (-2,)
+@test run == 0
+@test trait_function(1, 1) == (1,)
+@test run == 1
+@test trait_function(1, 2) == (1,)
+@test run == 1
+@test trait_function(2, 2) == (2,)
+@test run == 2
+@test trait_function(2, 2) == (2,)
+@test run == 2
+
+run = 0
+@memoize function trait_params(a, ::T) where {T}
+    global run += 1
+    (a, T)
+end
+@test trait_params(1, true) == (1, Bool)
+@test run == 1
+@test trait_params(1, false) == (1, Bool)
+@test run == 1
+@test trait_params(2, true) == (2, Bool)
+@test run == 2
+@test trait_params(2, false) == (2, Bool)
+@test run == 2
+@test trait_params(1, 3) == (1, Int)
+@test run == 3
+@test trait_params(1, 4) == (1, Int)
+@test run == 3
 
 run = 0
 struct callable_object
@@ -276,6 +312,27 @@ end
 @test callable_object(2)(3) == (2, 3)
 @test run == 3
 
+run = 0
+struct callable_trait_object{T}
+    a::T
+end
+@memoize function (::callable_trait_object{T})(b) where {T}
+    global run += 1
+    (T, b)
+end
+@test callable_trait_object(1)(2) == (Int, 2)
+@test run == 1
+@test callable_trait_object(2)(2) == (Int, 2)
+@test run == 1
+@test callable_trait_object(false)(2) == (Bool, 2)
+@test run == 2
+@test callable_trait_object(true)(3) == (Bool, 3)
+@test run == 3
+@test callable_trait_object(1)(3) == (Int, 3)
+@test run == 4
+@test callable_trait_object(2)(3) == (Int, 3)
+@test run == 4
+
 @memoize function typeinf(x)
     x + 1
 end

From f1427824db972727c25e97bfda9fb7ccdfd1d302 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 20:55:29 -0500
Subject: [PATCH 10/26] dead code

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

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 24f6856..c173880 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -62,14 +62,6 @@ macro memoize(args...)
             return arg_type
         end
     end
-    kwarg_sigs = map(def[:args]) do arg
-        arg_name, arg_type, is_splat, default = splitarg(arg)
-        if is_splat
-            return :(Vararg{$arg_type})
-        else
-            return arg_type
-        end
-    end
 
     # Set up identity arguments to pass to unmemoized function
     pass_args = map(args) do arg

From c303480888b60a46bb040c36623da11a9c15b95d Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 21:40:46 -0500
Subject: [PATCH 11/26] can we condense this?

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

diff --git a/src/Memoize.jl b/src/Memoize.jl
index c173880..7a0de32 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -39,7 +39,7 @@ macro memoize(args...)
 
     # Ensure that all args have names that can be passed to the inner function
     function tag_arg(arg)
-        arg_name, arg_type, is_splat, default = splitarg(arg)
+        arg_name, arg_type, slurp, default = splitarg(arg)
         if arg_name === nothing
             arg_name = gensym()
             push!(key_args, arg_type)
@@ -48,41 +48,39 @@ macro memoize(args...)
             push!(key_args, arg_name)
             push!(key_arg_types, arg_type)
         end
-        return combinearg(arg_name, arg_type, is_splat, default)
+        return combinearg(arg_name, arg_type, slurp, default)
     end
     args = def[:args] = map(tag_arg, def[:args])
     kwargs = def[:kwargs] = map(tag_arg, def[:kwargs])
 
     # Get argument types for function signature
-    arg_sigs = map(def[:args]) do arg
-        arg_name, arg_type, is_splat, default = splitarg(arg)
-        if is_splat
+    arg_sigs = Vector{Any}(map(def[:args]) do arg
+        arg_name, arg_type, slurp, default = splitarg(arg)
+        if slurp
             return :(Vararg{$arg_type})
         else
             return arg_type
         end
-    end
+    end)
 
     # Set up identity arguments to pass to unmemoized function
-    pass_args = map(args) do arg
-        arg_name, arg_type, is_splat, default = splitarg(arg)
-        if is_splat || namify(arg_type) === :Vararg
+    pass_args = Vector{Any}(map(args) do arg
+        arg_name, arg_type, slurp, default = splitarg(arg)
+        if slurp || namify(arg_type) === :Vararg
             Expr(:..., arg_name)
         else
             arg_name
         end
-    end
-    pass_kwargs = map(kwargs) do kwarg
-        kwarg_name, kwarg_type, is_splat, default = splitarg(kwarg)
-        if is_splat
+    end)
+    pass_arg_types = copy(arg_sigs)
+    pass_kwargs = Vector{Any}(map(kwargs) do kwarg
+        kwarg_name, kwarg_type, slurp, default = splitarg(kwarg)
+        if slurp
             Expr(:..., kwarg_name)
         else
             Expr(:kw, kwarg_name, kwarg_name)
         end
-    end
-
-    # A return type declaration of Any is a No-op because everything is <: Any
-    return_type = get(def, :rtype, Any)
+    end)
 
     @gensym inner
     inner_def = deepcopy(def)
@@ -99,6 +97,7 @@ macro memoize(args...)
             sig = :(Tuple{$cstr_type, $(arg_sigs...)} where {$(def[:whereparams]...)})
             pushfirst!(inner_def[:args], gensym())
             pushfirst!(pass_args, cstr_type)
+            pushfirst!(pass_arg_types, :(Type{cstr_type}))
             pushfirst!(key_args, cstr_type)
             pushfirst!(key_arg_types, :(Type{cstr_type}))
         elseif @capture(def[:name], obj_::obj_type_ | ::obj_type_)
@@ -115,6 +114,7 @@ macro memoize(args...)
             sig = :(Tuple{$obj_type, $(arg_sigs...)} where {$(def[:whereparams]...)})
             pushfirst!(inner_def[:args], :($obj::$obj_type))
             pushfirst!(pass_args, obj)
+            pushfirst!(pass_arg_types, obj_type)
         else
             sig = :(Tuple{typeof($(def[:name])), $(arg_sigs...)} where {$(def[:whereparams]...)})
         end
@@ -131,6 +131,9 @@ macro memoize(args...)
         end
     end
 
+    # A return type declaration of Any is a No-op because everything is <: Any
+    return_type = get(def, :rtype, Any)
+
     if length(kwargs) == 0
         def[:body] = quote
             $(def[:body])::Core.Compiler.return_type($inner, typeof(($(pass_args...),)))

From 9f77ba806eb54c81936e8a3c6d716a569a37d0e2 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 22:16:36 -0500
Subject: [PATCH 12/26] exporting just the one function.

---
 src/Memoize.jl   | 28 +++++++++++++++-------------
 test/runtests.jl |  6 +++---
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 7a0de32..9ca2f73 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -1,6 +1,6 @@
 module Memoize
 using MacroTools: isexpr, combinearg, combinedef, namify, splitarg, splitdef, @capture
-export @memoize, function_memories, method_memories
+export @memoize, memories
 
 # I would call which($sig) but it's only on 1.6 I think
 function _which(tt, world = typemax(UInt))
@@ -15,7 +15,8 @@ function _which(tt, world = typemax(UInt))
     end
 end
 
-const _memories = Dict()
+const _brain = Dict()
+brain() = _brain
 
 macro memoize(args...)
     if length(args) == 1
@@ -155,27 +156,32 @@ macro memoize(args...)
         # If overwriting a method, empty the old cache.
         $old_meth = $_which($sig, $world)
         if $old_meth !== nothing
-            empty!(pop!($_memories, $old_meth, []))
+            empty!(pop!($brain(), $old_meth, []))
         end
 
         # Store the cache so that it can be emptied later
         $meth = $_which($sig)
         @assert $meth !== nothing
-        $_memories[$meth] = $cache
+        $brain()[$meth] = $cache
         $result
     end)
     #println(res)
     res
 end
 
-function_memories(f) = _function_memories(methods(f))
-function_memories(f, types) = _function_memories(methods(f, types))
-function_memories(f, types, mod) = _function_memories(methods(f, types, mod))
+"""
+    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 _function_memories(ms)
+function _memories(ms::Base.MethodList)
     memories = []
     for m in ms
-        memory = method_memory(m)
+        memory = get(brain(), m, nothing)
         if memory !== nothing
             push!(memories, memory)
         end
@@ -183,8 +189,4 @@ function _function_memories(ms)
     return memories
 end
 
-function method_memory(m::Method)
-    return get(_memories, m, nothing)
-end
-
 end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 270c198..11ed870 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -29,7 +29,7 @@ end
 @test simple(6) == 6
 @test run == 2
 
-map(empty!, function_memories(simple))
+map(empty!, memories(simple))
 @test simple(6) == 6
 @test run == 3
 @test simple(6) == 6
@@ -407,13 +407,13 @@ end # module
 using .MemoizeTest
 using .MemoizeTest: custom_dict
 
-map(empty!, function_memories(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
 
-map(empty!, function_memories(MemoizeTest.custom_dict))
+map(empty!, memories(MemoizeTest.custom_dict))
 @test custom_dict(1) == 1
 @test MemoizeTest.run == 4
 

From 88ef8f947fd5ed6c1d9690577a35609fb6193472 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 22:18:23 -0500
Subject: [PATCH 13/26] checking finalize

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

diff --git a/test/runtests.jl b/test/runtests.jl
index 11ed870..c891d9a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -346,6 +346,9 @@ finalized = false
     x
 end
 method_rewrite()
+@memoize function method_rewrite(x) end
+GC.gc()
+@test !finalized
 @memoize function method_rewrite() end
 GC.gc()
 @test finalized

From cfb705d96b2275b42746309fcbe5e01a8d812ca2 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 22:23:54 -0500
Subject: [PATCH 14/26] key value types

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

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 9ca2f73..6eb6618 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -147,20 +147,24 @@ macro memoize(args...)
 
     res = esc(quote
         # The `local` qualifier will make this performant even in the global scope.
-        local $cache = $cache_constructor
+        local $cache = begin
+            local __Key__ = (Tuple{$(key_arg_types...)} where {$(def[:whereparams]...)})
+            local __Val__ = ($return_type where {$(def[:whereparams]...)})
+            $cache_constructor
+        end
 
-        $world = Base.get_world_counter()
+        local $world = Base.get_world_counter()
 
-        $result = Base.@__doc__($(combinedef(def)))
+        local $result = Base.@__doc__($(combinedef(def)))
         
         # If overwriting a method, empty the old cache.
-        $old_meth = $_which($sig, $world)
+        local $old_meth = $_which($sig, $world)
         if $old_meth !== nothing
             empty!(pop!($brain(), $old_meth, []))
         end
 
         # Store the cache so that it can be emptied later
-        $meth = $_which($sig)
+        local $meth = $_which($sig)
         @assert $meth !== nothing
         $brain()[$meth] = $cache
         $result

From 895bb026bad7d9f186b5b3893c7d8165be56d44d Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 22:27:13 -0500
Subject: [PATCH 15/26] cleaning up some more

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

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 6eb6618..58d4263 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -145,7 +145,7 @@ macro memoize(args...)
     @gensym old_meth
     @gensym meth
 
-    res = esc(quote
+    esc(quote
         # The `local` qualifier will make this performant even in the global scope.
         local $cache = begin
             local __Key__ = (Tuple{$(key_arg_types...)} where {$(def[:whereparams]...)})
@@ -167,10 +167,9 @@ macro memoize(args...)
         local $meth = $_which($sig)
         @assert $meth !== nothing
         $brain()[$meth] = $cache
+
         $result
     end)
-    #println(res)
-    res
 end
 
 """

From b5f37fb3a95547a956e6dbf0f945f26cfbb3aa55 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 22:38:37 -0500
Subject: [PATCH 16/26] callable types.

---
 src/Memoize.jl   | 21 +++++++++++++--------
 test/runtests.jl | 21 +++++++++++++++++++++
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 58d4263..04e20a1 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -90,18 +90,21 @@ macro memoize(args...)
 
     @gensym result
 
-    # If this is a method of a callable object, the definition returns nothing.
+    # If this is a method of a callable type or object, the definition returns nothing.
     # Thus, we must construct the type of the method on our own.
+    # We also need to pass the object to the inner function
     if haskey(def, :name)
         if haskey(def, :params)
-            cstr_type = :($(def[:name]){$(def[:params]...)})
-            sig = :(Tuple{$cstr_type, $(arg_sigs...)} where {$(def[:whereparams]...)})
-            pushfirst!(inner_def[:args], gensym())
-            pushfirst!(pass_args, cstr_type)
-            pushfirst!(pass_arg_types, :(Type{cstr_type}))
-            pushfirst!(key_args, cstr_type)
-            pushfirst!(key_arg_types, :(Type{cstr_type}))
+            # Callable type
+            typ = :($(def[:name]){$(def[:params]...)})
+            sig = :(Tuple{Type{$typ}, $(arg_sigs...)} where {$(def[:whereparams]...)})
+            pushfirst!(inner_def[:args], :(::Type{$typ}))
+            pushfirst!(pass_args, typ)
+            pushfirst!(pass_arg_types, :(Type{$typ}))
+            pushfirst!(key_args, typ)
+            pushfirst!(key_arg_types, :(Type{$typ}))
         elseif @capture(def[:name], obj_::obj_type_ | ::obj_type_)
+            # Callable object
             obj_type === nothing && (obj_type = Any)
             if obj === nothing
                 obj = gensym()
@@ -117,9 +120,11 @@ macro memoize(args...)
             pushfirst!(pass_args, obj)
             pushfirst!(pass_arg_types, obj_type)
         else
+            # Normal call
             sig = :(Tuple{typeof($(def[:name])), $(arg_sigs...)} where {$(def[:whereparams]...)})
         end
     else
+        # Anonymous function
         sig = :(Tuple{typeof($result), $(arg_sigs...)} where {$(def[:whereparams]...)})
     end
 
diff --git a/test/runtests.jl b/test/runtests.jl
index c891d9a..5bf09b9 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -333,6 +333,27 @@ end
 @test callable_trait_object(2)(3) == (Int, 3)
 @test run == 4
 
+run = 0
+struct callable_type{T}
+    a::T
+end
+@memoize function callable_type{T}(b) where {T}
+    global run += 1
+    (T, b)
+end
+@test callable_type{Int}(2) == (Int, 2)
+@test run == 1
+@test callable_type{Int}(2) == (Int, 2)
+@test run == 1
+@test callable_type{Int}(3) == (Int, 3)
+@test run == 2
+@test callable_type{Int}(3) == (Int, 3)
+@test run == 2
+@test callable_type{Bool}(3) == (Bool, 3)
+@test run == 3
+@test callable_type{Bool}(3) == (Bool, 3)
+@test run == 3
+
 @memoize function typeinf(x)
     x + 1
 end

From 9b0bccd36cf237d1b186493fb649b27b33e3af03 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 22:50:58 -0500
Subject: [PATCH 17/26] which just returns the most specific previous method.

---
 src/Memoize.jl   |  2 +-
 test/runtests.jl | 31 +++++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 04e20a1..d632444 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -164,7 +164,7 @@ macro memoize(args...)
         
         # If overwriting a method, empty the old cache.
         local $old_meth = $_which($sig, $world)
-        if $old_meth !== nothing
+        if $old_meth !== nothing && $old_meth.sig == $sig
             empty!(pop!($brain(), $old_meth, []))
         end
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 5bf09b9..c3086e9 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -354,6 +354,37 @@ end
 @test callable_type{Bool}(3) == (Bool, 3)
 @test run == 3
 
+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

From 29cb4d8e3fc5e8f8abb4be7fed4559397f93d5d9 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 23:02:00 -0500
Subject: [PATCH 18/26] auto key val

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

diff --git a/test/runtests.jl b/test/runtests.jl
index c3086e9..322dfff 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -484,7 +484,21 @@ end
 @test dict_call("bb") == 2
 @test run == 2
 @test dict_call("bb") == 2
+
+run = 0
+@memoize Dict{__Key__,__Val__}() function auto_dict_call(a::String)::Int
+    global run += 1
+    length(a)
+end
+@test auto_dict_call("a") == 1
+@test run == 1
+@test auto_dict_call("a") == 1
+@test run == 1
+@test auto_dict_call("bb") == 2
+@test run == 2
+@test auto_dict_call("bb") == 2
 @test run == 2
+@test memories(auto_dict_call)[1] isa Dict{Tuple{String}, Int}
 
 @memoize non_allocating(x) = x+1
 @test @allocated(non_allocating(10)) == 0

From 5d0263c7d5898a22b9fc162a3d0405d29a80fc04 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 23:24:36 -0500
Subject: [PATCH 19/26] not sure if we should automagically type the
 dictionaries

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

diff --git a/src/Memoize.jl b/src/Memoize.jl
index d632444..df314a2 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -18,9 +18,17 @@ end
 const _brain = Dict()
 brain() = _brain
 
+"""
+    @memoize [cache] declaration
+    
+    Turns the function declaration into 
+    Return an array of memoized method caches for the function f.
+    
+    This function takes the same arguments as the method methods.
+"""
 macro memoize(args...)
     if length(args) == 1
-        cache_constructor = :(IdDict())
+        cache_constructor = :(IdDict{__Key__}{__Val__}())
         ex = args[1]
     elseif length(args) == 2
         (cache_constructor, ex) = args
@@ -44,7 +52,10 @@ macro memoize(args...)
         if arg_name === nothing
             arg_name = gensym()
             push!(key_args, arg_type)
-            push!(key_arg_types, :(Type{$arg_type}))
+            push!(key_arg_types, :(DataType))
+        elseif namify(arg_type) === :Vararg
+            push!(key_args, arg_name)
+            push!(key_arg_types, :(Tuple{$arg_type}))
         else
             push!(key_args, arg_name)
             push!(key_arg_types, arg_type)
@@ -90,6 +101,8 @@ macro memoize(args...)
 
     @gensym result
 
+    println(key_arg_types)
+
     # If this is a method of a callable type or object, the definition returns nothing.
     # Thus, we must construct the type of the method on our own.
     # We also need to pass the object to the inner function
@@ -102,14 +115,14 @@ macro memoize(args...)
             pushfirst!(pass_args, typ)
             pushfirst!(pass_arg_types, :(Type{$typ}))
             pushfirst!(key_args, typ)
-            pushfirst!(key_arg_types, :(Type{$typ}))
+            pushfirst!(key_arg_types, :(DataType))
         elseif @capture(def[:name], obj_::obj_type_ | ::obj_type_)
             # Callable object
             obj_type === nothing && (obj_type = Any)
             if obj === nothing
                 obj = gensym()
                 pushfirst!(key_args, obj_type)
-                pushfirst!(key_arg_types, :(Type{$obj_type}))
+                pushfirst!(key_arg_types, :(DataType))
             else
                 pushfirst!(key_args, obj)
                 pushfirst!(key_arg_types, obj_type)
@@ -150,7 +163,7 @@ macro memoize(args...)
     @gensym old_meth
     @gensym meth
 
-    esc(quote
+    res = esc(quote
         # The `local` qualifier will make this performant even in the global scope.
         local $cache = begin
             local __Key__ = (Tuple{$(key_arg_types...)} where {$(def[:whereparams]...)})
@@ -175,6 +188,7 @@ macro memoize(args...)
 
         $result
     end)
+    return res
 end
 
 """

From aafd53c0e473144cfc972e8c23bb78e34e182190 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Mon, 4 Jan 2021 23:38:29 -0500
Subject: [PATCH 20/26] documenting.

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

diff --git a/src/Memoize.jl b/src/Memoize.jl
index df314a2..002934c 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -1,6 +1,6 @@
 module Memoize
 using MacroTools: isexpr, combinearg, combinedef, namify, splitarg, splitdef, @capture
-export @memoize, memories
+export @memoize, memories, memory
 
 # I would call which($sig) but it's only on 1.6 I think
 function _which(tt, world = typemax(UInt))
@@ -21,10 +21,11 @@ brain() = _brain
 """
     @memoize [cache] declaration
     
-    Turns the function declaration into 
-    Return an array of memoized method caches for the function f.
+    Transform any method declaration `declaration` (except for inner constructors) so that calls to the original method are cached by their arguments. When an argument is unnamed, its type is treated as an argument instead.
     
-    This function takes the same arguments as the method methods.
+    `cache` should be an expression which evaluates to a dictionary-like type that supports `get!` and `empty!`, and may depend on the local variables `__Key__` and `__Value__`, which evaluate to syntactically-determined bounds on the required key and value types the cache must support.
+
+    If the given cache contains values, it is assumed that they will agree with the values the method returns. Specializing a method will not empty the cache, but overwriting a method will. The caches corresponding to methods can be determined with `memory` or `memories.`
 """
 macro memoize(args...)
     if length(args) == 1
@@ -211,4 +212,13 @@ function _memories(ms::Base.MethodList)
     return memories
 end
 
+"""
+    memory(m)
+    
+    Return the memoized cache for the method m, or nothing if no such method exists
+"""
+function memory(m::Method)
+    return get(brain(), m, nothing)
+end
+
 end
\ No newline at end of file

From 7a6f2d40fd8d15b06d2201a0935bb2bf90b1a3b5 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Tue, 5 Jan 2021 12:51:31 -0500
Subject: [PATCH 21/26] cleaning up

---
 src/Memoize.jl | 142 ++++++++++++++++++++-----------------------------
 1 file changed, 57 insertions(+), 85 deletions(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 002934c..23527d2 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -42,111 +42,82 @@ macro memoize(args...)
     catch
         error("@memoize must be applied to a method definition")
     end
-
-    # Set up arguments for memo key
-    key_args = []
-    key_arg_types = []
-
-    # Ensure that all args have names that can be passed to the inner function
-    function tag_arg(arg)
+    
+    function split(arg, iskwarg=false)
         arg_name, arg_type, slurp, default = splitarg(arg)
-        if arg_name === nothing
-            arg_name = gensym()
-            push!(key_args, arg_type)
-            push!(key_arg_types, :(DataType))
-        elseif namify(arg_type) === :Vararg
-            push!(key_args, arg_name)
-            push!(key_arg_types, :(Tuple{$arg_type}))
-        else
-            push!(key_args, arg_name)
-            push!(key_arg_types, arg_type)
-        end
-        return combinearg(arg_name, arg_type, slurp, default)
+        trait = arg_name === nothing
+        trait && (arg_name = gensym())
+        vararg = namify(arg_type) === :Vararg
+        return (
+            arg_name = arg_name,
+            arg_type = arg_type,
+            arg_value = arg_name,
+            slurp = slurp,
+            vararg = vararg,
+            default = default,
+            trait = trait,
+            iskwarg = iskwarg)
     end
-    args = def[:args] = map(tag_arg, def[:args])
-    kwargs = def[:kwargs] = map(tag_arg, def[:kwargs])
 
-    # Get argument types for function signature
-    arg_sigs = Vector{Any}(map(def[:args]) do arg
-        arg_name, arg_type, slurp, default = splitarg(arg)
-        if slurp
-            return :(Vararg{$arg_type})
-        else
-            return arg_type
-        end
-    end)
+    combine(arg) = combinearg(arg.arg_name, arg.arg_type, arg.slurp, arg.default)
 
-    # Set up identity arguments to pass to unmemoized function
-    pass_args = Vector{Any}(map(args) do arg
-        arg_name, arg_type, slurp, default = splitarg(arg)
-        if slurp || namify(arg_type) === :Vararg
-            Expr(:..., arg_name)
-        else
-            arg_name
-        end
-    end)
-    pass_arg_types = copy(arg_sigs)
-    pass_kwargs = Vector{Any}(map(kwargs) do kwarg
-        kwarg_name, kwarg_type, slurp, default = splitarg(kwarg)
-        if slurp
-            Expr(:..., kwarg_name)
-        else
-            Expr(:kw, kwarg_name, kwarg_name)
-        end
-    end)
+    pass(arg) =
+        (arg.slurp || arg.vararg) ? Expr(:..., arg.arg_name) :
+            arg.iskwarg ? Expr(:kw, arg.arg_name, arg.arg_name) : arg.arg_name
 
+    dispatch(arg) = arg.slurp ? :(Vararg{$(arg.arg_type)}) : arg.arg_type
+
+    args = split.(def[:args])
+    kwargs = split.(def[:kwargs], true)
+    def[:args] = combine.(args)
+    def[:kwargs] = combine.(kwargs)
     @gensym inner
     inner_def = deepcopy(def)
     inner_def[:name] = inner
+    inner_args = copy(args)
+    inner_kwargs = copy(kwargs)
     pop!(inner_def, :params, nothing)
 
     @gensym result
 
-    println(key_arg_types)
-
     # If this is a method of a callable type or object, the definition returns nothing.
     # Thus, we must construct the type of the method on our own.
     # We also need to pass the object to the inner function
     if haskey(def, :name)
-        if haskey(def, :params)
-            # Callable type
-            typ = :($(def[:name]){$(def[:params]...)})
-            sig = :(Tuple{Type{$typ}, $(arg_sigs...)} where {$(def[:whereparams]...)})
-            pushfirst!(inner_def[:args], :(::Type{$typ}))
-            pushfirst!(pass_args, typ)
-            pushfirst!(pass_arg_types, :(Type{$typ}))
-            pushfirst!(key_args, typ)
-            pushfirst!(key_arg_types, :(DataType))
-        elseif @capture(def[:name], obj_::obj_type_ | ::obj_type_)
-            # Callable object
-            obj_type === nothing && (obj_type = Any)
-            if obj === nothing
-                obj = gensym()
-                pushfirst!(key_args, obj_type)
-                pushfirst!(key_arg_types, :(DataType))
-            else
-                pushfirst!(key_args, obj)
-                pushfirst!(key_arg_types, obj_type)
-            end
-            def[:name] = :($obj::$obj_type)
-            sig = :(Tuple{$obj_type, $(arg_sigs...)} where {$(def[:whereparams]...)})
-            pushfirst!(inner_def[:args], :($obj::$obj_type))
-            pushfirst!(pass_args, obj)
-            pushfirst!(pass_arg_types, obj_type)
-        else
-            # Normal call
-            sig = :(Tuple{typeof($(def[:name])), $(arg_sigs...)} where {$(def[:whereparams]...)})
+        if haskey(def, :params) # Callable type
+            typ = :($(def[:name]){$(pop!(def, :params)...)})
+            inner_args = [split(:(::Type{$typ})), inner_args...]
+            def[:name] = combine(inner_args[1])
+            head = :(Type{$typ})
+        elseif @capture(def[:name], obj_::obj_type_ | ::obj_type_) # Callable object
+            inner_args = [split(def[:name]), inner_args...]
+            def[:name] = combine(inner_args[1])
+            head = obj_type
+        else # Normal call
+            head = :(typeof($(def[:name])))
         end
-    else
-        # Anonymous function
-        sig = :(Tuple{typeof($result), $(arg_sigs...)} where {$(def[:whereparams]...)})
+    else # Anonymous function
+        head = :(typeof($result))
+    end
+    inner_def[:args] = combine.(inner_args)
+
+    # Set up arguments for memo key
+    key_names = map([inner_args; inner_kwargs]) do arg
+        arg.trait ? arg.arg_type : arg.arg_name
+    end
+    key_types = map([inner_args; inner_kwargs]) do arg
+        arg.trait ? DataType :
+        arg.vararg ? :(Tuple{$(arg.arg_type)}) : #TODO arg.slurp?
+            arg.arg_type
     end
 
     @gensym cache
 
+    pass_args = pass.(split.(inner_def[:args]))
+    pass_kwargs = pass.(split.(inner_def[:kwargs], true))
     def[:body] = quote
         $(combinedef(inner_def))
-        get!($cache, ($(key_args...),)) do
+        get!($cache, ($(key_names...),)) do
             $inner($(pass_args...); $(pass_kwargs...))
         end
     end
@@ -164,10 +135,12 @@ macro memoize(args...)
     @gensym old_meth
     @gensym meth
 
-    res = esc(quote
+    sig = :(Tuple{$head, $(dispatch.(args)...)} where {$(def[:whereparams]...)})
+
+    return esc(quote
         # The `local` qualifier will make this performant even in the global scope.
         local $cache = begin
-            local __Key__ = (Tuple{$(key_arg_types...)} where {$(def[:whereparams]...)})
+            local __Key__ = (Tuple{$(key_types...)} where {$(def[:whereparams]...)})
             local __Val__ = ($return_type where {$(def[:whereparams]...)})
             $cache_constructor
         end
@@ -189,7 +162,6 @@ macro memoize(args...)
 
         $result
     end)
-    return res
 end
 
 """

From 655db4a742761d1ca6fb97ed8f9b187fba876f0f Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Tue, 5 Jan 2021 12:53:27 -0500
Subject: [PATCH 22/26] finishing touches

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

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 23527d2..2f35183 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -107,14 +107,14 @@ macro memoize(args...)
     end
     key_types = map([inner_args; inner_kwargs]) do arg
         arg.trait ? DataType :
-        arg.vararg ? :(Tuple{$(arg.arg_type)}) : #TODO arg.slurp?
+        arg.vararg ? :(Tuple{$(arg.arg_type)}) :
             arg.arg_type
     end
 
     @gensym cache
 
-    pass_args = pass.(split.(inner_def[:args]))
-    pass_kwargs = pass.(split.(inner_def[:kwargs], true))
+    pass_args = pass.(inner_args)
+    pass_kwargs = pass.(inner_kwargs)
     def[:body] = quote
         $(combinedef(inner_def))
         get!($cache, ($(key_names...),)) do

From 0d8c9d5bc50b6900eb68a71cfba8058b386db1a8 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Tue, 5 Jan 2021 13:31:20 -0500
Subject: [PATCH 23/26] update README

---
 README.md        | 55 ++++++++++++++++++++++++++++++++++--------------
 src/Memoize.jl   |  2 +-
 test/runtests.jl | 20 ++++++++++++++++++
 3 files changed, 60 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 2876a31..e6b4685 100644
--- a/README.md
+++ b/README.md
@@ -6,7 +6,7 @@
 [ci-img]: https://github.com/JuliaCollections/Memoize.jl/workflows/CI/badge.svg
 [ci-url]: https://github.com/JuliaCollections/Memoize.jl/actions
 
-Easy memoization for Julia.
+Easy method memoization for Julia.
 
 ## Usage
 
@@ -23,15 +23,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{Tuple{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
@@ -41,22 +42,22 @@ julia> x(1)
 2
 ```
 
-By default, Memoize.jl uses an [`IdDict`](https://docs.julialang.org/en/v1/base/collections/#Base.IdDict) as a cache, but it's also possible to specify the type of the cache. If you want to cache vectors based on the values they contain, you probably want this:
+By default, Memoize.jl uses an [`IdDict`](https://docs.julialang.org/en/v1/base/collections/#Base.IdDict) as a cache, but it's also possible to specify your own cache that supports the methods `Base.get!` and `Base.empty!`. If you want to cache vectors based on the values they contain, you probably want this:
 
 ```julia
 using Memoize
-@memoize Dict function x(a)
+@memoize Dict() function x(a)
 	println("Running")
 	a
 end
 ```
 
-You can also specify the full function call for constructing the dictionary. For example, to use LRUCache.jl:
+You can also specify the full expression for constructing the cache. The variables `__Key__` and `__Val__` are available to the constructor expression, containing the syntactically determined type bounds on the keys and values used by Memoize.jl.  For example, to use LRUCache.jl:
 
 ```julia
 using Memoize
 using LRUCache
-@memoize LRU{Tuple{Any,Any},Any}(maxsize=2) function x(a, b)
+@memoize LRU{__Key__,__Val__}(maxsize=2) function x(a, b)
     println("Running")
     a + b
 end
@@ -86,12 +87,34 @@ julia> x(2,3)
 5
 ```
 
-## Notes
+Memoize works on *almost* every method declaration in global and local scope, including lambdas and callable objects. When only the type of an argument is given, memoize caches the type.
 
-Note that the `@memoize` macro treats the type argument differently depending on its syntactical form: in the expression
-```julia
-@memoize CacheType function x(a, b)
-    # ...
+julia```
+struct F{A}
+	a::A
+end
+@memoize function (::F{A})(b, ::C) where {A, C}
+	println("Running")
+	(A, b, C)
 end
 ```
-the expression `CacheType` must be either a non-function-call that evaluates to a type, or a function call that evaluates to an _instance_ of the desired cache type.  Either way, the methods `Base.get!` and `Base.empty!` must be defined for the supplied cache type.
+
+```
+julia> F(1)(1, 1)
+Running
+(Int64, 1, Int64)
+
+julia> F(1)(1, 2)
+(Int64, 1, Int64)
+
+julia> F(1)(2, 2)
+Running
+(Int64, 2, Int64)
+
+julia> F(2)(2, 2)
+(Int64, 2, Int64)
+
+julia> F(false)(2, 2)
+Running
+(Bool, 2, Int64)
+```
diff --git a/src/Memoize.jl b/src/Memoize.jl
index 2f35183..1668315 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -127,7 +127,7 @@ macro memoize(args...)
 
     if length(kwargs) == 0
         def[:body] = quote
-            $(def[:body])::Core.Compiler.return_type($inner, typeof(($(pass_args...),)))
+            $(def[:body])::Core.Compiler.widenconst(Core.Compiler.return_type($inner, typeof(($(pass_args...),))))
         end
     end
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 322dfff..e304388 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -35,6 +35,26 @@ map(empty!, memories(simple))
 @test simple(6) == 6
 @test run == 3
 
+run = 0
+lambda = @memoize (a) -> begin
+    global run += 1
+    a
+end
+@test lambda(5) == 5
+@test run == 1
+@test lambda(5) == 5
+@test run == 1
+@test lambda(6) == 6
+@test run == 2
+@test lambda(6) == 6
+@test run == 2
+
+map(empty!, memories(lambda))
+@test lambda(6) == 6
+@test run == 3
+@test lambda(6) == 6
+@test run == 3
+
 run = 0
 @memoize function typed(a::Int)
     global run += 1

From 426b97d26ebfda3e7e28e69f1a0b0f42df42eca6 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Tue, 5 Jan 2021 14:23:11 -0500
Subject: [PATCH 24/26] spell out __Value__

---
 README.md        | 4 ++--
 src/Memoize.jl   | 4 ++--
 test/runtests.jl | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index e6b4685..7e4dc0d 100644
--- a/README.md
+++ b/README.md
@@ -52,12 +52,12 @@ using Memoize
 end
 ```
 
-You can also specify the full expression for constructing the cache. The variables `__Key__` and `__Val__` are available to the constructor expression, containing the syntactically determined type bounds on the keys and values used by Memoize.jl.  For example, to use LRUCache.jl:
+You can also specify the full expression for constructing the cache. The variables `__Key__` and `__Value__` are available to the constructor expression, containing the syntactically determined type bounds on the keys and values used by Memoize.jl.  For example, to use LRUCache.jl:
 
 ```julia
 using Memoize
 using LRUCache
-@memoize LRU{__Key__,__Val__}(maxsize=2) function x(a, b)
+@memoize LRU{__Key__,__Value__}(maxsize=2) function x(a, b)
     println("Running")
     a + b
 end
diff --git a/src/Memoize.jl b/src/Memoize.jl
index 1668315..cc9d828 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -29,7 +29,7 @@ brain() = _brain
 """
 macro memoize(args...)
     if length(args) == 1
-        cache_constructor = :(IdDict{__Key__}{__Val__}())
+        cache_constructor = :(IdDict{__Key__}{__Value__}())
         ex = args[1]
     elseif length(args) == 2
         (cache_constructor, ex) = args
@@ -141,7 +141,7 @@ macro memoize(args...)
         # The `local` qualifier will make this performant even in the global scope.
         local $cache = begin
             local __Key__ = (Tuple{$(key_types...)} where {$(def[:whereparams]...)})
-            local __Val__ = ($return_type where {$(def[:whereparams]...)})
+            local __Value__ = ($return_type where {$(def[:whereparams]...)})
             $cache_constructor
         end
 
diff --git a/test/runtests.jl b/test/runtests.jl
index e304388..2d2c5cf 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -506,7 +506,7 @@ end
 @test dict_call("bb") == 2
 
 run = 0
-@memoize Dict{__Key__,__Val__}() function auto_dict_call(a::String)::Int
+@memoize Dict{__Key__,__Value__}() function auto_dict_call(a::String)::Int
     global run += 1
     length(a)
 end

From 6087295e9b13f040aa3f6b6b5f5759a222389dc2 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Tue, 5 Jan 2021 15:48:43 -0500
Subject: [PATCH 25/26] precompilation safe, I think.

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

diff --git a/src/Memoize.jl b/src/Memoize.jl
index cc9d828..94738ad 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -15,9 +15,6 @@ function _which(tt, world = typemax(UInt))
     end
 end
 
-const _brain = Dict()
-brain() = _brain
-
 """
     @memoize [cache] declaration
     
@@ -134,6 +131,7 @@ macro memoize(args...)
     @gensym world
     @gensym old_meth
     @gensym meth
+    @gensym brain
 
     sig = :(Tuple{$head, $(dispatch.(args)...)} where {$(def[:whereparams]...)})
 
@@ -148,17 +146,24 @@ macro memoize(args...)
         local $world = Base.get_world_counter()
 
         local $result = Base.@__doc__($(combinedef(def)))
+
+        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
-            empty!(pop!($brain(), $old_meth, []))
+            empty!(pop!($brain, $old_meth.sig, []))
         end
 
         # Store the cache so that it can be emptied later
         local $meth = $_which($sig)
         @assert $meth !== nothing
-        $brain()[$meth] = $cache
+        $brain[$meth.sig] = $cache
 
         $result
     end)
@@ -176,9 +181,12 @@ memories(f, args...) = _memories(methods(f, args...))
 function _memories(ms::Base.MethodList)
     memories = []
     for m in ms
-        memory = get(brain(), m, nothing)
-        if memory !== nothing
-            push!(memories, memory)
+        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
@@ -190,7 +198,10 @@ end
     Return the memoized cache for the method m, or nothing if no such method exists
 """
 function memory(m::Method)
-    return get(brain(), m, nothing)
+    if isdefined(m.module, :__Memoize_brain__)
+        brain = getfield(m.module, :__Memoize_brain__)
+        return get(brain, m.sig, nothing)
+    end
 end
 
 end
\ No newline at end of file

From f4e697ea11b0fd488a7579bf8c355b279140f314 Mon Sep 17 00:00:00 2001
From: Peter Ahrens <ptrahrens@gmail.com>
Date: Tue, 5 Jan 2021 16:16:10 -0500
Subject: [PATCH 26/26] more precompilation fixes

---
 src/Memoize.jl | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/Memoize.jl b/src/Memoize.jl
index 94738ad..0ba5e0d 100644
--- a/src/Memoize.jl
+++ b/src/Memoize.jl
@@ -132,6 +132,7 @@ macro memoize(args...)
     @gensym old_meth
     @gensym meth
     @gensym brain
+    @gensym old_brain
 
     sig = :(Tuple{$head, $(dispatch.(args)...)} where {$(def[:whereparams]...)})
 
@@ -157,7 +158,10 @@ macro memoize(args...)
         # Notice that methods are hashed by their stored signature
         local $old_meth = $_which($sig, $world)
         if $old_meth !== nothing && $old_meth.sig == $sig
-            empty!(pop!($brain, $old_meth.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