From 59884f253f93fda7a71d64186b60e554818198f1 Mon Sep 17 00:00:00 2001
From: Michael Miller <icy.arctic.fox@gmail.com>
Date: Sun, 3 Nov 2019 13:04:44 -0700
Subject: [PATCH] Better handling of casting with covariance and contravariance

---
 src/spectator/mocks/double.cr              | 12 ++----------
 src/spectator/mocks/generic_method_stub.cr |  2 --
 src/spectator/mocks/method_stub.cr         |  2 ++
 src/spectator/mocks/nil_method_stub.cr     |  8 ++++++--
 src/spectator/mocks/proc_method_stub.cr    |  9 +++++++--
 src/spectator/mocks/value_method_stub.cr   |  9 +++++++--
 6 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/src/spectator/mocks/double.cr b/src/spectator/mocks/double.cr
index ccba68b..a0ca3f4 100644
--- a/src/spectator/mocks/double.cr
+++ b/src/spectator/mocks/double.cr
@@ -41,11 +41,7 @@ module Spectator::Mocks
         %call = ::Spectator::Mocks::GenericMethodCall.new({{name.symbolize}}, %args)
         @spectator_stub_calls << %call
         if (%stub = @spectator_stubs.find(&.callable?(%call)))
-          if (%cast = %stub.as?(::Spectator::Mocks::GenericMethodStub(typeof(%method({{args.splat}})))))
-            %cast.call(%args)
-          else
-            raise "The return type of stub #{%stub} doesn't match the expected type #{typeof(%method({{args.splat}}))}"
-          end
+          %stub.call(%args, typeof(%method({{args.splat}})))
         else
           %method({{args.splat}})
         end
@@ -56,11 +52,7 @@ module Spectator::Mocks
         %call = ::Spectator::Mocks::GenericMethodCall.new({{name.symbolize}}, %args)
         @spectator_stub_calls << %call
         if (%stub = @spectator_stubs.find(&.callable?(%call)))
-          if (%cast = %stub.as?(::Spectator::Mocks::GenericMethodStub(typeof(%method({{args.splat}}) { |*%yield_args| yield *%yield_args }))))
-            %cast.call(%args)
-          else
-            raise "The return type of stub #{%stub} doesn't match the expected type #{typeof(%method({{args.splat}}) { |*%ya| yield *%ya })}"
-          end
+          %stub.call(%args, typeof(%method({{args.splat}}) { |*%ya| yield *%ya }))
         else
           %method({{args.splat}}) do |*%yield_args|
             yield *%yield_args
diff --git a/src/spectator/mocks/generic_method_stub.cr b/src/spectator/mocks/generic_method_stub.cr
index f5e650e..0ac09ed 100644
--- a/src/spectator/mocks/generic_method_stub.cr
+++ b/src/spectator/mocks/generic_method_stub.cr
@@ -13,8 +13,6 @@ module Spectator::Mocks
       super && (!@args || @args === call.args)
     end
 
-    abstract def call(args : GenericArguments(T, NT)) : ReturnType forall T, NT
-
     def to_s(io)
       super(io)
       if @args
diff --git a/src/spectator/mocks/method_stub.cr b/src/spectator/mocks/method_stub.cr
index 98bc268..01020ac 100644
--- a/src/spectator/mocks/method_stub.cr
+++ b/src/spectator/mocks/method_stub.cr
@@ -10,6 +10,8 @@ module Spectator::Mocks
       @name == call.name
     end
 
+    abstract def call(args : GenericArguments(T, NT), rt : RT.class) forall T, NT, RT
+
     def to_s(io)
       io << '#'
       io << @name
diff --git a/src/spectator/mocks/nil_method_stub.cr b/src/spectator/mocks/nil_method_stub.cr
index 798d378..9a8954c 100644
--- a/src/spectator/mocks/nil_method_stub.cr
+++ b/src/spectator/mocks/nil_method_stub.cr
@@ -4,8 +4,12 @@ require "./value_method_stub"
 
 module Spectator::Mocks
   class NilMethodStub < GenericMethodStub(Nil)
-    def call(_args : GenericArguments(T, NT)) : Nil forall T, NT
-      nil
+    def call(_args : GenericArguments(T, NT), rt : RT.class) forall T, NT, RT
+      if (cast = nil.as?(RT))
+        cast
+      else
+        raise "The return type of stub #{to_s} : Nil doesn't match the expected type #{RT}"
+      end
     end
 
     def and_return(value : T) forall T
diff --git a/src/spectator/mocks/proc_method_stub.cr b/src/spectator/mocks/proc_method_stub.cr
index f3141ed..1960a04 100644
--- a/src/spectator/mocks/proc_method_stub.cr
+++ b/src/spectator/mocks/proc_method_stub.cr
@@ -7,8 +7,13 @@ module Spectator::Mocks
       super(name, source, args)
     end
 
-    def call(_args : GenericArguments(T, NT)) : ReturnType forall T, NT
-      @proc.call
+    def call(_args : GenericArguments(T, NT), rt : RT.class) forall T, NT, RT
+      result = @proc.call
+      if (cast = result.as?(RT))
+        cast
+      else
+        raise "The return type of stub #{to_s} : #{ReturnType} doesn't match the expected type #{RT}"
+      end
     end
   end
 end
diff --git a/src/spectator/mocks/value_method_stub.cr b/src/spectator/mocks/value_method_stub.cr
index fbf7421..f1e989a 100644
--- a/src/spectator/mocks/value_method_stub.cr
+++ b/src/spectator/mocks/value_method_stub.cr
@@ -7,8 +7,13 @@ module Spectator::Mocks
       super(name, source, args)
     end
 
-    def call(_args : GenericArguments(T, NT)) : ReturnType forall T, NT
-      @value
+    def call(_args : GenericArguments(T, NT), rt : RT.class) forall T, NT, RT
+      result = @value
+      if (cast = result.as?(RT))
+        cast
+      else
+        raise "The return type of stub #{to_s} : #{ReturnType} doesn't match the expected type #{RT}"
+      end
     end
   end
 end