-- | Helps to map enum types to postgresql enums.
module Database.PostgreSQL.Query.TH.Enum
  ( derivePgEnum
  , InflectorFunc
  ) where

import Data.FileEmbed
import Database.PostgreSQL.Query.TH.Common
import Database.PostgreSQL.Simple.FromField
import Database.PostgreSQL.Simple.ToField
import Language.Haskell.TH

import qualified Data.Text.Encoding as T
import qualified Data.Text as T

-- | Function to transform constructor name into its PG enum conterpart.
type InflectorFunc = String -> String

{-| derives 'FromField' and 'ToField' instances for a sum-type enum like

@
data Entity = Red | Green | Blue
@
-}
derivePgEnum
  :: InflectorFunc
     -- ^ mapping function from haskell constructor name to PG enum label
  -> Name
     -- ^ type to derive instances for
  -> DecsQ
derivePgEnum :: InflectorFunc -> Name -> DecsQ
derivePgEnum InflectorFunc
infl Name
typeName = do
  [Con]
constructors <- Info -> [Con]
dataConstructors (Info -> [Con]) -> Q Info -> Q [Con]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Q Info
reify Name
typeName
  Dec
tfInstance <- InflectorFunc -> Name -> [Con] -> DecQ
makeToField InflectorFunc
infl Name
typeName [Con]
constructors
  Dec
ffInstance <- InflectorFunc -> Name -> [Con] -> DecQ
makeFromField InflectorFunc
infl Name
typeName [Con]
constructors
  [Dec] -> DecsQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Dec
tfInstance, Dec
ffInstance]

makeToField :: InflectorFunc
            -> Name
            -> [Con]
            -> DecQ
makeToField :: InflectorFunc -> Name -> [Con] -> DecQ
makeToField InflectorFunc
i Name
typeName [Con]
constr = do
  [Clause]
clauses <- (Con -> Q Clause) -> [Con] -> Q [Clause]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (InflectorFunc -> Con -> Q Clause
makeToFieldClause InflectorFunc
i) [Con]
constr
  CxtQ -> TypeQ -> [DecQ] -> DecQ
instanceD
    ([Pred] -> CxtQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure [])
    (TypeQ -> TypeQ -> TypeQ
appT (Name -> TypeQ
conT ''ToField) (Name -> TypeQ
conT Name
typeName))
    [Name -> [Q Clause] -> DecQ
funD 'toField ([Q Clause] -> DecQ) -> [Q Clause] -> DecQ
forall a b. (a -> b) -> a -> b
$ (Clause -> Q Clause) -> [Clause] -> [Q Clause]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Clause -> Q Clause
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Clause]
clauses]

makeFromField :: InflectorFunc
              -> Name
              -> [Con]
              -> Q Dec
makeFromField :: InflectorFunc -> Name -> [Con] -> DecQ
makeFromField InflectorFunc
i Name
typeName [Con]
enumCons = do
  Name
f <- String -> Q Name
newName String
"f"
  Name
mb <- String -> Q Name
newName String
"mb"
  Name
byteSt <- String -> Q Name
newName String
"bs"
  Name
hName <- String -> Q Name
newName String
"h"
  let
    otherw :: Q (Guard, Exp)
otherw  = (,)
          (Guard -> Exp -> (Guard, Exp))
-> Q Guard -> Q (Exp -> (Guard, Exp))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpQ -> Q Guard
normalG [|otherwise|]
          Q (Exp -> (Guard, Exp)) -> ExpQ -> Q (Guard, Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [|returnError ConversionFailed $(varE f) (show $(varE mb))|]
    guards :: [Q (Guard, Exp)]
guards  = (Con -> Q (Guard, Exp)) -> [Con] -> [Q (Guard, Exp)]
forall a b. (a -> b) -> [a] -> [b]
map (InflectorFunc -> Name -> Con -> Q (Guard, Exp)
makeFromFieldGuard InflectorFunc
i Name
hName) [Con]
enumCons [Q (Guard, Exp)] -> [Q (Guard, Exp)] -> [Q (Guard, Exp)]
forall a. [a] -> [a] -> [a]
++ [Q (Guard, Exp)
otherw]
    helper :: DecQ
helper =
      Name -> [Q Clause] -> DecQ
funD
        Name
hName
        [[PatQ] -> BodyQ -> [DecQ] -> Q Clause
clause
          [Name -> PatQ
varP Name
byteSt]
          (ExpQ -> BodyQ
normalB [|((Just True) ==) (fmap (== $(varE byteSt)) $(varE mb))|])
          []
        ]
  CxtQ -> TypeQ -> [DecQ] -> DecQ
instanceD
    ([Pred] -> CxtQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure [])
    (TypeQ -> TypeQ -> TypeQ
appT (Name -> TypeQ
conT ''FromField) (Name -> TypeQ
conT Name
typeName))
    [Name -> [Q Clause] -> DecQ
funD 'fromField [[PatQ] -> BodyQ -> [DecQ] -> Q Clause
clause [Name -> PatQ
varP Name
f, Name -> PatQ
varP Name
mb] ([Q (Guard, Exp)] -> BodyQ
guardedB [Q (Guard, Exp)]
guards) [DecQ
helper]]]

