{-# LANGUAGE CPP, LambdaCase, TemplateHaskell #-}
module Singlethongs.TH (singlethongs) where

import Data.Type.Equality
import Language.Haskell.TH
import Singlethongs.Internal

{-| Generate 'Sing', 'SingI', 'SingKind' and 'TestEquality' instances for a
datatype.

Given a datatype like @Foo@ below, having one or more unary constructors:

@
data Foo = Bar | Qux

'singlethongs' ''Foo
@

The following code will be generated:

@
data instance 'Sing' (x :: Foo) where
  SBar :: 'Sing' 'Bar
  SQux :: 'Sing' 'Qux

instance 'SingKind' Fobo where
  type 'Demote' Foo = Foo
  'fromSing' SBar = Bar
  'fromSing' SQux = Qux
  'toSing' Bar = 'SomeSing' SBar
  'toSing' Qux = 'SomeSing' SQux

instance 'SingI' 'Bar where 'sing' = SBar
instance 'SingI' 'Qux where 'sing' = SQux

instance 'TestEquality' ('Sing' :: Foo -> *) where
  'testEquality' SBar SBar = 'Just' 'Refl'
  'testEquality' SQux SQux = 'Just' 'Refl'
  'testEquality' _ _ = 'Nothing'
@

In order to use this 'singlethongs' function, you will need to enable the
following GHC extensions:

@
\{\-\# LANGUAGE DataKinds, GADTs, KindSignatures, TemplateHaskell, TypeFamilies \#\-\}
@
-}

singlethongs :: Name -> Q [Dec]
singlethongs :: Name -> Q [Dec]
singlethongs n0 :: Name
n0 = Name -> Q Info
reify Name
n0 Q Info -> (Info -> Q [Dec]) -> Q [Dec]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  TyConI (DataD [] n1 :: Name
n1 [] Nothing cons :: [Con]
cons@(_:_) []) -> do
    [Name]
nCons <- (Con -> Q Name) -> [Con] -> Q [Name]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Con -> Q Name
conName [Con]
cons
    [Dec]
out0 <- Name -> [Name] -> Q [Dec]
genDataInstSing Name
n1 [Name]
nCons
    [Dec]
out1 <- Name -> [Name] -> Q [Dec]
genInstanceSingKind Name
n1 [Name]
nCons
    [Dec]
out2 <- Name -> [Name] -> Q [Dec]
genInstanceTestEquality Name
n1 [Name]
nCons
    [Dec]
out3 <- [[Dec]] -> [Dec]
forall a. Monoid a => [a] -> a
mconcat ([[Dec]] -> [Dec]) -> Q [[Dec]] -> Q [Dec]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Name -> Q [Dec]) -> [Name] -> Q [[Dec]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Name -> Q [Dec]
genInstanceSingI [Name]
nCons
    [Dec] -> Q [Dec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Dec]
out0 [Dec] -> [Dec] -> [Dec]
forall a. Semigroup a => a -> a -> a
<> [Dec]
out1 [Dec] -> [Dec] -> [Dec]
forall a. Semigroup a => a -> a -> a
<> [Dec]
out2 [Dec] -> [Dec] -> [Dec]
forall a. Semigroup a => a -> a -> a
<> [Dec]
out3)
  _ -> String -> Q [Dec]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Only enum types are acceptable"

conName :: Con -> Q Name
conName :: Con -> Q Name
conName = \case
  NormalC n :: Name
n [] -> Name -> Q Name
forall (f :: * -> *) a. Applicative f => a -> f a
pure Name
n
  _ -> String -> Q Name
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Only enum types are acceptable"

sName :: Name -> Name
sName :: Name -> Name
sName a :: Name
a = String -> Name
mkName ("S" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Name -> String
nameBase Name
a)

genDataInstSing :: Name -> [Name] -> Q [Dec]
genDataInstSing :: Name -> [Name] -> Q [Dec]
genDataInstSing nTy :: Name
nTy nCons :: [Name]
nCons = do
  let cons1 :: [Con]
cons1 = ((Name -> Con) -> [Name] -> [Con])
-> [Name] -> (Name -> Con) -> [Con]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Name -> Con) -> [Name] -> [Con]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Name]
nCons ((Name -> Con) -> [Con]) -> (Name -> Con) -> [Con]
forall a b. (a -> b) -> a -> b
$ \nCon :: Name
nCon ->
        [Name] -> [BangType] -> Type -> Con
