Record wrapped functions instead of keyword wrappers when possible

Previously, if `f(; x)` called g(x), then the `current_function` would
be something like `getfield(Main, Symbol("#kw##f"))` because of the
way that keyword functions are implemented. Now, it is simply `f` as
expected.
The only time that this seems to not work is when `f` is a closure.

src/SimpleMock.jl

@@ -5,7 +5,7 @@ module SimpleMock
 
 using Base: Callable, invokelatest, unwrap_unionall
 using Base.Iterators: Pairs
-using Core: Builtin, IntrinsicFunction, kwftype
+using Core: Builtin, kwftype
 
 using Cassette: Cassette, Context, overdub, posthook, prehook, recurse, @context
 

src/metadata.jl

@@ -55,17 +55,42 @@ should_mock(m::Metadata, method::Tuple) = method in m.methods && all(f -> f(m),
 
 # Update the function/module stacks.
 function update!(m::Metadata, ::typeof(prehook), @nospecialize(f), @nospecialize(args...))
+    f_uw = unwrap_fun(f)
     Ts = Tuple{map(typeof, args)...}
-    mod = if f isa Union{Builtin, IntrinsicFunction} || !hasmethod(f, Ts)
-        parentmodule(f)
+
+    mod = if f_uw isa Builtin || !hasmethod(f_uw, Ts)
+        parentmodule(f_uw)
     else
-        parentmodule(f, Ts)
+        parentmodule(f_uw, Ts)
     end
-    push!(m.funcs, f)
+
+    push!(m.funcs, f_uw)
     push!(m.mods, mod)
 end
 
 function update!(m::Metadata, ::typeof(posthook), @nospecialize(args...))
     pop!(m.funcs)
     pop!(m.mods)
 end
+
+# If a function is a keyword wrapper, try to get the wrapped function.
+# This garbage is the result of random experimentation and is very sketchy.
+# It fails in the case of closures.
+unwrap_fun(f) = f
+unwrap_fun(f::Builtin) = f
+function unwrap_fun(f::F) where F <: Function
+    startswith(string(F.name.name), "#kw##") || return f
+
+    name = Symbol(string(F.name.name)[6:end])
+    name_hash = Symbol("#", name)
+    mod = F.name.module
+
+    return if isdefined(mod, name)
+        getfield(mod, name)
+    elseif isdefined(mod, name_hash)
+        gf = getfield(mod, name_hash)
+        isdefined(gf, :instance) ? gf.instance : f
+    else
+        f
+    end
+end

test/mock_fun.jl

@@ -177,4 +177,32 @@ end
             @test ncalls(f) == 2 && has_calls(f, Call(), Call())
         end
     end
+
+    @testset "Metadata records original functions and not keyword wrappers" begin
+        @eval begin
+            kw_f(; x) = kw_g(; x=x)
+            kw_g(; x) = 2x
+            kw_h = (; x) -> kw_i(; x=x)
+            kw_i = (; x) -> 2x
+        end
+        kw_j(; x) = kw_k(; x=x)
+        kw_k(; x) = 2x
+
+        test(f; broken::Bool=false) = function(m::SimpleMock.Metadata)
+            eq = SimpleMock.current_function(m) === f
+            if broken
+                @test_broken eq
+            else
+                @test eq
+            end
+            return true
+        end
+
+        # Regular functions.
+        @test mock(_g -> kw_f(; x=1), kw_g; filters=[test(kw_f)]) != 2
+        # Anonymous functions.
+        @test mock(_i -> kw_h(; x=1), kw_i; filters=[test(kw_h)]) != 2
+        # Closures.
+        @test mock(_k -> kw_j(; x=1), kw_k; filters=[test(kw_j; broken=true)]) != 2
+    end
 end

test/runtests.jl

@@ -1,6 +1,6 @@
 using Base: JLOptions
 
-using Test: @test, @testset, @test_logs, @test_throws
+using Test: @test, @testset, @test_broken, @test_logs, @test_throws
 
 using Suppressor: @capture_err, @suppress