From 77d2eaca73c03dcff310b8987e857c1db0f5aab8 Mon Sep 17 00:00:00 2001
From: Mauro Werder <mauro_lc@runbox.com>
Date: Mon, 19 Feb 2018 10:32:09 +0100
Subject: [PATCH] Fixed issue with dispatch_cache and modules

---
 src/SimpleTraits.jl | 10 +++++++---
 test/runtests.jl    | 11 +++++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/SimpleTraits.jl b/src/SimpleTraits.jl
index 2d09b55..4c06bd0 100644
--- a/src/SimpleTraits.jl
+++ b/src/SimpleTraits.jl
@@ -187,7 +187,7 @@ end
 dispatch_cache = Dict()  # to ensure that the trait-dispatch function is defined only once per pair
 let
     global traitfn
-    function traitfn(tfn)
+    function traitfn(tfn, cur_module)
         # Need
         # f{X,Y}(x::X,Y::Y) = f(trait(Tr1{X,Y}), x, y)
         # f(::False, x, y)= ...
@@ -300,7 +300,7 @@ let
         end
         # Create the trait dispatch function
         ex = fn
-        key = (@__MODULE__, fname, typs0, strip_kw(args0), trait0_opposite)
+        key = (cur_module, fname, typs0, strip_kw(args0), trait0_opposite)
         if !(key ∈ keys(dispatch_cache)) # define trait dispatch function
             if !haskwargs
                 ex = quote
@@ -354,7 +354,11 @@ Defines a function dispatching on a trait. Examples:
 Note that the second example is just syntax sugar for `@traitfn f{X,Y; Not{Tr1{X,Y}}}(x::X,y::Y) = ...`.
 """
 macro traitfn(tfn)
-    esc(traitfn(tfn))
+    @static if isdefined(Base, Symbol("@__MODULE__"))
+        esc(traitfn(tfn, __module__))
+    else
+        esc(traitfn(tfn, current_module()))
+    end
 end
 
 ######
diff --git a/test/runtests.jl b/test/runtests.jl
index 4939b35..5f885cf 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -245,6 +245,17 @@ isarrow(X) = eltype(X)<:Integer ? true : false
 @test !istrait(TrArrow2{Vector{Int}})
 @test istrait(TrArrow2{Vector{Float64}})
 
+####
+# issue with key in dispatch_cache
+@traitfn f_dc(::::Tr1) = 1
+module Mod
+using SimpleTraits
+using Base.Test
+@traitdef Tr1{X}
+@traitimpl Tr1{Integer}
+@traitfn f_dc(::::Tr1) = 2
+@test f_dc(1) == 2
+end
 
 ######
 # Other tests