GadtC [Name -> Name
sName Name
nCon] [] (Type -> Type -> Type
AppT (Name -> Type
ConT ''Sing) (Name -> Type
PromotedT Name
nCon))
  [Dec] -> Q [Dec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Name -> [Con] -> Dec
mkSingDataInstD Name
nTy [Con]
cons1]

genInstanceSingI :: Name -> Q [Dec]
genInstanceSingI :: Name -> Q [Dec]
genInstanceSingI nCon :: Name
nCon = do
  let singD :: Dec
singD = Name -> [Clause] -> Dec
FunD (String -> Name
mkName "sing") [[Pat] -> Body -> [Dec] -> Clause
Clause [] (Exp -> Body
NormalB (Name -> Exp
ConE (Name -> Name
sName Name
nCon))) []]
  [Dec] -> Q [Dec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Maybe Overlap -> Cxt -> Type -> [Dec] -> Dec
InstanceD Maybe Overlap
forall a. Maybe a
Nothing [] (Type -> Type -> Type
AppT (Name -> Type
ConT ''SingI) (Name -> Type
PromotedT Name
nCon)) [Dec
singD]]

genInstanceSingKind :: Name -> [Name] -> Q [Dec]
genInstanceSingKind :: Name -> [Name] -> Q [Dec]
genInstanceSingKind nTy :: Name
nTy nCons :: [Name]
nCons = do
  let fromSingD :: Dec
fromSingD = Name -> [Clause] -> Dec
FunD (String -> Name
mkName "fromSing") ([Clause] -> Dec) -> [Clause] -> Dec
forall a b. (a -> b) -> a -> b
$ ((Name -> Clause) -> [Name] -> [Clause])
-> [Name] -> (Name -> Clause) -> [Clause]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Name -> Clause) -> [Name] -> [Clause]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Name]
nCons ((Name -> Clause) -> [Clause]) -> (Name -> Clause) -> [Clause]
forall a b. (a -> b) -> a -> b
$ \nCon :: Name
nCon ->
        [Pat] -> Body -> [Dec] -> Clause
Clause [Name -> [Pat] -> Pat
ConP (Name -> Name
sName Name
nCon) []] (Exp -> Body
NormalB (Name -> Exp
ConE Name
nCon)) []
      toSingD :: Dec
toSingD = Name -> [Clause] -> Dec
FunD (String -> Name
mkName "toSing") ([Clause] -> Dec) -> [Clause] -> Dec
forall a b. (a -> b) -> a -> b
$ ((Name -> Clause) -> [Name] -> [Clause])
-> [Name] -> (Name -> Clause) -> [Clause]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Name -> Clause) -> [Name] -> [Clause]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Name]
nCons ((Name -> Clause) -> [Clause]) -> (Name -> Clause) -> [Clause]
forall a b. (a -> b) -> a -> b
$ \nCon :: Name
nCon ->
        [Pat] -> Body -> [Dec] -> Clause
Clause [Name -> [Pat] -> Pat
ConP Name
nCon []]
               (Exp -> Body
NormalB (Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'SomeSing) (Name -> Exp
ConE (Name -> Name
sName Name
nCon)))) []
  [Dec] -> Q [Dec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Maybe Overlap -> Cxt -> Type -> [Dec] -> Dec
InstanceD Maybe Overlap
forall a. Maybe a
Nothing [] (Type -> Type -> Type
AppT (Name -> Type
ConT ''SingKind) (Name -> Type
ConT Name
nTy))
         [Name -> Dec
mkDemoteD Name
nTy, Dec
fromSingD, Dec
toSingD] ]

