{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE UndecidableInstances #-}

module Data.DeferIsPrime where

import GHC.TypeLits
import Data.Type.Equality
import Data.Proxy
import Unsafe.Coerce
import Control.Applicative

type family Not (x :: Bool) where
  Not 'True  = 'False
  Not 'False = 'True

type family (x :: Nat) % (y :: Nat) :: Nat where
  x % y = Mod' x y (y <=? x)

type family Mod' (x :: Nat) (y :: Nat) (xGeqY :: Bool) :: Nat where
  Mod' x y 'True  = Mod' (x - y) y (y <=? (x - y))
  Mod' x y 'False = x

type IsPrime p = IsPrimeB p ~ 'True

type family IsPrimeB (p :: Nat) :: Bool where
  IsPrimeB 0 = 'False
  IsPrimeB 1 = 'False
  IsPrimeB 2 = 'True
  IsPrimeB p = IsPrimeB' p 3 (p % 2 == 0) (Not (3 ^ 2 <=? p))

type family IsPrimeB' (p :: Nat) (i :: Nat) (hasFactor :: Bool) (searchEnd :: Bool) where
  IsPrimeB' _ _ 'True _ = 'False
  IsPrimeB' _ _ _ 'True = 'True
  IsPrimeB' p i _ _     = IsPrimeB' p (i + 2) (p % i == 0) (Not ((i + 2) ^ 2 <=? p))

data SomeIsPrime where
  SomeIsPrime :: (KnownNat p, IsPrime p) => Proxy p -> SomeIsPrime

unsafeMagic :: Proxy p -> IsPrimeB p :~: 'True
unsafeMagic _ = unsafeCoerce Refl

-- ここに型族のIsPrimeと同じ実装を書く
isPrime :: KnownNat n => Proxy n -> Bool
isPrime n = case natVal n of
  0 -> False -- テスト用
  1 -> True  -- テスト用
  _ -> undefined

-- |
-- >>> import Data.Maybe
-- >>> isNothing $ deferIsPrime 0
-- True
-- >>> isNothing $ deferIsPrime 1
-- False
--
deferIsPrime :: Integer -> Maybe SomeIsPrime
deferIsPrime i = do
  n <- someNatVal i
  case n of
    SomeNat p | isPrime p -> case unsafeMagic p of
      Refl -> pure $ SomeIsPrime p
    _  -> empty