Improve typeCheck?

1 files changed, 22 insertions, 12 deletions

diff --git a/src/Data/Dependent/EnumMap/Strict/Internal.hs b/src/Data/Dependent/EnumMap/Strict/Internal.hs
index 79d2b52..b2bef1a 100644
--- a/src/Data/Dependent/EnumMap/Strict/Internal.hs
+++ b/src/Data/Dependent/EnumMap/Strict/Internal.hs
@@ -742,24 +742,34 @@ maxViewWithKey (DEnumMap m) = bimap' kVToDSum DEnumMap <$!> IM.maxViewWithKey m
 coe1 :: v a -> v b
 coe1 = unsafeCoerce
 
-typeCheck1 :: (Enum1 k, TestEquality k)
+typeCheck1 :: forall k a b. (Enum1 k, TestEquality k)
            => k a -> Int -> Enum1Info k b -> k a :~: k b
 typeCheck1 k1 i inf2 =
-  assert (let k2 = toEnum1 i inf2
-          in case testEquality k1 k2 of
-               Just Refl -> True
-               Nothing -> False)
-         (unsafeCoerce Refl)
+  let ret :: k a :~: k b
+      ret = unsafeCoerce Refl
+  in assert (let k2 = toEnum1 i inf2
+             in case testEquality k1 k2 of
+                  Just r -> justifies r ret True
+                  Nothing -> False)
+            ret
+  where
+    justifies :: a :~: b -> k a :~: k b -> r -> r
+    justifies Refl _ x = x
 
 typeCheck2 :: forall k proxy a b. (Enum1 k, TestEquality k)
            => proxy k -> Int -> Enum1Info k a -> Enum1Info k b -> k a :~: k b
 typeCheck2 _ i inf1 inf2 =
-  assert (let k1 = toEnum1 @k i inf1
-              k2 = toEnum1 i inf2
-          in case testEquality k1 k2 of
-               Just Refl -> True
-               Nothing -> False)
-         (unsafeCoerce Refl)
+  let ret :: k a :~: k b
+      ret = unsafeCoerce Refl
+  in assert (let k1 = toEnum1 @k i inf1
+                 k2 = toEnum1 i inf2
+             in case testEquality k1 k2 of
+                  Just r -> justifies r ret True
+                  Nothing -> False)
+            ret
+  where
+    justifies :: a :~: b -> k a :~: k b -> r -> r
+    justifies Refl _ x = x
 
 bimap' :: (a -> b) -> (c -> d) -> (a, c) -> (b, d)
 bimap' f g (a, c) =