genInstanceTestEquality :: Name -> [Name] -> Q [Dec]
genInstanceTestEquality :: Name -> [Name] -> Q [Dec]
genInstanceTestEquality nTy :: Name
nTy nCons :: [Name]
nCons = do
  let teD :: Dec
teD = Name -> [Clause] -> Dec
FunD (String -> Name
mkName "testEquality") ([Clause] -> Dec) -> [Clause] -> Dec
forall a b. (a -> b) -> a -> b
$ [[Clause]] -> [Clause]
forall a. Monoid a => [a] -> a
mconcat
       [ ((Name -> Clause) -> [Name] -> [Clause])
-> [Name] -> (Name -> Clause) -> [Clause]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Name -> Clause) -> [Name] -> [Clause]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Name]
nCons ((Name -> Clause) -> [Clause]) -> (Name -> Clause) -> [Clause]
forall a b. (a -> b) -> a -> b
$ \nCon :: Name
nCon ->
           let p :: Pat
p = Name -> [Pat] -> Pat
ConP (Name -> Name
sName Name
nCon) []
           in [Pat] -> Body -> [Dec] -> Clause
Clause [Pat
p, Pat
p] (Exp -> Body
NormalB (Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Just) (Name -> Exp
ConE 'Refl))) []
       , case [Name]
nCons of
           [_] -> []
           _ -> [[Pat] -> Body -> [Dec] -> Clause
Clause [Pat
WildP, Pat
WildP] (Exp -> Body
NormalB (Name -> Exp
ConE 'Nothing)) []]
       ]
  [Dec] -> Q [Dec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Maybe Overlap -> Cxt -> Type -> [Dec] -> Dec
InstanceD Maybe Overlap
forall a. Maybe a
Nothing []
          (Type -> Type -> Type
AppT (Name -> Type
ConT ''TestEquality)
                (Type -> Type -> Type
SigT (Name -> Type
ConT ''Sing)
                      (Type -> Type -> Type
AppT (Type -> Type -> Type
AppT Type
ArrowT (Name -> Type
ConT Name
nTy)) Type
StarT)))
          [Dec
teD ]]


mkDemoteD :: Name -> Dec
mkDemoteD :: Name -> Dec
mkDemoteD nTy :: Name
nTy =
#if MIN_VERSION_template_haskell(2,15,0)
  TySynEqn -> Dec
TySynInstD (Maybe [TyVarBndr] -> Type -> Type -> TySynEqn
TySynEqn Maybe [TyVarBndr]
forall a. Maybe a
Nothing (Type -> Type -> Type
AppT (Name -> Type
ConT ''Demote) (Name -> Type
ConT Name
nTy)) (Name -> Type
ConT Name
nTy))
#else
  TySynInstD ''Demote (TySynEqn [ConT nTy] (ConT nTy))
#endif

mkSingDataInstD :: Name -> [Con] -> Dec
mkSingDataInstD :: Name -> [Con] -> Dec
mkSingDataInstD nTy :: Name
nTy cons :: [Con]
cons =
#if MIN_VERSION_template_haskell(2,15,0)
  Cxt
-> Maybe [TyVarBndr]
-> Type
-> Maybe Type
-> [Con]
-> [DerivClause]
-> Dec
DataInstD [] Maybe [TyVarBndr]
forall a. Maybe a
Nothing (Type -> Type -> Type
AppT (Name -> Type
ConT ''Sing) (Type -> Type -> Type
SigT (Name -> Type
VarT (String -> Name
mkName "x")) (Name -> Type
ConT Name
nTy)))
            Maybe Type
forall a. Maybe a
Nothing [Con]
cons []
#else
  DataInstD [] ''Sing [SigT (VarT (mkName "x")) (ConT nTy)] Nothing cons []
#endif