Specializing: Cache a non-bignum version of module nat.typeable

1 files changed, 9 insertions, 5 deletions

diff --git a/packages/base/src/Internal/Specialized.hs b/packages/base/src/Internal/Specialized.hs
index c85d44f..6143f65 100644
--- a/packages/base/src/Internal/Specialized.hs
+++ b/packages/base/src/Internal/Specialized.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE ConstraintKinds #-}
@@ -116,10 +117,12 @@ idint = IntegralRep id id id id id Nothing
 coerceint :: Coercible t a => IntegralRep t a
 coerceint = IntegralRep coerce coerce coerce coerce coerce Nothing
 
+data CachedNat = forall (n::Nat). Typeable n => CachedNat !Int64 !(Proxy n)
+
 -- This exists to hopefully save me from parsing a string since Type.Reflection
 -- currently has no direct way to extract a Nat value from a TypeRep.
-cachedNat :: IORef SomeNat
-cachedNat = unsafePerformIO $ newIORef (SomeNat (Proxy :: Proxy 3))
+cachedNat :: IORef CachedNat
+cachedNat = unsafePerformIO $ newIORef (CachedNat 3 (Proxy::Proxy 3))
 {-# NOINLINE cachedNat #-}
 
 withTypes :: p (a::k) -> q (b::h) -> f a b -> f a b
@@ -130,14 +133,15 @@ modint r = IntegralRep i2f i2fM f2i f2iM unMod (n `seq` Just n)
  where
     -- n = withTypeable r $ fromIntegral . natVal $ (undefined :: Proxy n) -- If only..
     n = unsafePerformIO $ do
-        SomeNat c <- readIORef cachedNat
+        CachedNat ci c <- readIORef cachedNat
         withTypeable r $ do
             case withTypes c r <$> eqT of
-                Just Refl -> return $ fromIntegral $ natVal c
+                Just Refl -> return $ fromIntegral ci
                 _ -> do
                     let newnat = read . show $ r -- XXX: Hack to get nat value from Type.Reflection
                     case someNatVal $ fromIntegral newnat of
-                        Just somenat@(SomeNat nt) -> nt `seq` writeIORef cachedNat somenat
+                        Just somenat@(SomeNat nt) -> do
+                            nt `seq` writeIORef cachedNat $ CachedNat (fromIntegral newnat) nt
                         _ -> return ()
                     return newnat