makeFromFieldGuard :: InflectorFunc
                   -> Name           -- ^ shared helper function
                   -> Con            -- ^ constructor name
                   -> Q (Guard, Exp)
makeFromFieldGuard :: InflectorFunc -> Name -> Con -> Q (Guard, Exp)
makeFromFieldGuard InflectorFunc
i Name
typeName Con
con =
  ((Name -> ExpQ -> Q (Guard, Exp)) -> Con -> Q (Guard, Exp))
-> Con -> (Name -> ExpQ -> Q (Guard, Exp)) -> Q (Guard, Exp)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (InflectorFunc
-> (Name -> ExpQ -> Q (Guard, Exp)) -> Con -> Q (Guard, Exp)
forall a. InflectorFunc -> (Name -> ExpQ -> Q a) -> Con -> Q a
withEnumConstructor InflectorFunc
i) Con
con ((Name -> ExpQ -> Q (Guard, Exp)) -> Q (Guard, Exp))
-> (Name -> ExpQ -> Q (Guard, Exp)) -> Q (Guard, Exp)
forall a b. (a -> b) -> a -> b
$ \Name
nam ExpQ
ec -> do
    let constr :: ExpQ
constr             = Name -> ExpQ
conE Name
nam
    Guard
guard <- ExpQ -> Q Guard
normalG (ExpQ -> Q Guard) -> ExpQ -> Q Guard
forall a b. (a -> b) -> a -> b
$ ExpQ -> ExpQ -> ExpQ
appE (Name -> ExpQ
varE Name
typeName) ExpQ
ec
    Exp
expr <- ExpQ -> ExpQ -> ExpQ
appE (Name -> ExpQ
varE 'pure) ExpQ
constr
    (Guard, Exp) -> Q (Guard, Exp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Guard
guard, Exp
expr)

makeToFieldClause :: InflectorFunc
                  -> Con
                  -> ClauseQ
makeToFieldClause :: InflectorFunc -> Con -> Q Clause
makeToFieldClause InflectorFunc
i Con
con =
  ((Name -> ExpQ -> Q Clause) -> Con -> Q Clause)
-> Con -> (Name -> ExpQ -> Q Clause) -> Q Clause
forall a b c. (a -> b -> c) -> b -> a -> c
flip (InflectorFunc -> (Name -> ExpQ -> Q Clause) -> Con -> Q Clause
forall a. InflectorFunc -> (Name -> ExpQ -> Q a) -> Con -> Q a
withEnumConstructor InflectorFunc
i) Con
con ((Name -> ExpQ -> Q Clause) -> Q Clause)
-> (Name -> ExpQ -> Q Clause) -> Q Clause
forall a b. (a -> b) -> a -> b
$ \Name
nam ExpQ
ec -> do
    [PatQ] -> BodyQ -> [DecQ] -> Q Clause
clause [Name -> [PatQ] -> PatQ
conP Name
nam []] (ExpQ -> BodyQ
normalB [|Escape $ec|]) []

-- | Takes constructor w/o arguments and apply callback function.
-- Ejects with 'error' if called with wrong type of constructor.
withEnumConstructor :: InflectorFunc
                    -- ^ function to transform the constructor name
                    -> (Name -> ExpQ -> Q a)
                    -- ^ callback function from:
                    --   1. haskell constructor name and
                    --   2. PG enum option (ByteString)
                    -> Con
                    -- ^ constructor to decompose
                    -> Q a
withEnumConstructor :: InflectorFunc -> (Name -> ExpQ -> Q a) -> Con -> Q a
withEnumConstructor InflectorFunc
i Name -> ExpQ -> Q a
f = \case
  (NormalC Name
_    (BangType
_:[BangType]
_))   ->
    String -> Q a
forall a. HasCallStack => String -> a
error String
"constructors with arguments are not supported in makeToFieldClause"
  (NormalC Name
nam  []   ) -> Name -> ExpQ -> Q a
f Name
nam ExpQ
inflectedBs
    where inflectedT :: Text
inflectedT  = String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ InflectorFunc
i InflectorFunc -> InflectorFunc
forall a b. (a -> b) -> a -> b
$ Name -> String
nameBase Name
nam
          inflectedBs :: ExpQ
inflectedBs = ByteString -> ExpQ
bsToExp (ByteString -> ExpQ) -> ByteString -> ExpQ
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
T.encodeUtf8 Text
inflectedT
  Con
_                      ->
    String -> Q a
forall a. HasCallStack => String -> a
error String
"unsupported constructor in makeFromFieldClause"