From 74720f495c1420f46e21fb6aa7ea338d1fc97d0f Mon Sep 17 00:00:00 2001
From: KristofferC <kcarlsson89@gmail.com>
Date: Tue, 8 Sep 2020 17:19:18 +0200
Subject: [PATCH 1/2] fix some issues with module prefixed functions

---
 src/utils.jl     | 38 +++++++++++++++++++-------------------
 test/runtests.jl | 11 +++++++++++
 test/script.jl   | 12 ++++++++++++
 3 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/src/utils.jl b/src/utils.jl
index d88a64c..ae56ee1 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -1,27 +1,27 @@
-function isfuncexpr(ex, name=nothing)
-    checkname(fdef::Expr, name)            = checkname(fdef.args[1], name)
-    checkname(fname::Symbol, name::Symbol) = begin
-        fname === name && return true
-        startswith(string(name), string('#', fname, '#')) && return true
-        return false
+function checkname(fdef::Expr, name)
+    fdef.head === :where && return checkname(fdef.args[1], name)
+    if fdef.args[1] isa Expr
+        # E.g. `function Mod.bar.foo(a, b)`
+        return checkname(fdef.args[1].args[end], name)
     end
-    checkname(fname::Symbol, ::Nothing)    = true
+    return checkname(fdef.args[1], name)
+end
+checkname(fname::Symbol, name::Symbol) = begin
+    fname === name && return true
+    startswith(string(name), string('#', fname, '#')) && return true
+    return false
+end
+checkname(fname::Symbol, ::Nothing) = true
+checkname(fname::QuoteNode, name) = checkname(fname.value, name)
 
+function isfuncexpr(ex, name=nothing)
     # Strip any macros that wrap the method definition
-    while isexpr(ex, :macrocall) && length(ex.args) == 3
+    while ex isa Expr && ex.head === :macrocall && length(ex.args) == 3
         ex = ex.args[3]
     end
     isa(ex, Expr) || return false
-    ex.head == :function && return checkname(ex, name)
-    if ex.head == :(=)
-        a = ex.args[1]
-        if isa(a, Expr)
-            while a.head == :where
-                a = a.args[1]
-                isa(a, Expr) || return false
-            end
-            a.head == :call && return checkname(a, name)
-        end
+    if ex.head === :function || ex.head === :(=)
+        return checkname(ex.args[1], name)
     end
     return false
 end
@@ -35,7 +35,7 @@ end
 linerange(arg) = linerange(convert(Expr, arg))  # Handle Revise's RelocatableExpr
 
 function findline(ex, order)
-    ex.head == :line && return ex.args[1], true
+    ex.head === :line && return ex.args[1], true
     for a in order(ex.args)
         a isa LineNumberNode && return a.line, true
         if a isa Expr
diff --git a/test/runtests.jl b/test/runtests.jl
index 91de44c..1c69f48 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -190,3 +190,14 @@ end
     @test loc == 28
     @test body == "func_2nd_kwarg(; kw=2) = true"
 end
+
+@testset "method extensions" begin
+    body, _ = CodeTracking.definition(String, @which Foo.Bar.fit(1))
+    @test body == """
+    function Foo.Bar.fit(m)
+        return m
+    end"""
+    body, _ = CodeTracking.definition(String, @which Foo.Bar.fit(1, 2))
+    @show body
+    @test body == "Foo.Bar.fit(a, b) = a + b"
+end
diff --git a/test/script.jl b/test/script.jl
index af157bf..5562579 100644
--- a/test/script.jl
+++ b/test/script.jl
@@ -26,3 +26,15 @@ end
 
 func_1st_nokwarg() = true
 func_2nd_kwarg(; kw=2) = true
+
+module Foo
+module Bar
+function fit end
+end
+end
+
+function Foo.Bar.fit(m)
+    return m
+end
+
+Foo.Bar.fit(a, b) = a + b

From ea14df6534256fe45113b85d10dfa314c6b7577d Mon Sep 17 00:00:00 2001
From: Kristoffer <kcarlsson89@gmail.com>
Date: Wed, 9 Sep 2020 11:29:53 +0200
Subject: [PATCH 2/2] fixups

---
 src/utils.jl     | 11 +++++++----
 test/runtests.jl |  1 -
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/utils.jl b/src/utils.jl
index ae56ee1..7934a19 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -1,10 +1,13 @@
 function checkname(fdef::Expr, name)
-    fdef.head === :where && return checkname(fdef.args[1], name)
-    if fdef.args[1] isa Expr
+    fproto = fdef.args[1]
+    fdef.head === :where && return checkname(fproto, name)
+    if fproto isa Expr
+        # Is the check below redundant?
+        fproto.head === :. || return false
         # E.g. `function Mod.bar.foo(a, b)`
-        return checkname(fdef.args[1].args[end], name)
+        return checkname(fproto.args[end], name)
     end
-    return checkname(fdef.args[1], name)
+    return checkname(fproto, name)
 end
 checkname(fname::Symbol, name::Symbol) = begin
     fname === name && return true
diff --git a/test/runtests.jl b/test/runtests.jl
index 1c69f48..33e6441 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -198,6 +198,5 @@ end
         return m
     end"""
     body, _ = CodeTracking.definition(String, @which Foo.Bar.fit(1, 2))
-    @show body
     @test body == "Foo.Bar.fit(a, b) = a + b"
 end