Generating CmpNat constraint at runtime

CmpNat.hs

#!/usr/bin/env stack
-- stack --resolver lts-14.15 script --package reflection
{-# LANGUAGE DataKinds, RankNTypes, TypeOperators, TypeFamilies, ScopedTypeVariables, TypeApplications #-}
import GHC.TypeNats
import Data.Type.Equality
import Unsafe.Coerce
import Data.Proxy
import Data.Reflection

cmpNat :: (KnownNat a, KnownNat b)
  => Proxy a -> Proxy b
  -> (CmpNat a b :~: 'LT -> r)
  -> (CmpNat a b :~: 'EQ -> r)
  -> (CmpNat a b :~: 'GT -> r)
  -> r
cmpNat proxyA proxyB rLT rEQ rGT = case compare (natVal proxyA) (natVal proxyB) of
  LT -> rLT (unsafeCoerce Refl)
  EQ -> rEQ (unsafeCoerce Refl)
  GT -> rGT (unsafeCoerce Refl)

lt :: forall a b. (KnownNat a, KnownNat b, CmpNat a b ~ 'LT) => Proxy a -> Proxy b -> String
lt proxyA proxyB = concat [show (natVal proxyA), " < ", show (natVal proxyB)]

eq :: forall a b. (KnownNat a, KnownNat b, CmpNat a b ~ 'EQ) => Proxy a -> Proxy b -> String
eq proxyA proxyB = concat [show (natVal proxyA), " = ", show (natVal proxyB)]

gt :: forall a b. (KnownNat a, KnownNat b, CmpNat a b ~ 'GT) => Proxy a -> Proxy b -> String
gt proxyA proxyB = concat [show (natVal proxyA), " > ", show (natVal proxyB)]

compareAndDoSomething :: (KnownNat a, KnownNat b) => Proxy a -> Proxy b -> String
compareAndDoSomething proxyA proxyB = cmpNat proxyA proxyB (\Refl -> lt) (\Refl -> eq) (\Refl -> gt) proxyA proxyB

main :: IO ()
main = do
  putStrLn $ compareAndDoSomething (Proxy @3) (Proxy @5)
  putStrLn "Enter something:"
  n <- readLn
  putStrLn $ reifyNat n compareAndDoSomething (Proxy @10)

Prime.hs

#!/usr/bin/env stack
-- stack --resolver lts-14.15 script --package reflection
{-# LANGUAGE DataKinds, RankNTypes, TypeOperators, TypeFamilies, ScopedTypeVariables, TypeApplications, UndecidableInstances #-}
import GHC.TypeNats
import Data.Type.Equality
import Unsafe.Coerce
import Data.Proxy
import Data.Reflection

type family TrialDivision (n :: Nat) (m :: Nat) (l :: Nat) where
  TrialDivision n _ 0 = 'False
  TrialDivision n 1 _ = 'True
  TrialDivision n m _ = TrialDivision n (m - 1) (Mod n m)

type family IsPrime (n :: Nat) where
  IsPrime 0 = 'False
  IsPrime 1 = 'False
  IsPrime 2 = 'True
  IsPrime n = TrialDivision n (n - 1) 1

checkPrime :: (KnownNat a) => Proxy a -> Maybe (IsPrime a :~: 'True)
checkPrime proxy = let n = natVal proxy
                   in if n >= 2 && and [ n `rem` m /= 0 | m <- [2..n], m*m <= n ] then
                        Just (unsafeCoerce Refl)
                      else
                        Nothing

doSomethingWithPrime :: (KnownNat a, IsPrime a ~ 'True) => Proxy a -> IO ()
doSomethingWithPrime proxy = putStrLn $ "Yes! " ++ show (natVal proxy) ++ " is a prime!!!"

main = do
  print $ checkPrime (Proxy @0)
  print $ checkPrime (Proxy @1)
  print $ checkPrime (Proxy @2)
  print $ checkPrime (Proxy @3)
  print $ checkPrime (Proxy @4)
  print $ checkPrime (Proxy @5)
  print $ checkPrime (Proxy @6)
  print $ checkPrime (Proxy @7)
  print $ checkPrime (Proxy @8)
  print $ checkPrime (Proxy @9)
  print $ checkPrime (Proxy @10)
  doSomethingWithPrime (Proxy @7)
  -- doSomethingWithPrime (Proxy @8) -- type error
  -- doSomethingWithPrime (Proxy @(10^9+7)) -- Reduction stack overflow
  putStrLn "Enter a number:"
  n <- readLn
  reifyNat n $ \proxy -> case checkPrime proxy of
                           Just Refl -> doSomethingWithPrime proxy
                           Nothing -> putStrLn "not a prime"