{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE PolyKinds            #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}

module Debug.RecoverRTTI.Util.TypeLevel (
    -- * Singletons
    Sing(..)
  , SingI(..)
  , DecidableEquality(..)
    -- ** Natural numbers
  , Nat(..)
  , knownNat
  , Length
    -- * General purpose type level functions
  , Or
  , Equal
  , Elem
  , Assert
    -- * Type-level membership check
  , IsElem(..)
  , checkIsElem
    -- * Phantom type parameters
  , Phantom(..)
  , Poly(..)
  , maybePoly
  ) where

import Data.Kind
import Data.Proxy
import Data.Type.Equality
import GHC.TypeLits (ErrorMessage, Symbol, KnownSymbol, TypeError, sameSymbol)

{-------------------------------------------------------------------------------
  Singletons
-------------------------------------------------------------------------------}

data family Sing :: k -> Type

class SingI (a :: k) where
  sing :: Sing a

class DecidableEquality k where
  decideEquality :: Sing (a :: k) -> Sing (b :: k) -> Maybe (a :~: b)

{-------------------------------------------------------------------------------
  For kind 'Type', Sing is just a proxy
-------------------------------------------------------------------------------}

data instance Sing (a :: Type) where
  SProxy :: Sing (a :: Type)

instance SingI (a :: Type) where
  sing :: Sing a
sing = Sing a
forall a. Sing a
SProxy

{-------------------------------------------------------------------------------
  Natural numbers

  Unlike @ghc@'s, these are inductively defined.
-------------------------------------------------------------------------------}

data Nat = Z | S Nat

data instance Sing (n :: Nat) where
  SZ :: Sing 'Z
  SS :: Sing n -> Sing ('S n)

instance            SingI 'Z     where sing :: Sing 'Z
sing = Sing 'Z
SZ
instance SingI n => SingI ('S n) where sing :: Sing ('S n)
sing = Sing n -> Sing ('S n)
forall (n :: Nat). Sing n -> Sing ('S n)
SS Sing n
forall k (a :: k). SingI a => Sing a
sing

knownNat :: Sing (n :: Nat) -> Int
knownNat :: Sing n -> Int
knownNat Sing n
SZ     = Int
0
knownNat (SS n) = Sing n -> Int
forall (n :: Nat). Sing n -> Int
knownNat Sing n
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

type family Length (xs :: [k]) :: Nat where
  Length '[]       = 'Z
  Length (_ ': xs) = 'S (Length xs)

{-------------------------------------------------------------------------------
  Singleton instance for type-level symbols
-------------------------------------------------------------------------------}

data instance Sing (n :: Symbol) where
  SSymbol :: KnownSymbol n => Sing n

instance KnownSymbol n => SingI (n :: Symbol) where
  sing :: Sing n
sing = Sing n
forall (n :: Symbol). KnownSymbol n => Sing n
SSymbol

instance DecidableEquality Symbol where
  decideEquality :: Sing a -> Sing b -> Maybe (a :~: b)
decideEquality Sing a
SSymbol Sing b
SSymbol = Proxy a -> Proxy b -> Maybe (a :~: b)
forall (a :: Symbol) (b :: Symbol).
(KnownSymbol a, KnownSymbol b) =>
Proxy a -> Proxy b -> Maybe (a :~: b)
sameSymbol Proxy a
forall k (t :: k). Proxy t
Proxy Proxy b
forall k (t :: k). Proxy t
Proxy

{-------------------------------------------------------------------------------
  Singleton instance for lists
-------------------------------------------------------------------------------}

data instance Sing (xs :: [k]) where
  SN :: Sing '[]
  SC :: Sing x -> Sing xs -> Sing (x ': xs)

instance                        SingI '[]       where sing :: Sing '[]
sing = Sing '[]
forall k. Sing '[]
SN
instance (SingI x, SingI xs) => SingI (x ': xs) where sing :: Sing (x : xs)
sing = Sing x -> Sing xs -> Sing (x : xs)
forall a (x :: a) (xs :: [a]). Sing x -> Sing xs -> Sing (x : xs)
SC Sing x
forall k (a :: k). SingI a => Sing a
sing Sing xs
forall k (a :: k). SingI a => Sing a
sing

{-------------------------------------------------------------------------------
  General purpose type level functions
-------------------------------------------------------------------------------}

type family Or (a :: Bool) (b :: Bool) where
  Or 'True b     = 'True
  Or a     'True = 'True
  Or _     _     = 'False

type family Equal (x :: k) (y :: k) where
  Equal x x = 'True
  Equal x y = 'False

type family Elem (x :: k) (xs :: [k]) where
  Elem x '[]       = 'False
  Elem x (y ': ys) = Or (Equal x y) (Elem x ys)

-- | Assert type-level predicate
--
-- We cannot define this in terms of a more general @If@ construct, because
-- @ghc@'s type-level language has an undefined reduction order and so we get
-- no short-circuiting.
type family Assert (b :: Bool) (err :: ErrorMessage) :: Constraint where
  Assert 'True  err = ()
  Assert 'False err = TypeError err

{-------------------------------------------------------------------------------
  Decidable equality gives a decidable membership check
-------------------------------------------------------------------------------}

data IsElem (x :: k) (xs :: [k]) where
  IsElem :: Elem x xs ~ 'True => IsElem x xs

shiftIsElem :: IsElem x ys -> IsElem x (y ': ys)
shiftIsElem :: IsElem x ys -> IsElem x (y : ys)
shiftIsElem IsElem x ys
IsElem = IsElem x (y : ys)
forall k (x :: k) (xs :: [k]). (Elem x xs ~ 'True) => IsElem x xs
IsElem

checkIsElem ::
     DecidableEquality k
  => Sing (x :: k) -> Sing (xs :: [k]) -> Maybe (IsElem x xs)
checkIsElem :: Sing x -> Sing xs -> Maybe (IsElem x xs)
checkIsElem Sing x
_ Sing xs
SN         = Maybe (IsElem x xs)
forall a. Maybe a
Nothing
checkIsElem Sing x
x (SC y ys) = case Sing x -> Sing x -> Maybe (x :~: x)
forall k (a :: k) (b :: k).
DecidableEquality k =>
Sing a -> Sing b -> Maybe (a :~: b)
decideEquality Sing x
x Sing x
y of
                            Just x :~: x
Refl -> IsElem x xs -> Maybe (IsElem x xs)
forall a. a -> Maybe a
Just IsElem x xs
forall k (x :: k) (xs :: [k]). (Elem x xs ~ 'True) => IsElem x xs
IsElem
                            Maybe (x :~: x)
Nothing   -> IsElem x xs -> IsElem x (x : xs)
forall a (x :: a) (ys :: [a]) (y :: a).
IsElem x ys -> IsElem x (y : ys)
shiftIsElem (IsElem x xs -> IsElem x (x : xs))
-> Maybe (IsElem x xs) -> Maybe (IsElem x (x : xs))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sing x -> Sing xs -> Maybe (IsElem x xs)
forall k (x :: k) (xs :: [k]).
DecidableEquality k =>
Sing x -> Sing xs -> Maybe (IsElem x xs)
checkIsElem Sing x
x Sing xs
ys

{-------------------------------------------------------------------------------
  Phantom type parameters
-------------------------------------------------------------------------------}

-- | Functors with phantom arguments
class Phantom (f :: k -> Type) where
  -- | Similar to 'Data.Functor.Contravariant.phantom', but without requiring
  -- 'Functor' or 'Contravariant'
  phantom :: forall a b. f a -> f b

data Poly (f :: k -> Type) = Poly (forall (a :: k). f a)

-- | Commute @Maybe@ and @forall@
maybePoly :: Phantom f => Maybe (f a) -> Maybe (Poly f)
maybePoly :: Maybe (f a) -> Maybe (Poly f)
maybePoly = (f a -> Poly f) -> Maybe (f a) -> Maybe (Poly f)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\f a
v -> (forall (a :: k). f a) -> Poly f
forall k (f :: k -> *). (forall (a :: k). f a) -> Poly f
Poly (f a -> f a
forall k (f :: k -> *) (a :: k) (b :: k). Phantom f => f a -> f b
phantom f a
v))