Type Level (extraordinarily) naive prime test

type-level-prime-test.hs

{-# LANGUAGE ConstraintKinds, DataKinds, FlexibleContexts                  #-}
{-# LANGUAGE FlexibleInstances, GADTs, MultiParamTypeClasses, PolyKinds    #-}
{-# LANGUAGE TypeFamilies, TypeInType, TypeOperators, UndecidableInstances #-}
module Lib where
import Data.Type.Equality
import GHC.Exts
import GHC.TypeLits

type Mod n p = Mod0 n p (n + 1 <=? p)
type family Mod0 n p b where
  Mod0 n p 'True = n
  Mod0 n p 'False = Mod0 (n - p) p (n - p + 1 <=? p)

data PrimeResult (n :: Nat) where
  Prime :: PrimeResult n
  DivisibleBy :: Nat -> PrimeResult n
  Unit :: PrimeResult n
  Zero :: PrimeResult n

type family If b (t :: k) f :: k where
  If 'True  t f = t
  If 'False t f = f

type family TestPrime p  :: PrimeResult p where
  TestPrime p  = TestPrimeAux p 2
type family TestPrimeAux p k :: PrimeResult p where
  TestPrimeAux 1 _ = 'Unit
  TestPrimeAux 0 _ = 'Zero
  TestPrimeAux p p = 'Prime
  TestPrimeAux p m =
    If (p + 1 <=? m * m)
       'Prime
       (If (Mod p m == 0) ('DivisibleBy m) (TestPrimeAux p (m + 1)))

type IsPrime p = IsPrimeAux (TestPrime p)
type family IsPrimeAux (pr :: PrimeResult p) :: Constraint where
  IsPrimeAux 'Prime = ()
  IsPrimeAux ('DivisibleBy n :: PrimeResult p) =
    TypeError ('ShowType p ':<>: 'Text " is not a prime; divisible by " ':<>: 'ShowType n ':<>: 'Text ".")
  IsPrimeAux 'Zero = TypeError ('Text "0 is not a prime.")
  IsPrimeAux ('Unit :: PrimeResult p) = TypeError ('ShowType p ':<>: 'Text " is not a prime; it's a multiplictive unit.")